From 5818f0786af82d1986b71a114956e8f610e36b23 Mon Sep 17 00:00:00 2001
From: Ilya Chernikov <ilya.chernikov@jetbrains.com>
Date: Wed, 3 Nov 2021 14:05:25 +0100
Subject: [PATCH] Fix for loop variable detection in fir2ir for LT case

---
 .../jetbrains/kotlin/fir/backend/Fir2IrVisitor.kt    | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrVisitor.kt b/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrVisitor.kt
index 861f40989c4..4b420253d75 100644
--- a/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrVisitor.kt
+++ b/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrVisitor.kt
@@ -5,9 +5,7 @@
 
 package org.jetbrains.kotlin.fir.backend
 
-import org.jetbrains.kotlin.KtFakeSourceElementKind
-import org.jetbrains.kotlin.KtNodeTypes
-import org.jetbrains.kotlin.KtSourceElement
+import org.jetbrains.kotlin.*
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.Visibilities
@@ -49,7 +47,6 @@ import org.jetbrains.kotlin.ir.util.parentClassOrNull
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.name.SpecialNames
-import org.jetbrains.kotlin.psi
 import org.jetbrains.kotlin.psi.KtBinaryExpression
 import org.jetbrains.kotlin.psi.KtForExpression
 import org.jetbrains.kotlin.util.OperatorNameConventions
@@ -284,7 +281,7 @@ class Fir2IrVisitor(
         val initializer = variable.initializer
         val isNextVariable = initializer is FirFunctionCall &&
                 (initializer.calleeReference.resolvedSymbol as? FirNamedFunctionSymbol)?.callableId?.isIteratorNext() == true &&
-                variable.source.psi?.parent is KtForExpression
+                variable.source?.isChildOfForLoop == true
         val irVariable = declarationStorage.createIrVariable(
             variable, conversionScope.parentFromStack(),
             if (isNextVariable) {
@@ -1198,3 +1195,8 @@ class Fir2IrVisitor(
         }
     }
 }
+
+val KtSourceElement.isChildOfForLoop: Boolean
+    get() =
+        if (this is KtPsiSourceElement) psi.parent is KtForExpression
+        else treeStructure.getParent(lighterASTNode)?.tokenType == KtNodeTypes.FOR
\ No newline at end of file