From 142421da5b966049b4eab44ce6856eb172cf122a Mon Sep 17 00:00:00 2001
From: Anna Kozlova <anna.kozlova@jetbrains.com>
Date: Tue, 4 Jul 2023 13:31:44 +0200
Subject: [PATCH] [AA] support psi-resolve for loops labels

required for navigation, rename (KTIJ-25762), find usages.
Resolve through symbols can't work, because loops are not declarations
---
 .../api/fir/references/FirSimpleNameReference.kt         | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/references/FirSimpleNameReference.kt b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/references/FirSimpleNameReference.kt
index 21339aed122..1204eb800e9 100644
--- a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/references/FirSimpleNameReference.kt
+++ b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/references/FirSimpleNameReference.kt
@@ -12,6 +12,9 @@ import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirNamedClassOrObjectSymb
 import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirSyntheticJavaPropertySymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KtClassKind
 import org.jetbrains.kotlin.analysis.api.symbols.KtSymbol
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFir
+import org.jetbrains.kotlin.fir.expressions.FirLoopJump
+import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.psi.*
 
 internal class KtFirSimpleNameReference(
@@ -45,6 +48,12 @@ internal class KtFirSimpleNameReference(
     }
 
     override fun getResolvedToPsi(analysisSession: KtAnalysisSession): Collection<PsiElement> = with(analysisSession) {
+        if (expression is KtLabelReferenceExpression) {
+            val fir = expression.getOrBuildFir((analysisSession as KtFirAnalysisSession).firResolveSession)
+            if (fir is FirLoopJump) {
+                return listOfNotNull(fir.target.labeledElement.psi)
+            }
+        }
         val referenceTargetSymbols = resolveToSymbols()
         val psiOfReferenceTarget = super.getResolvedToPsi(analysisSession, referenceTargetSymbols)
         if (psiOfReferenceTarget.isNotEmpty()) return psiOfReferenceTarget