[FIR IDE] Properly support loop variables in KtFirSymbolProvider

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/symbols/KtFirSymbolProvider.kt

@@ -35,12 +35,19 @@ internal class KtFirSymbolProvider(
 ) : KtSymbolProvider(), ValidityTokenOwner {
     private val firSymbolProvider by weakRef(firSymbolProvider)
 
-    override fun getParameterSymbol(psi: KtParameter): KtValueParameterSymbol = withValidityAssertion {
-        if (psi.isFunctionTypeParameter) {
-            error("Creating KtValueParameterSymbol for function type parameter is not possible. Please see the KDoc of getParameterSymbol")
-        }
-        psi.withFirDeclarationOfType<FirValueParameter, KtValueParameterSymbol>(resolveState) {
-            firSymbolBuilder.variableLikeBuilder.buildValueParameterSymbol(it)
+    override fun getParameterSymbol(psi: KtParameter): KtVariableLikeSymbol = withValidityAssertion {
+        when {
+            psi.isFunctionTypeParameter -> error(
+                "Creating KtValueParameterSymbol for function type parameter is not possible. Please see the KDoc of getParameterSymbol"
+            )
+
+            psi.isLoopParameter -> psi.withFirDeclarationOfType<FirProperty, KtLocalVariableSymbol>(resolveState) {
+                firSymbolBuilder.variableLikeBuilder.buildLocalVariableSymbol(it)
+            }
+
+            else -> psi.withFirDeclarationOfType<FirValueParameter, KtValueParameterSymbol>(resolveState) {
+                firSymbolBuilder.variableLikeBuilder.buildValueParameterSymbol(it)
+            }
         }
     }
 

analysis/analysis-api-fir/tests/org/jetbrains/kotlin/analysis/api/fir/symbols/SymbolByPsiTestGenerated.java

@@ -102,6 +102,12 @@ public class SymbolByPsiTestGenerated extends AbstractSymbolByPsiTest {
         runTest("analysis/analysis-api/testData/symbols/symbolByPsi/extensionFunction.kt");
     }
 
+    @Test
+    @TestMetadata("forLoopVariable.kt")
+    public void testForLoopVariable() throws Exception {
+        runTest("analysis/analysis-api/testData/symbols/symbolByPsi/forLoopVariable.kt");
+    }
+
     @Test
     @TestMetadata("function.kt")
     public void testFunction() throws Exception {

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/symbols/KtSymbolProvider.kt

@@ -31,7 +31,7 @@ public abstract class KtSymbolProvider : KtAnalysisSessionComponent() {
         else -> error("Cannot build symbol for ${psi::class}")
     }
 
-    public abstract fun getParameterSymbol(psi: KtParameter): KtValueParameterSymbol
+    public abstract fun getParameterSymbol(psi: KtParameter): KtVariableLikeSymbol
     public abstract fun getFileSymbol(psi: KtFile): KtFileSymbol
     public abstract fun getFunctionLikeSymbol(psi: KtNamedFunction): KtFunctionLikeSymbol
     public abstract fun getConstructorSymbol(psi: KtConstructor<*>): KtConstructorSymbol
@@ -60,12 +60,20 @@ public interface KtSymbolProviderMixIn : KtAnalysisSessionMixIn {
         analysisSession.symbolProvider.getSymbol(this)
 
     /**
-     * Creates [KtValueParameterSymbol] by [KtParameter]
+     * Creates [KtVariableLikeSymbol] by [KtParameter].
+     *
+     * Unfortunately, [KtParameter] in PSI stands for many things, and not all of them are represented by a single type of symbol,
+     * so this function does not work for all possible [KtParameter]s.
      *
      * If [KtParameter.isFunctionTypeParameter] is `true`, i.e., if the given [KtParameter] is used as a function type parameter,
      * it is not possible to create [KtValueParameterSymbol], hence an error will be raised.
+     *
+     * If [KtParameter.isLoopParameter] is `true`, i.e. if the given [KtParameter] is a loop variable in `for` expression, then the function
+     * returns [KtLocalVariableSymbol].
+     *
+     * Otherwise, returns [KtValueParameterSymbol].
      */
-    public fun KtParameter.getParameterSymbol(): KtValueParameterSymbol =
+    public fun KtParameter.getParameterSymbol(): KtVariableLikeSymbol =
         analysisSession.symbolProvider.getParameterSymbol(this)
 
     /**

analysis/analysis-api/testData/symbols/symbolByPsi/forLoopVariable.kt

@@ -0,0 +1,3 @@
+fun usage() {
+    for (loopVariable in 1..10) {}
+}

analysis/analysis-api/testData/symbols/symbolByPsi/forLoopVariable.txt

@@ -0,0 +1,35 @@
+KtFirLocalVariableSymbol:
+  annotatedType: [] kotlin/Int
+  callableIdIfNonLocal: null
+  isExtension: false
+  isVal: true
+  name: loopVariable
+  origin: SOURCE
+  receiverType: null
+  symbolKind: LOCAL
+  deprecationStatus: null
+
+KtFirFunctionSymbol:
+  annotatedType: [] kotlin/Unit
+  annotationClassIds: []
+  annotations: []
+  callableIdIfNonLocal: /usage
+  dispatchType: null
+  hasStableParameterNames: true
+  isExtension: false
+  isExternal: false
+  isInfix: false
+  isInline: false
+  isOperator: false
+  isOverride: false
+  isStatic: false
+  isSuspend: false
+  modality: FINAL
+  name: usage
+  origin: SOURCE
+  receiverType: null
+  symbolKind: TOP_LEVEL
+  typeParameters: []
+  valueParameters: []
+  visibility: Public
+  deprecationStatus: null