From 1729eff31b5d02285268c21a9c7364f054b94eec Mon Sep 17 00:00:00 2001
From: Jinseong Jeon <jsjeon@google.com>
Date: Mon, 8 Feb 2021 11:26:19 -0800
Subject: [PATCH] FIR checker: reincarnate FIR source child lookup utils

This is a partial revert of commit 94ddb712130f
---
 .../kotlin/fir/analysis/FirSourceChild.kt     | 41 +++++++++++++++++++
 1 file changed, 41 insertions(+)
 create mode 100644 compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/FirSourceChild.kt

diff --git a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/FirSourceChild.kt b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/FirSourceChild.kt
new file mode 100644
index 00000000000..e173eec7cb2
--- /dev/null
+++ b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/FirSourceChild.kt
@@ -0,0 +1,41 @@
+/*
+ * Copyright 2010-2021 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.getChild(type: IElementType, index: Int = 0, depth: Int = -1): FirSourceElement? {
+    return getChild(setOf(type), index, depth)
+}
+
+fun FirSourceElement.getChild(types: TokenSet, index: Int = 0, depth: Int = -1): FirSourceElement? {
+    return getChild(types.types.toSet(), index, depth)
+}
+
+fun FirSourceElement.getChild(types: Set<IElementType>, index: Int = 0, depth: Int = -1): FirSourceElement? {
+    return when (this) {
+        is FirPsiSourceElement<*> -> {
+            getChild(types, index, depth)
+        }
+        is FirLightSourceElement -> {
+            getChild(types, index, depth)
+        }
+        else -> null
+    }
+}
+
+private fun FirPsiSourceElement<*>.getChild(types: Set<IElementType>, index: Int, depth: Int): FirSourceElement? {
+    val visitor = PsiElementFinderByType(types, index, depth)
+    return visitor.find(psi)?.toFirPsiSourceElement()
+}
+
+private fun FirLightSourceElement.getChild(types: Set<IElementType>, index: Int, depth: Int): FirSourceElement? {
+    val visitor = LighterTreeElementFinderByType(treeStructure, types, index, depth)
+
+    return visitor.find(lighterASTNode)?.toFirLightSourceElement(treeStructure)
+}
\ No newline at end of file