FIR2IR: provide correct origins for 'for' loops

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrDeclarationStorage.kt

@@ -779,11 +779,15 @@ class Fir2IrDeclarationStorage(
         }
     }
 
-    fun createAndSaveIrVariable(variable: FirVariable<*>): IrVariable {
+    fun createAndSaveIrVariable(variable: FirVariable<*>, givenOrigin: IrDeclarationOrigin? = null): IrVariable {
         val type = variable.returnTypeRef.toIrType(session, this)
         // Some temporary variables are produced in RawFirBuilder, but we consistently use special names for them.
-        val origin =
-            if (variable.name.isSpecial) IrDeclarationOrigin.IR_TEMPORARY_VARIABLE else IrDeclarationOrigin.DEFINED
+        val origin = when {
+            givenOrigin != null -> givenOrigin
+            variable.name == Name.special("<iterator>") -> IrDeclarationOrigin.FOR_LOOP_ITERATOR
+            variable.name.isSpecial -> IrDeclarationOrigin.IR_TEMPORARY_VARIABLE
+            else -> IrDeclarationOrigin.DEFINED
+        }
         val irVariable = variable.convertWithOffsets { startOffset, endOffset ->
             declareIrVariable(
                 startOffset, endOffset, origin,

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrVisitor.kt

@@ -469,9 +469,14 @@ class Fir2IrVisitor(
 
     private fun visitLocalVariable(variable: FirProperty): IrElement {
         assert(variable.isLocal)
-        val irVariable = declarationStorage.createAndSaveIrVariable(variable)
+        val initializer = variable.initializer
+        val isNextVariable = initializer is FirFunctionCall &&
+                initializer.resolvedNamedFunctionSymbol()?.callableId?.isIteratorNext() == true &&
+                variable.source.psi?.parent is KtForExpression
+        val irVariable = declarationStorage.createAndSaveIrVariable(
+            variable, if (isNextVariable) IrDeclarationOrigin.FOR_LOOP_VARIABLE else null
+        )
         return irVariable.setParentByParentStack().apply {
-            val initializer = variable.initializer
             if (initializer != null) {
                 this.initializer = initializer.toIrExpression()
             }
@@ -671,7 +676,13 @@ class Fir2IrVisitor(
             is FirPropertyFromParameterResolvedNamedReference -> IrStatementOrigin.INITIALIZE_PROPERTY_FROM_PARAMETER
             is FirResolvedNamedReference -> when (resolvedSymbol) {
                 is AccessorSymbol, is SyntheticPropertySymbol -> IrStatementOrigin.GET_PROPERTY
-                is FirNamedFunctionSymbol -> if (resolvedSymbol.callableId.isInvoke()) IrStatementOrigin.INVOKE else null
+                is FirNamedFunctionSymbol -> when {
+                    resolvedSymbol.callableId.isInvoke() -> IrStatementOrigin.INVOKE
+                    source.psi is KtForExpression && resolvedSymbol.callableId.isIteratorNext() -> IrStatementOrigin.FOR_LOOP_NEXT
+                    source.psi is KtForExpression && resolvedSymbol.callableId.isIteratorHasNext() -> IrStatementOrigin.FOR_LOOP_HAS_NEXT
+                    source.psi is KtForExpression && resolvedSymbol.callableId.isIterator() -> IrStatementOrigin.FOR_LOOP_ITERATOR
+                    else -> null
+                }
                 else -> null
             }
             else -> null
@@ -857,16 +868,24 @@ class Fir2IrVisitor(
         val convertibleCall = if (functionCall.toResolvedCallableSymbol()?.fir is FirIntegerOperator) {
             functionCall.copy().transformSingle(integerApproximator, null)
         } else {
-            functionCall.replaceCalleeReferenceWithOverridden()
+            val resolvedSymbol = functionCall.resolvedNamedFunctionSymbol()
+            if (resolvedSymbol != null) {
+                functionCall.replaceCalleeReferenceWithOverridden(resolvedSymbol)
+            } else {
+                functionCall
+            }
         }
         return convertibleCall.toIrExpression(convertibleCall.typeRef)
             .applyCallArguments(convertibleCall).applyTypeArguments(convertibleCall).applyReceivers(convertibleCall)
     }
 
+    private fun FirFunctionCall.resolvedNamedFunctionSymbol(): FirNamedFunctionSymbol? {
+        val calleeReference = (calleeReference as? FirResolvedNamedReference) ?: return null
+        return calleeReference.resolvedSymbol as? FirNamedFunctionSymbol
+    }
+
     // Use the generic invoke & next methods to match bridges generated by the backend.
-    private fun FirFunctionCall.replaceCalleeReferenceWithOverridden(): FirFunctionCall {
-        val calleeReference = (calleeReference as? FirResolvedNamedReference) ?: return this
-        val resolvedSymbol = (calleeReference.resolvedSymbol as? FirNamedFunctionSymbol) ?: return this
+    private fun FirFunctionCall.replaceCalleeReferenceWithOverridden(resolvedSymbol: FirNamedFunctionSymbol): FirFunctionCall {
         val overriddenSymbol = resolvedSymbol.overriddenSymbol ?: return this
         if (resolvedSymbol.callableId.isInvoke() || resolvedSymbol.callableId.isIteratorNext()) {
             return copy(
@@ -1090,7 +1109,7 @@ class Fir2IrVisitor(
         }
     }
 
-    private fun FirBlock.convertToIrExpressionOrBlock(): IrExpression {
+    private fun FirBlock.convertToIrExpressionOrBlock(origin: IrStatementOrigin? = null): IrExpression {
         if (statements.size == 1) {
             val firStatement = statements.single()
             if (firStatement is FirExpression) {
@@ -1101,7 +1120,7 @@ class Fir2IrVisitor(
             (statements.lastOrNull() as? FirExpression)?.typeRef?.toIrType(this@Fir2IrVisitor.session, declarationStorage) ?: unitType
         return convertWithOffsets { startOffset, endOffset ->
             IrBlockImpl(
-                startOffset, endOffset, type, null,
+                startOffset, endOffset, type, origin,
                 statements.mapNotNull { it.toIrStatement() }
             )
         }
@@ -1204,15 +1223,13 @@ class Fir2IrVisitor(
 
     override fun visitWhileLoop(whileLoop: FirWhileLoop, data: Any?): IrElement {
         return whileLoop.convertWithOffsets { startOffset, endOffset ->
-            IrWhileLoopImpl(
-                startOffset, endOffset, unitType,
-                if (whileLoop.psi is KtForExpression) IrStatementOrigin.FOR_LOOP_INNER_WHILE
-                else IrStatementOrigin.WHILE_LOOP
-            ).apply {
+            val origin = if (whileLoop.psi is KtForExpression) IrStatementOrigin.FOR_LOOP_INNER_WHILE
+            else IrStatementOrigin.WHILE_LOOP
+            IrWhileLoopImpl(startOffset, endOffset, unitType, origin).apply {
                 loopMap[whileLoop] = this
                 label = whileLoop.label?.name
                 condition = whileLoop.condition.toIrExpression()
-                body = whileLoop.block.convertToIrExpressionOrBlock()
+                body = whileLoop.block.convertToIrExpressionOrBlock(origin)
                 loopMap.remove(whileLoop)
             }
         }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/ResolveUtils.kt

@@ -483,4 +483,12 @@ fun CallableId.isKFunctionInvoke() =
             && packageName.asString() == "kotlin.reflect"
 
 fun CallableId.isIteratorNext() =
-    callableName.asString() == "next" && className?.asString()?.equals("Iterator") == true && packageName.asString() == "kotlin.collections"
+    callableName.asString() == "next" && className?.asString()?.endsWith("Iterator") == true
+            && packageName.asString() == "kotlin.collections"
+
+fun CallableId.isIteratorHasNext() =
+    callableName.asString() == "hasNext" && className?.asString()?.endsWith("Iterator") == true
+            && packageName.asString() == "kotlin.collections"
+
+fun CallableId.isIterator() =
+    callableName.asString() == "iterator" && packageName.asString() == "kotlin.collections"