[FIR] Add methods to access child/children of FirSourceElement

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

@@ -0,0 +1,41 @@
+/*
+ * Copyright 2010-2020 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.fir.analysis
+
+import com.intellij.psi.tree.IElementType
+import com.intellij.psi.tree.TokenSet
+import org.jetbrains.kotlin.fir.*
+
+fun FirSourceElement.getChildren(type: IElementType, index: Int = 0, depth: Int = -1): FirSourceElement? {
+    return getChildren(setOf(type), index, depth)
+}
+
+fun FirSourceElement.getChildren(types: TokenSet, index: Int = 0, depth: Int = -1): FirSourceElement? {
+    return getChildren(types.types.toSet(), index, depth)
+}
+
+fun FirSourceElement.getChildren(types: Set<IElementType>, index: Int = 0, depth: Int = -1): FirSourceElement? {
+    return when (this) {
+        is FirPsiSourceElement<*> -> {
+            getChildren(types, index, depth)
+        }
+        is FirLightSourceElement -> {
+            getChildren(types, index, depth)
+        }
+        else -> null
+    }
+}
+
+private fun FirPsiSourceElement<*>.getChildren(types: Set<IElementType>, index: Int, depth: Int): FirSourceElement? {
+    val visitor = PsiElementFinderByType(types, index, depth)
+    return visitor.find(psi)?.toFirPsiSourceElement()
+}
+
+private fun FirLightSourceElement.getChildren(types: Set<IElementType>, index: Int, depth: Int): FirSourceElement? {
+    val visitor = LighterTreeElementFinderByType(tree, types, index, depth)
+
+    return visitor.find(lightNode)?.let { it.toFirLightSourceElement(it.startOffset, it.endOffset, tree) }
+}

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

@@ -0,0 +1,47 @@
+/*
+ * Copyright 2010-2020 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.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
+
+class LighterTreeElementFinderByType(
+    private val tree: FlyweightCapableTreeStructure<LighterASTNode>,
+    private var types: Collection<IElementType>,
+    private var index: Int,
+    private val depth: Int
+) {
+    fun find(node: LighterASTNode?): LighterASTNode? {
+        if (node == null) return null
+        return visitNode(node, 0)
+    }
+
+    fun visitNode(node: LighterASTNode, currentDepth: Int): LighterASTNode? {
+        if (node.tokenType in types) {
+            if (index == 0) {
+                return node
+            }
+            index--
+        }
+
+        if (currentDepth == depth) return null
+
+        for (child in node.getChildren()) {
+            val result = visitNode(child, currentDepth + 1)
+            if (result != null) return result
+        }
+
+        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/PsiElementFinderByType.kt

@@ -0,0 +1,38 @@
+/*
+ * Copyright 2010-2020 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.fir.analysis
+
+import com.intellij.psi.PsiElement
+import com.intellij.psi.tree.IElementType
+import org.jetbrains.kotlin.psi.psiUtil.allChildren
+
+class PsiElementFinderByType(
+    private val types: Collection<IElementType>,
+    private var index: Int,
+    private val depth: Int
+) {
+    fun find(root: PsiElement): PsiElement? {
+        return visitElement(root, 0)
+    }
+
+    fun visitElement(element: PsiElement, currentDepth: Int): PsiElement? {
+        if (element.node.elementType in types) {
+            if (index == 0) {
+                return element
+            }
+            index--
+        }
+
+        if (currentDepth == depth) return null
+
+        for (children in element.allChildren) {
+            val result = visitElement(children, currentDepth + 1)
+            if (result != null) return result
+        }
+
+        return null
+    }
+}