[FIR LT] Extract getChildren function into common utils

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/LighterTreeElementFinderByType.kt

@@ -6,9 +6,9 @@
 package org.jetbrains.kotlin.fir.analysis
 
 import com.intellij.lang.LighterASTNode
-import com.intellij.openapi.util.Ref
 import com.intellij.psi.tree.IElementType
 import com.intellij.util.diff.FlyweightCapableTreeStructure
+import org.jetbrains.kotlin.util.getChildren
 
 class LighterTreeElementFinderByType(
     private val tree: FlyweightCapableTreeStructure<LighterASTNode>,
@@ -34,7 +34,7 @@ class LighterTreeElementFinderByType(
 
         if (currentDepth == depth) return null
 
-        val children = if (reverse) node.getChildren().asReversed() else node.getChildren()
+        val children = if (reverse) node.getChildren(tree).asReversed() else node.getChildren(tree)
         for (child in children) {
             val result = visitNode(child, currentDepth + 1)
             if (result != null) return result
@@ -42,11 +42,4 @@ class LighterTreeElementFinderByType(
 
         return null
     }
-
-    private fun LighterASTNode.getChildren(): List<LighterASTNode> {
-        val ref = Ref<Array<LighterASTNode?>>()
-        tree.getChildren(this, ref)
-        return ref.get()?.filterNotNull() ?: emptyList()
-    }
-
 }

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/FirHelpers.kt

@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.fir.analysis.checkers
 
 import com.intellij.lang.LighterASTNode
-import com.intellij.openapi.util.Ref
 import org.jetbrains.kotlin.*
 import org.jetbrains.kotlin.builtins.StandardNames.HASHCODE_NAME
 import org.jetbrains.kotlin.descriptors.ClassKind
@@ -44,6 +43,7 @@ import org.jetbrains.kotlin.types.model.KotlinTypeMarker
 import org.jetbrains.kotlin.types.model.TypeCheckerProviderContext
 import org.jetbrains.kotlin.util.ImplementationStatus
 import org.jetbrains.kotlin.util.OperatorNameConventions
+import org.jetbrains.kotlin.util.getChildren
 
 private val INLINE_ONLY_ANNOTATION_CLASS_ID: ClassId = ClassId.topLevel(FqName("kotlin.internal.InlineOnly"))
 
@@ -666,14 +666,11 @@ internal val KtSourceElement.defaultValueForParameter: KtSourceElement?
     }
 
 private fun findDefaultValue(source: KtLightSourceElement): KtLightSourceElement? {
-    val childrenRef = Ref<Array<LighterASTNode?>>()
-    source.treeStructure.getChildren(source.lighterASTNode, childrenRef)
-
     var defaultValue: LighterASTNode? = null
     var defaultValueOffset = source.startOffset
 
-    for (node in childrenRef.get()) {
+    val nodes = source.lighterASTNode.getChildren(source.treeStructure)
-        if (node == null) continue
+    for (node in nodes) {
         if (node.isExpression()) {
             defaultValue = node
             break

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/FirKeywordUtils.kt

@@ -10,14 +10,15 @@ import com.intellij.lang.LighterASTNode
 import com.intellij.psi.tree.TokenSet
 import com.intellij.util.diff.FlyweightCapableTreeStructure
 import org.jetbrains.kotlin.*
-import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.diagnostics.valOrVarKeyword
+import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.lexer.KtKeywordToken
 import org.jetbrains.kotlin.lexer.KtModifierKeywordToken
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.KtModifierList
 import org.jetbrains.kotlin.psi.KtModifierListOwner
 import org.jetbrains.kotlin.psi.KtValVarKeywordOwner
+import org.jetbrains.kotlin.util.getChildren
 
 // DO
 // - use this to retrieve modifiers on the source and confirm a certain modifier indeed appears

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/SourceHelpers.kt

@@ -15,12 +15,6 @@ import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.KtModifierList
 import org.jetbrains.kotlin.psi.psiUtil.visibilityModifierType
 
-internal fun LighterASTNode.getChildren(tree: FlyweightCapableTreeStructure<LighterASTNode>): List<LighterASTNode> {
-    val children = Ref<Array<LighterASTNode?>>()
-    val count = tree.getChildren(this, children)
-    return if (count > 0) children.get().filterNotNull() else emptyList()
-}
-
 /**
  * Returns the visibility by given KtModifierList
  */

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/SourceNavigator.kt

@@ -21,6 +21,7 @@ import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.stubs.elements.KtNameReferenceExpressionElementType
 import org.jetbrains.kotlin.psi.stubs.elements.KtTypeProjectionElementType
+import org.jetbrains.kotlin.util.getChildren
 
 /**
  * Service to answer source-related questions in generic fashion.

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/extended/CanBeReplacedWithOperatorAssignmentChecker.kt

@@ -9,20 +9,18 @@ import com.intellij.lang.LighterASTNode
 import com.intellij.psi.tree.IElementType
 import org.jetbrains.kotlin.KtFakeSourceElementKind
 import org.jetbrains.kotlin.KtLightSourceElement
-import org.jetbrains.kotlin.KtPsiSourceElement
 import org.jetbrains.kotlin.KtNodeTypes
+import org.jetbrains.kotlin.KtPsiSourceElement
+import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
+import org.jetbrains.kotlin.diagnostics.reportOn
 import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
 import org.jetbrains.kotlin.fir.analysis.checkers.expression.FirVariableAssignmentChecker
-import org.jetbrains.kotlin.fir.analysis.checkers.getChildren
-import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
-import org.jetbrains.kotlin.diagnostics.reportOn
 import org.jetbrains.kotlin.fir.dispatchReceiverClassTypeOrNull
 import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
 import org.jetbrains.kotlin.fir.expressions.FirVariableAssignment
 import org.jetbrains.kotlin.fir.expressions.calleeReference
 import org.jetbrains.kotlin.fir.expressions.toResolvedCallableSymbol
-import org.jetbrains.kotlin.fir.references.toResolvedCallableSymbol
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.types.coneType
@@ -30,6 +28,7 @@ import org.jetbrains.kotlin.fir.types.isPrimitive
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.KtBinaryExpression
 import org.jetbrains.kotlin.psi.KtNameReferenceExpression
+import org.jetbrains.kotlin.util.getChildren
 
 object CanBeReplacedWithOperatorAssignmentChecker : FirVariableAssignmentChecker() {
     override fun check(expression: FirVariableAssignment, context: CheckerContext, reporter: DiagnosticReporter) {
@@ -75,7 +74,7 @@ object CanBeReplacedWithOperatorAssignmentChecker : FirVariableAssignmentChecker
         prevOperator: LighterASTNode? = null
     ): Boolean {
         val tree = source.treeStructure
-        val children = expression.getChildren(tree).filterNotNull()
+        val children = expression.getChildren(tree)
 
         val operator = children.firstOrNull { it.tokenType == KtNodeTypes.OPERATION_REFERENCE }
         if (prevOperator != null && !isLightNodesHierarchicallyTrue(prevOperator, operator)) return false

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/extended/CanBeValChecker.kt

@@ -13,13 +13,13 @@ import org.jetbrains.kotlin.diagnostics.reportOn
 import org.jetbrains.kotlin.fir.analysis.cfa.AbstractFirPropertyInitializationChecker
 import org.jetbrains.kotlin.fir.analysis.cfa.requiresInitialization
 import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
-import org.jetbrains.kotlin.fir.analysis.checkers.getChildren
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
 import org.jetbrains.kotlin.fir.analysis.cfa.util.PropertyInitializationInfoData
 import org.jetbrains.kotlin.fir.expressions.calleeReference
 import org.jetbrains.kotlin.fir.references.toResolvedPropertySymbol
 import org.jetbrains.kotlin.fir.resolve.dfa.cfg.*
 import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
+import org.jetbrains.kotlin.util.getChildren
 
 object CanBeValChecker : AbstractFirPropertyInitializationChecker() {
     override fun analyze(data: PropertyInitializationInfoData, reporter: DiagnosticReporter, context: CheckerContext) {

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/extended/RedundantSingleExpressionStringTemplateChecker.kt

@@ -13,7 +13,6 @@ import org.jetbrains.kotlin.KtPsiSourceElement
 import org.jetbrains.kotlin.KtNodeTypes
 import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
 import org.jetbrains.kotlin.fir.analysis.checkers.expression.FirStringConcatenationCallChecker
-import org.jetbrains.kotlin.fir.analysis.checkers.getChildren
 import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.REDUNDANT_SINGLE_EXPRESSION_STRING_TEMPLATE
 import org.jetbrains.kotlin.diagnostics.reportOn
@@ -24,6 +23,7 @@ import org.jetbrains.kotlin.fir.types.classId
 import org.jetbrains.kotlin.fir.types.coneType
 import org.jetbrains.kotlin.name.StandardClassIds
 import org.jetbrains.kotlin.psi.KtStringTemplateExpression
+import org.jetbrains.kotlin.util.getChildren
 
 object RedundantSingleExpressionStringTemplateChecker : FirStringConcatenationCallChecker() {
     override fun check(expression: FirStringConcatenationCall, context: CheckerContext, reporter: DiagnosticReporter) {

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/syntax/FirConfusingWhenBranchSyntaxChecker.kt

@@ -18,11 +18,11 @@ import org.jetbrains.kotlin.KtRealPsiSourceElement
 import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.diagnostics.reportOn
 import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
-import org.jetbrains.kotlin.fir.analysis.checkers.getChildren
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
 import org.jetbrains.kotlin.fir.expressions.FirWhenExpression
 import org.jetbrains.kotlin.lexer.KtTokens.*
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.util.getChildren
 
 object FirConfusingWhenBranchSyntaxChecker : FirExpressionSyntaxChecker<FirWhenExpression, PsiElement>() {
     private val prohibitedTokens = TokenSet.create(

compiler/frontend.common-psi/src/org/jetbrains/kotlin/diagnostics/LightTreePositioningStrategies.kt

@@ -17,12 +17,10 @@ import org.jetbrains.kotlin.KtSourceElement
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.lexer.KtTokens.*
 import org.jetbrains.kotlin.psi.KtParameter.VAL_VAR_TOKEN_SET
-import org.jetbrains.kotlin.psi.KtValueArgument
-import org.jetbrains.kotlin.psi.stubs.elements.KtClassElementType
 import org.jetbrains.kotlin.psi.stubs.elements.KtConstantExpressionElementType
 import org.jetbrains.kotlin.psi.stubs.elements.KtStringTemplateExpressionElementType
-import org.jetbrains.kotlin.psi.stubs.elements.KtStubElementTypes
 import org.jetbrains.kotlin.psi.stubs.elements.KtTokenSets
+import org.jetbrains.kotlin.util.getChildren
 import org.jetbrains.kotlin.utils.addToStdlib.runUnless
 
 object LightTreePositioningStrategies {

compiler/frontend.common-psi/src/org/jetbrains/kotlin/diagnostics/LightTreePositioningStrategy.kt

@@ -15,6 +15,7 @@ import org.jetbrains.kotlin.KtSourceElement
 import org.jetbrains.kotlin.lexer.KtSingleValueToken
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.lexer.KtTokens.WHITE_SPACE
+import org.jetbrains.kotlin.util.getChildren
 
 open class LightTreePositioningStrategy {
     open fun markKtDiagnostic(element: KtSourceElement, diagnostic: KtDiagnostic): List<TextRange> {
@@ -93,10 +94,8 @@ internal fun LighterASTNode.isFiller() = tokenType in FILLER_TOKENS
 private fun hasSyntaxErrors(node: LighterASTNode, tree: FlyweightCapableTreeStructure<LighterASTNode>): Boolean {
     if (node.tokenType == TokenType.ERROR_ELEMENT) return true
 
-    val childrenRef = Ref<Array<LighterASTNode?>?>()
+    val children = node.getChildren(tree)
-    tree.getChildren(node, childrenRef)
+    return children.lastOrNull {
-    val children = childrenRef.get() ?: return false
-    return children.filterNotNull().lastOrNull {
         val tokenType = it.tokenType
         tokenType !is KtSingleValueToken && tokenType !in DOC_AND_COMMENT_TOKENS
     }?.let { hasSyntaxErrors(it, tree) } == true
@@ -104,14 +103,11 @@ private fun hasSyntaxErrors(node: LighterASTNode, tree: FlyweightCapableTreeStru
 
 val KtLightSourceElement.startOffsetSkippingComments: Int
     get() {
-        val kidsRef = Ref<Array<LighterASTNode?>>()
+        val children = lighterASTNode.getChildren(treeStructure)
-        treeStructure.getChildren(this.lighterASTNode, kidsRef)
-        val children = kidsRef.get()
 
         // The solution to find first non comment children will not work here. `treeStructure` can have different root
         // than original program. Because of that `startOffset` is relative and not in absolute value.
-        val comments = children.filterNotNull()
+        val comments = children.takeWhile { it.tokenType in DOC_AND_COMMENT_TOKENS }
-            .takeWhile { it.tokenType in DOC_AND_COMMENT_TOKENS }
         return startOffset + comments.sumOf { it.textLength }
     }
 

compiler/frontend.common-psi/src/org/jetbrains/kotlin/diagnostics/UnreachableCodeLightTreeHelper.kt

@@ -6,10 +6,10 @@
 package org.jetbrains.kotlin.diagnostics
 
 import com.intellij.lang.LighterASTNode
-import com.intellij.openapi.util.Ref
 import com.intellij.openapi.util.TextRange
 import com.intellij.util.diff.FlyweightCapableTreeStructure
 import org.jetbrains.kotlin.lexer.KtTokens
+import org.jetbrains.kotlin.util.getChildren
 
 typealias Node = LighterASTNode
 
@@ -84,9 +84,3 @@ class UnreachableCodeLightTreeHelper(val tree: FlyweightCapableTreeStructure<Nod
         return result
     }
 }
-
-internal fun LighterASTNode.getChildren(tree: FlyweightCapableTreeStructure<LighterASTNode>): List<LighterASTNode> {
-    val children = Ref<Array<LighterASTNode?>>()
-    val count = tree.getChildren(this, children)
-    return if (count > 0) children.get().filterNotNull() else emptyList()
-}

compiler/frontend.common/src/org/jetbrains/kotlin/util/LightTreeUtils.kt

@@ -0,0 +1,16 @@
+/*
+ * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.util
+
+import com.intellij.lang.LighterASTNode
+import com.intellij.openapi.util.Ref
+import com.intellij.util.diff.FlyweightCapableTreeStructure
+
+fun LighterASTNode.getChildren(tree: FlyweightCapableTreeStructure<LighterASTNode>): List<LighterASTNode> {
+    val children = Ref<Array<LighterASTNode?>>()
+    val count = tree.getChildren(this, children)
+    return if (count > 0) children.get().filterNotNull() else emptyList()
+}

compiler/tests-common-new/tests/org/jetbrains/kotlin/test/frontend/fir/handlers/lightTreeDiagnostics.kt

@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.test.frontend.fir.handlers
 
 import com.intellij.lang.LighterASTNode
-import com.intellij.openapi.util.Ref
 import com.intellij.psi.TokenType
 import com.intellij.util.diff.FlyweightCapableTreeStructure
 import org.jetbrains.kotlin.KtLightSourceElement
@@ -15,20 +14,14 @@ import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.visitors.FirVisitor
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.toKtLightSourceElement
+import org.jetbrains.kotlin.util.getChildren
 
 private typealias Tree = FlyweightCapableTreeStructure<LighterASTNode>
 
 private class LightTreeErrorsCollector(private val tree: Tree) {
-
-    private fun LighterASTNode.getChildrenAsArray(): Array<out LighterASTNode?> {
-        val kidsRef = Ref<Array<LighterASTNode?>>()
-        return if (tree.getChildren(this, kidsRef) > 0) kidsRef.get() else emptyArray()
-    }
-
     private inline fun LighterASTNode.forEachChildren(f: (LighterASTNode) -> Unit) {
-        val kidsArray = this.getChildrenAsArray()
+        val kidsArray = this.getChildren(tree)
         for (kid in kidsArray) {
-            if (kid == null) break
             val tokenType = kid.tokenType
             if (KtTokens.COMMENTS.contains(tokenType) || tokenType == KtTokens.WHITE_SPACE || tokenType == KtTokens.SEMICOLON) continue
             f(kid)