Write loops parsing (light tree to FIR)

compiler/fir/lightTree/src/org/jetbrains/kotlin/fir/lightTree/converter/ConverterUtil.kt

@@ -203,6 +203,7 @@ object FunctionUtil {
     val firFunctions = mutableListOf<FirFunction>()
     val firFunctionCalls = mutableListOf<FirFunctionCall>()
     val firLabels = mutableListOf<FirLabel>()
+    val firLoops = mutableListOf<FirLoop>()
 
     fun <T> MutableList<T>.removeLast() {
         removeAt(size - 1)

compiler/fir/lightTree/src/org/jetbrains/kotlin/fir/lightTree/converter/DeclarationsConverter.kt

@@ -46,7 +46,6 @@ import org.jetbrains.kotlin.fir.types.impl.*
 import org.jetbrains.kotlin.lexer.KtModifierKeywordToken
 import org.jetbrains.kotlin.lexer.KtTokens.*
 import org.jetbrains.kotlin.name.FqName
-import org.jetbrains.kotlin.name.SpecialNames
 
 class DeclarationsConverter(
     private val session: FirSession,
@@ -557,7 +556,7 @@ class DeclarationsConverter(
         var firBlock: FirBlock? = null
         anonymousInitializer.forEachChildren {
             when (it.tokenType) {
-                BLOCK -> firBlock = visitBlock(it)
+                BLOCK -> firBlock = convertBlock(it)
             }
         }
 
@@ -601,7 +600,7 @@ class DeclarationsConverter(
         firConstructor.annotations += modifiers.annotations
         firConstructor.typeParameters += ConverterUtil.typeParametersFromSelfType(classWrapper.delegatedSelfTypeRef)
         firConstructor.valueParameters += firValueParameters.map { it.firValueParameter }
-        firConstructor.body = visitFunctionBody(block, null)
+        firConstructor.body = convertFunctionBody(block, null)
         FunctionUtil.firFunctions.removeLast()
         return firConstructor
     }
@@ -822,7 +821,7 @@ class DeclarationsConverter(
             firAccessor.valueParameters += firValueParameters
         }
 
-        firAccessor.body = visitFunctionBody(block, firExpression)
+        firAccessor.body = convertFunctionBody(block, firExpression)
         FunctionUtil.firFunctions.removeLast()
         return firAccessor
     }
@@ -901,7 +900,7 @@ class DeclarationsConverter(
                 null,
                 FirFunctionSymbol(ClassNameUtil.callableIdForName(functionName, isLocal)),
                 functionName,
-                modifiers.getVisibility(),
+                if (isLocal) Visibilities.LOCAL else modifiers.getVisibility(),
                 modifiers.getModality(),
                 modifiers.hasExpect(),
                 modifiers.hasActual(),
@@ -926,7 +925,7 @@ class DeclarationsConverter(
         }
 
         firFunction.valueParameters += valueParametersList.map { it.firValueParameter }
-        firFunction.body = visitFunctionBody(block, firExpression)
+        firFunction.body = convertFunctionBody(block, firExpression)
         FunctionUtil.firFunctions.removeLast()
         return firFunction
     }
@@ -935,10 +934,10 @@ class DeclarationsConverter(
      * @see org.jetbrains.kotlin.parsing.KotlinParsing.parseFunctionBody
      * @see org.jetbrains.kotlin.fir.builder.RawFirBuilder.Visitor.buildFirBody
      */
-    private fun visitFunctionBody(blockNode: LighterASTNode?, firExpression: FirExpression?): FirBlock? {
+    private fun convertFunctionBody(blockNode: LighterASTNode?, firExpression: FirExpression?): FirBlock? {
         return when {
             blockNode != null -> if (!stubMode) {
-                return visitBlock(blockNode)
+                return convertBlock(blockNode)
             } else {
                 FirSingleExpressionBlock(
                     session,
@@ -958,7 +957,7 @@ class DeclarationsConverter(
     /**
      * @see org.jetbrains.kotlin.parsing.KotlinParsing.parseBlock
      */
-    private fun visitBlock(block: LighterASTNode): FirBlock {
+    fun convertBlock(block: LighterASTNode): FirBlock {
         return if (!stubMode) {
             val blockTree = LightTree2Fir.buildLightTreeBlockExpression(block.getAsString())
             return DeclarationsConverter(session, stubMode, blockTree).convertBlockExpression(blockTree.root)
@@ -1266,7 +1265,7 @@ class DeclarationsConverter(
     /**
      * @see org.jetbrains.kotlin.parsing.KotlinParsing.parseValueParameter
      */
-    private fun convertValueParameter(valueParameter: LighterASTNode): ValueParameter {
+    fun convertValueParameter(valueParameter: LighterASTNode): ValueParameter {
         var modifiers = Modifier(session)
         var isVal = false
         var isVar = false

compiler/fir/lightTree/src/org/jetbrains/kotlin/fir/lightTree/converter/ExpressionsConverter.kt

@@ -41,7 +41,6 @@ import org.jetbrains.kotlin.fir.types.impl.FirImplicitTypeRefImpl
 import org.jetbrains.kotlin.ir.expressions.IrConstKind
 import org.jetbrains.kotlin.lexer.KtTokens.*
 import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.psi.KtLambdaExpression
 import org.jetbrains.kotlin.psi.stubs.elements.KtConstantExpressionElementType
 import org.jetbrains.kotlin.resolve.constants.evaluate.*
 import org.jetbrains.kotlin.types.expressions.OperatorConventions
@@ -81,6 +80,9 @@ class ExpressionsConverter(
                 STRING_TEMPLATE -> convertStringTemplate(expression)
                 is KtConstantExpressionElementType -> convertConstantExpression(expression)
                 REFERENCE_EXPRESSION -> convertSimpleNameExpression(expression)
+                DO_WHILE -> convertDoWhile(expression)
+                WHILE -> convertWhile(expression)
+                FOR -> convertFor(expression)
                 RETURN -> convertReturn(expression)
                 PARENTHESIZED, PROPERTY_DELEGATE, INDICES -> convertExpression(expression.getExpressionInParentheses())
                 THIS_EXPRESSION -> convertThisExpression(expression)
@@ -608,6 +610,146 @@ class ExpressionsConverter(
         }
     }
 
+    /**
+     * @see org.jetbrains.kotlin.parsing.KotlinExpressionParsing.parseDoWhile
+     * @see org.jetbrains.kotlin.fir.builder.RawFirBuilder.Visitor.visitDoWhileExpression
+     */
+    private fun convertDoWhile(doWhileLoop: LighterASTNode): FirElement {
+        var block: LighterASTNode? = null
+        var condition: FirExpression? = null
+        doWhileLoop.forEachChildren {
+            when (it.tokenType) {
+                BODY -> block = it
+                CONDITION -> condition = convertCondition(it)
+            }
+        }
+
+        return FirDoWhileLoopImpl(
+            session, null, condition ?: FirErrorExpressionImpl(session, null, "No condition in do-while loop")
+        ).configure { convertLoopBody(block) }
+    }
+
+    /**
+     * @see org.jetbrains.kotlin.parsing.KotlinExpressionParsing.parseWhile
+     * @see org.jetbrains.kotlin.fir.builder.RawFirBuilder.Visitor.visitWhileExpression
+     */
+    private fun convertWhile(whileLoop: LighterASTNode): FirElement {
+        var block: LighterASTNode? = null
+        var condition: FirExpression? = null
+        whileLoop.forEachChildren {
+            when (it.tokenType) {
+                BODY -> block = it
+                CONDITION -> condition = convertCondition(it)
+            }
+        }
+
+        return FirWhileLoopImpl(
+            session, null, condition ?: FirErrorExpressionImpl(session, null, "No condition in while loop")
+        ).configure { convertLoopBody(block) }
+    }
+
+    /**
+     * @see org.jetbrains.kotlin.parsing.KotlinExpressionParsing.parseFor
+     * @see org.jetbrains.kotlin.fir.builder.RawFirBuilder.Visitor.visitForExpression
+     */
+    private fun convertFor(forLoop: LighterASTNode): FirElement {
+        var parameter: ValueParameter? = null
+        var rangeExpression: FirExpression = FirErrorExpressionImpl(session, null, "No range in for loop")
+        var blockNode: LighterASTNode? = null
+        forLoop.forEachChildren {
+            when (it.tokenType) {
+                VALUE_PARAMETER -> parameter = declarationsConverter.convertValueParameter(it)
+                LOOP_RANGE -> convertCondition(it)?.apply { rangeExpression = this }
+                BODY -> blockNode = it
+            }
+        }
+
+        return FirBlockImpl(session, null).apply {
+            val rangeVal =
+                generateTemporaryVariable(this@ExpressionsConverter.session, null, Name.special("<range>"), rangeExpression)
+            statements += rangeVal
+            val iteratorVal = generateTemporaryVariable(
+                this@ExpressionsConverter.session, null, Name.special("<iterator>"),
+                FirFunctionCallImpl(this@ExpressionsConverter.session, null).apply {
+                    calleeReference = FirSimpleNamedReference(this@ExpressionsConverter.session, null, Name.identifier("iterator"))
+                    explicitReceiver = generateResolvedAccessExpression(this@ExpressionsConverter.session, null, rangeVal)
+                }
+            )
+            statements += iteratorVal
+            statements += FirWhileLoopImpl(
+                this@ExpressionsConverter.session, null,
+                FirFunctionCallImpl(this@ExpressionsConverter.session, null).apply {
+                    calleeReference = FirSimpleNamedReference(this@ExpressionsConverter.session, null, Name.identifier("hasNext"))
+                    explicitReceiver = generateResolvedAccessExpression(this@ExpressionsConverter.session, null, iteratorVal)
+                }
+            ).configure {
+                // NB: just body.toFirBlock() isn't acceptable here because we need to add some statements
+                val block = FirBlockImpl(this@ExpressionsConverter.session, null).apply {
+                    statements += convertLoopBody(blockNode).statements
+                }
+                if (parameter != null) {
+                    val multiDeclaration = parameter!!.destructuringDeclaration
+                    val firLoopParameter = generateTemporaryVariable(
+                        this@ExpressionsConverter.session, null,
+                        if (multiDeclaration != null) Name.special("<destruct>") else parameter!!.firValueParameter.name,
+                        FirFunctionCallImpl(this@ExpressionsConverter.session, null).apply {
+                            calleeReference = FirSimpleNamedReference(this@ExpressionsConverter.session, null, Name.identifier("next"))
+                            explicitReceiver = generateResolvedAccessExpression(this@ExpressionsConverter.session, null, iteratorVal)
+                        }
+                    )
+                    if (multiDeclaration != null) {
+                        val destructuringBlock = generateDestructuringBlock(
+                            this@ExpressionsConverter.session,
+                            multiDeclaration,
+                            firLoopParameter,
+                            tmpVariable = true
+                        )
+                        if (destructuringBlock is FirBlock) {
+                            for ((index, statement) in destructuringBlock.statements.withIndex()) {
+                                block.statements.add(index, statement)
+                            }
+                        }
+                    } else {
+                        block.statements.add(0, firLoopParameter)
+                    }
+                }
+                block
+            }
+        }
+    }
+
+    /**
+     * @see org.jetbrains.kotlin.parsing.KotlinExpressionParsing.parseLoopBody
+     * @see org.jetbrains.kotlin.fir.builder.RawFirBuilder.Visitor.toFirBlock
+     */
+    private fun convertLoopBody(body: LighterASTNode?): FirBlock {
+        var firBlock: FirBlock? = null
+        var firStatement: FirStatement? = null
+        body?.forEachChildren {
+            when (it.tokenType) {
+                BLOCK -> firBlock = declarationsConverter.convertBlockExpression(it)
+                else -> if (it.isExpression()) firStatement = getAsFirExpression(it)
+            }
+        }
+
+        return when {
+            firStatement != null -> FirSingleExpressionBlock(session, firStatement!!)
+            firBlock == null -> FirEmptyExpressionBlock(session)
+            else -> firBlock!!
+        }
+    }
+
+    /**
+     * @see org.jetbrains.kotlin.parsing.KotlinExpressionParsing.parseCondition
+     */
+    private fun convertCondition(condition: LighterASTNode): FirExpression? {
+        condition.forEachChildren {
+            if (it.isExpression()) return getAsFirExpression(it)
+        }
+
+        return null
+    }
+
     /**
      * @see org.jetbrains.kotlin.parsing.KotlinExpressionParsing.parseReturn
      * @see org.jetbrains.kotlin.fir.builder.RawFirBuilder.Visitor.visitReturnExpression

compiler/fir/lightTree/src/org/jetbrains/kotlin/fir/lightTree/converter/utils/ExpressionUtil.kt

@@ -16,7 +16,10 @@ import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.impl.*
 import org.jetbrains.kotlin.fir.lightTree.converter.ConverterUtil.getAsString
 import org.jetbrains.kotlin.fir.lightTree.converter.ConverterUtil.nameAsSafeName
+import org.jetbrains.kotlin.fir.lightTree.converter.FunctionUtil.pop
+import org.jetbrains.kotlin.fir.lightTree.converter.FunctionUtil.removeLast
 import org.jetbrains.kotlin.fir.lightTree.converter.ExpressionsConverter
+import org.jetbrains.kotlin.fir.lightTree.converter.FunctionUtil
 import org.jetbrains.kotlin.fir.lightTree.fir.DestructuringDeclaration
 import org.jetbrains.kotlin.fir.references.FirSimpleNamedReference
 import org.jetbrains.kotlin.fir.symbols.impl.FirVariableSymbol
@@ -194,3 +197,11 @@ private fun getSelectorType(qualifiedAccessChildren: Array<LighterASTNode?>): IE
     }
     return null
 }
+
+fun FirAbstractLoop.configure(generateBlock: () -> FirBlock): FirAbstractLoop {
+    label = FunctionUtil.firLabels.pop()
+    FunctionUtil.firLoops += this
+    block = generateBlock()
+    FunctionUtil.firLoops.removeLast()
+    return this
+}

compiler/fir/lightTree/tests/org/jetbrains/kotlin/fir/lightTree/compare/FirPartialTransformer.kt

@@ -31,7 +31,7 @@ class FirPartialTransformer(
     private val visitContinueExpression: Boolean = false,
     private val visitReturnExpression: Boolean = false,
     private val visitThrowExpression: Boolean = false,
-    private val visitForLoop: Boolean = false,
+    private val visitLoops: Boolean = false,
     private val visitConstExpression: Boolean = false,
     private val visitQualifiedAccessExpression: Boolean = false,
     private val visitCallableReferenceAccess: Boolean = false,
@@ -39,8 +39,6 @@ class FirPartialTransformer(
     private val visitWhenExpression: Boolean = false,
     private val visitLambdaArgumentExpression: Boolean = false,
     private val visitAnonymousObject: Boolean = false,
-    private val visitDoWhileLoop: Boolean = false,
-    private val visitWhileLoop: Boolean = false,
     private val visitAssignment: Boolean = false
 ) : FirAbstractTreeTransformer() {
     override fun transformAnonymousFunction(
@@ -167,7 +165,7 @@ class FirPartialTransformer(
     }
 
     override fun transformBlock(block: FirBlock, data: Nothing?): CompositeTransformResult<FirStatement> {
-        return if (!visitForLoop && block.psi is KtForExpression) {
+        return if (!visitLoops && block.psi is KtForExpression) {
             DummyFirStatement().compose()
         } else {
             super.transformBlock(block, data)
@@ -243,7 +241,7 @@ class FirPartialTransformer(
     }
 
     override fun transformDoWhileLoop(doWhileLoop: FirDoWhileLoop, data: Nothing?): CompositeTransformResult<FirStatement> {
-        return if (visitDoWhileLoop) {
+        return if (visitLoops) {
             (doWhileLoop.transformChildren(this, data) as FirDoWhileLoop).compose()
         } else {
             DummyFirStatement().compose()
@@ -251,7 +249,7 @@ class FirPartialTransformer(
     }
 
     override fun transformWhileLoop(whileLoop: FirWhileLoop, data: Nothing?): CompositeTransformResult<FirStatement> {
-        return if (visitWhileLoop) {
+        return if (visitLoops) {
             (whileLoop.transformChildren(this, data) as FirWhileLoop).compose()
         } else {
             DummyFirStatement().compose()

compiler/fir/lightTree/tests/org/jetbrains/kotlin/fir/lightTree/compare/TreesCompareTest.kt

@@ -71,7 +71,7 @@ class TreesCompareTest : AbstractRawFirBuilderTestCase() {
         visitContinueExpression: Boolean = false,
         visitReturnExpression: Boolean = false,
         visitThrowExpression: Boolean = false,
-        visitForLoop: Boolean = false,
+        visitLoops: Boolean = false,
         visitConstExpression: Boolean = false,
         visitQualifiedAccessExpression: Boolean = false,
         visitCallableReferenceAccess: Boolean = false,
@@ -79,8 +79,6 @@ class TreesCompareTest : AbstractRawFirBuilderTestCase() {
         visitWhenExpression: Boolean = false,
         visitLambdaArgumentExpression: Boolean = false,
         visitAnonymousObject: Boolean = false,
-        visitDoWhileLoop: Boolean = false,
-        visitWhileLoop: Boolean = false,
         visitAssignment: Boolean = false
     ) {
         val firVisitor = FirPartialTransformer(
@@ -98,7 +96,7 @@ class TreesCompareTest : AbstractRawFirBuilderTestCase() {
             visitContinueExpression = visitContinueExpression,
             visitReturnExpression = visitReturnExpression,
             visitThrowExpression = visitThrowExpression,
-            visitForLoop = visitForLoop,
+            visitLoops = visitLoops,
             visitConstExpression = visitConstExpression,
             visitQualifiedAccessExpression = visitQualifiedAccessExpression,
             visitCallableReferenceAccess = visitCallableReferenceAccess,
@@ -106,8 +104,6 @@ class TreesCompareTest : AbstractRawFirBuilderTestCase() {
             visitWhenExpression = visitWhenExpression,
             visitLambdaArgumentExpression = visitLambdaArgumentExpression,
             visitAnonymousObject = visitAnonymousObject,
-            visitDoWhileLoop = visitDoWhileLoop,
-            visitWhileLoop = visitWhileLoop,
             visitAssignment = visitAssignment
         )
         val lightTreeConverter = LightTree2Fir(stubMode, myProject)
@@ -279,4 +275,17 @@ class TreesCompareTest : AbstractRawFirBuilderTestCase() {
             visitAnonymousObject = true
         )
     }
+
+    fun testLoops() {
+        compare(
+            stubMode = false,
+            visitLambdaExpression = true,
+            visitAnnotation = true,
+            visitFunctionCall = true,
+            visitReturnExpression = true,
+            visitLoops = true,
+            visitConstExpression = true,
+            visitQualifiedAccessExpression = true
+        )
+    }
 }