From f51b108821dbaa38f58b1f9119ecd1b454e8ac8e Mon Sep 17 00:00:00 2001
From: Tianyu Geng <tgeng@google.com>
Date: Mon, 25 Oct 2021 16:31:40 -0700
Subject: [PATCH] FIR: fix PsiElementFinderByType

Current implementation would return the input node if the type matches.
This is clear wrong since the method is supposed to return children.
---
 .../fir/analysis/LighterTreeElementFinderByType.kt     | 10 ++++++----
 .../kotlin/fir/analysis/PsiElementFinderByType.kt      | 10 ++++++----
 2 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/LighterTreeElementFinderByType.kt b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/LighterTreeElementFinderByType.kt
index ab6b90c4f86..91ed3db4edc 100644
--- a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/LighterTreeElementFinderByType.kt
+++ b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/LighterTreeElementFinderByType.kt
@@ -22,11 +22,13 @@ class LighterTreeElementFinderByType(
     }
 
     private fun visitNode(node: LighterASTNode, currentDepth: Int): LighterASTNode? {
-        if (node.tokenType in types) {
-            if (index == 0) {
-                return node
+        if (currentDepth != 0) {
+            if (node.tokenType in types) {
+                if (index == 0) {
+                    return node
+                }
+                index--
             }
-            index--
         }
 
         if (currentDepth == depth) return null
diff --git a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/PsiElementFinderByType.kt b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/PsiElementFinderByType.kt
index 0fc7dee9a8f..912854ef1c7 100644
--- a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/PsiElementFinderByType.kt
+++ b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/PsiElementFinderByType.kt
@@ -19,11 +19,13 @@ class PsiElementFinderByType(
     }
 
     private fun visitElement(element: PsiElement, currentDepth: Int): PsiElement? {
-        if (element.node.elementType in types) {
-            if (index == 0) {
-                return element
+        if (currentDepth != 0) {
+            if (element.node.elementType in types) {
+                if (index == 0) {
+                    return element
+                }
+                index--
             }
-            index--
         }
 
         if (currentDepth == depth) return null