[LT, Parsing] Parse blocks and lambdas eagerly for Light Tree

compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/LightTree2Fir.kt

@@ -6,7 +6,7 @@
 package org.jetbrains.kotlin.fir.lightTree
 
 import com.intellij.lang.LighterASTNode
-import com.intellij.lang.impl.PsiBuilderFactoryImpl
+import com.intellij.lang.PsiBuilderFactory
 import com.intellij.util.diff.FlyweightCapableTreeStructure
 import org.jetbrains.kotlin.KtIoFileSourceFile
 import org.jetbrains.kotlin.KtSourceFile
@@ -36,22 +36,9 @@ class LightTree2Fir(
         private val parserDefinition = KotlinParserDefinition()
         private fun makeLexer() = KotlinLexer()
 
-        fun buildLightTreeBlockExpression(code: String): FlyweightCapableTreeStructure<LighterASTNode> {
-            val builder = PsiBuilderFactoryImpl().createBuilder(parserDefinition, makeLexer(), code)
-            KotlinLightParser.parseBlockExpression(builder)
-            return builder.lightTree
-        }
-
-        fun buildLightTreeLambdaExpression(code: String): FlyweightCapableTreeStructure<LighterASTNode> {
-            val builder = PsiBuilderFactoryImpl().createBuilder(parserDefinition, makeLexer(), code)
-            KotlinLightParser.parseLambdaExpression(builder)
-            return builder.lightTree
-        }
-
         fun buildLightTree(code: CharSequence): FlyweightCapableTreeStructure<LighterASTNode> {
-            val builder = PsiBuilderFactoryImpl().createBuilder(parserDefinition, makeLexer(), code)
+            val builder = PsiBuilderFactory.getInstance().createBuilder(parserDefinition, makeLexer(), code)
-            KotlinLightParser.parse(builder)
+            return KotlinLightParser.parse(builder)
-            return builder.lightTree
         }
     }
 

compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/converter/BaseConverter.kt

@@ -27,15 +27,13 @@ abstract class BaseConverter(
     val tree: FlyweightCapableTreeStructure<LighterASTNode>,
     context: Context<LighterASTNode> = Context()
 ) : BaseFirBuilder<LighterASTNode>(baseSession, context) {
-    abstract val offset: Int
-
     protected val implicitType = FirImplicitTypeRefImplWithoutSource
 
     protected open fun reportSyntaxError(node: LighterASTNode) {}
 
     override fun LighterASTNode.toFirSourceElement(kind: KtFakeSourceElementKind?): KtLightSourceElement {
-        val startOffset = offset + tree.getStartOffset(this)
+        val startOffset = tree.getStartOffset(this)
-        val endOffset = offset + tree.getEndOffset(this)
+        val endOffset = tree.getEndOffset(this)
         return toKtLightSourceElement(tree, kind ?: context.forcedElementSourceKind ?: KtRealSourceElementKind, startOffset, endOffset)
     }
 

compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/converter/DeclarationsConverter.kt

@@ -69,23 +69,11 @@ class DeclarationsConverter(
     session: FirSession,
     internal val baseScopeProvider: FirScopeProvider,
     tree: FlyweightCapableTreeStructure<LighterASTNode>,
-    @set:PrivateForInline override var offset: Int = 0,
     context: Context<LighterASTNode> = Context(),
     private val diagnosticsReporter: DiagnosticReporter? = null,
-    private val diagnosticContext: DiagnosticContext? = null
+    private val diagnosticContext: DiagnosticContext? = null,
 ) : BaseConverter(session, tree, context) {
 
-    @OptIn(PrivateForInline::class)
-    inline fun <R> withOffset(newOffset: Int, block: () -> R): R {
-        val oldOffset = offset
-        offset = newOffset
-        return try {
-            block()
-        } finally {
-            offset = oldOffset
-        }
-    }
-
     private val expressionConverter = ExpressionsConverter(session, tree, this, context)
 
     override fun reportSyntaxError(node: LighterASTNode) {
@@ -1809,11 +1797,7 @@ class DeclarationsConverter(
             )
         }
 
-        val blockTree = LightTree2Fir.buildLightTreeBlockExpression(block.asText)
+        return convertBlockExpression(block)
-        return DeclarationsConverter(
-            baseSession, baseScopeProvider, blockTree, offset = offset + tree.getStartOffset(block), context,
-            diagnosticsReporter, diagnosticContext
-        ).convertBlockExpression(blockTree.root)
     }
 
     /**

compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/converter/ExpressionsConverter.kt

@@ -58,10 +58,8 @@ class ExpressionsConverter(
     session: FirSession,
     tree: FlyweightCapableTreeStructure<LighterASTNode>,
     private val declarationsConverter: DeclarationsConverter,
-    context: Context<LighterASTNode> = Context()
+    context: Context<LighterASTNode> = Context(),
 ) : BaseConverter(session, tree, context) {
-    override val offset: Int
-        get() = declarationsConverter.offset
 
     inline fun <reified R : FirElement> getAsFirExpression(expression: LighterASTNode?, errorReason: String = ""): R {
         val converted = expression?.let { convertExpression(it, errorReason) }
@@ -76,17 +74,7 @@ class ExpressionsConverter(
     /*****    EXPRESSIONS    *****/
     fun convertExpression(expression: LighterASTNode, errorReason: String): FirElement {
         return when (expression.tokenType) {
-            LAMBDA_EXPRESSION -> {
+            LAMBDA_EXPRESSION -> convertLambdaExpression(expression)
-                val lambdaTree = LightTree2Fir.buildLightTreeLambdaExpression(expression.asText)
-                // Pass on label user to the lambda root
-                context.forwardLabelUsagePermission(expression, lambdaTree.root)
-                val lambdaDeclarationsConverter = DeclarationsConverter(
-                    baseSession, declarationsConverter.baseScopeProvider, lambdaTree,
-                    offset = offset + expression.startOffset, context
-                )
-                ExpressionsConverter(baseSession, lambdaTree, lambdaDeclarationsConverter, context)
-                    .convertLambdaExpression(lambdaTree.root)
-            }
             BINARY_EXPRESSION -> convertBinaryExpression(expression)
             BINARY_WITH_TYPE -> convertBinaryWithTypeRHSExpression(expression) {
                 this.getOperationSymbol().toFirOperation()
@@ -200,29 +188,27 @@ class ExpressionsConverter(
             }
 
             body = if (block != null) {
-                declarationsConverter.withOffset(expressionSource.startOffset) {
+                declarationsConverter.convertBlockExpressionWithoutBuilding(block!!).apply {
-                    declarationsConverter.convertBlockExpressionWithoutBuilding(block!!).apply {
+                    statements.firstOrNull()?.let {
-                        statements.firstOrNull()?.let {
+                        if (it.isContractBlockFirCheck()) {
-                            if (it.isContractBlockFirCheck()) {
+                            this@buildAnonymousFunction.contractDescription = it.toLegacyRawContractDescription()
-                                this@buildAnonymousFunction.contractDescription = it.toLegacyRawContractDescription()
+                            statements[0] = FirContractCallBlock(it)
-                                statements[0] = FirContractCallBlock(it)
-                            }
                         }
+                    }
 
-                        if (statements.isEmpty()) {
+                    if (statements.isEmpty()) {
-                            statements.add(
+                        statements.add(
-                                buildReturnExpression {
+                            buildReturnExpression {
-                                    source = expressionSource.fakeElement(KtFakeSourceElementKind.ImplicitReturn.FromExpressionBody)
+                                source = expressionSource.fakeElement(KtFakeSourceElementKind.ImplicitReturn.FromExpressionBody)
-                                    this.target = target
+                                this.target = target
-                                    result = buildUnitExpression {
+                                result = buildUnitExpression {
-                                        source = expressionSource.fakeElement(KtFakeSourceElementKind.ImplicitUnit)
+                                    source = expressionSource.fakeElement(KtFakeSourceElementKind.ImplicitUnit)
-                                    }
                                 }
-                            )
+                            }
-                        }
+                        )
-                        statements.addAll(0, destructuringStatements)
+                    }
-                    }.build()
+                    statements.addAll(0, destructuringStatements)
-                }
+                }.build()
             } else {
                 buildSingleExpressionBlock(buildErrorExpression(null, ConeSimpleDiagnostic("Lambda has no body", DiagnosticKind.Syntax)))
             }

compiler/psi/src/org/jetbrains/kotlin/parsing/AbstractKotlinParsing.java

@@ -51,9 +51,15 @@ import static org.jetbrains.kotlin.lexer.KtTokens.*;
     }
 
     protected final SemanticWhitespaceAwarePsiBuilder myBuilder;
+    protected final boolean isLazy;
 
     public AbstractKotlinParsing(SemanticWhitespaceAwarePsiBuilder builder) {
+        this(builder, true);
+    }
+
+    public AbstractKotlinParsing(SemanticWhitespaceAwarePsiBuilder builder, boolean isLazy) {
         this.myBuilder = builder;
+        this.isLazy = isLazy;
     }
 
     protected IElementType getLastToken() {

compiler/psi/src/org/jetbrains/kotlin/parsing/KotlinExpressionParsing.java

@@ -280,7 +280,11 @@ public class KotlinExpressionParsing extends AbstractKotlinParsing {
     private final KotlinParsing myKotlinParsing;
 
     public KotlinExpressionParsing(SemanticWhitespaceAwarePsiBuilder builder, KotlinParsing kotlinParsing) {
-        super(builder);
+        this(builder, kotlinParsing, true);
+    }
+
+    public KotlinExpressionParsing(SemanticWhitespaceAwarePsiBuilder builder, KotlinParsing kotlinParsing, boolean isLazy) {
+        super(builder, isLazy);
         myKotlinParsing = kotlinParsing;
     }
 
@@ -1191,7 +1195,7 @@ public class KotlinExpressionParsing extends AbstractKotlinParsing {
             return;
         }
 
-        if (collapse) {
+        if (collapse && isLazy) {
             myKotlinParsing.advanceBalancedBlock();
             literal.done(FUNCTION_LITERAL);
             literalExpression.collapse(LAMBDA_EXPRESSION);

compiler/psi/src/org/jetbrains/kotlin/parsing/KotlinLightParser.java

@@ -11,21 +11,9 @@ import com.intellij.util.diff.FlyweightCapableTreeStructure;
 
 public class KotlinLightParser {
     public static FlyweightCapableTreeStructure<LighterASTNode> parse(PsiBuilder builder) {
-        KotlinParsing ktParsing = KotlinParsing.createForTopLevel(new SemanticWhitespaceAwarePsiBuilderImpl(builder));
+        KotlinParsing ktParsing = KotlinParsing.createForTopLevelNonLazy(new SemanticWhitespaceAwarePsiBuilderImpl(builder));
         ktParsing.parseFile();
 
         return builder.getLightTree();
     }
-
-    public static FlyweightCapableTreeStructure<LighterASTNode> parseLambdaExpression(PsiBuilder psiBuilder) {
-        KotlinParsing ktParsing = KotlinParsing.createForTopLevel(new SemanticWhitespaceAwarePsiBuilderImpl(psiBuilder));
-        ktParsing.parseLambdaExpression();
-        return psiBuilder.getLightTree();
-    }
-
-    public static FlyweightCapableTreeStructure<LighterASTNode> parseBlockExpression(PsiBuilder builder) {
-        KotlinParsing ktParsing = KotlinParsing.createForTopLevel(new SemanticWhitespaceAwarePsiBuilderImpl(builder));
-        ktParsing.parseBlockExpression();
-        return builder.getLightTree();
-    }
 }

compiler/psi/src/org/jetbrains/kotlin/parsing/KotlinParsing.java

@@ -106,11 +106,15 @@ public class KotlinParsing extends AbstractKotlinParsing {
     private final static TokenSet CLASS_INTERFACE_SET = TokenSet.create(CLASS_KEYWORD, INTERFACE_KEYWORD);
 
     static KotlinParsing createForTopLevel(SemanticWhitespaceAwarePsiBuilder builder) {
-        return new KotlinParsing(builder, true);
+        return new KotlinParsing(builder, true, true);
     }
 
-    private static KotlinParsing createForByClause(SemanticWhitespaceAwarePsiBuilder builder) {
+    static KotlinParsing createForTopLevelNonLazy(SemanticWhitespaceAwarePsiBuilder builder) {
-        return new KotlinParsing(new SemanticWhitespaceAwarePsiBuilderForByClause(builder), false);
+        return new KotlinParsing(builder, true, false);
+    }
+
+    private static KotlinParsing createForByClause(SemanticWhitespaceAwarePsiBuilder builder, boolean isLazy) {
+        return new KotlinParsing(new SemanticWhitespaceAwarePsiBuilderForByClause(builder), false, isLazy);
     }
 
     private final KotlinExpressionParsing myExpressionParsing;
@@ -142,11 +146,11 @@ public class KotlinParsing extends AbstractKotlinParsing {
                 }
             });
 
-    private KotlinParsing(SemanticWhitespaceAwarePsiBuilder builder, boolean isTopLevel) {
+    private KotlinParsing(SemanticWhitespaceAwarePsiBuilder builder, boolean isTopLevel, boolean isLazy) {
-        super(builder);
+        super(builder, isLazy);
         myExpressionParsing = isTopLevel
-                              ? new KotlinExpressionParsing(builder, this)
+                              ? new KotlinExpressionParsing(builder, this, isLazy)
-                              : new KotlinExpressionParsing(builder, this) {
+                              : new KotlinExpressionParsing(builder, this, isLazy) {
                                   @Override
                                   protected boolean parseCallWithClosure() {
                                       if (((SemanticWhitespaceAwarePsiBuilderForByClause) builder).getStackSize() > 0) {
@@ -157,7 +161,7 @@ public class KotlinParsing extends AbstractKotlinParsing {
 
                                   @Override
                                   protected KotlinParsing create(SemanticWhitespaceAwarePsiBuilder builder) {
-                                      return createForByClause(builder);
+                                      return createForByClause(builder, isLazy);
                                   }
                               };
     }
@@ -1914,13 +1918,12 @@ public class KotlinParsing extends AbstractKotlinParsing {
     }
 
     private void parseBlock(boolean collapse) {
-
         PsiBuilder.Marker lazyBlock = mark();
 
         myBuilder.enableNewlines();
 
         boolean hasOpeningBrace = expect(LBRACE, "Expecting '{' to open a block");
-        boolean canCollapse = collapse && hasOpeningBrace;
+        boolean canCollapse = collapse && hasOpeningBrace && isLazy;
 
         if (canCollapse) {
             advanceBalancedBlock();
@@ -1997,7 +2000,7 @@ public class KotlinParsing extends AbstractKotlinParsing {
         if (at(BY_KEYWORD)) {
             reference.drop();
             advance(); // BY_KEYWORD
-            createForByClause(myBuilder).myExpressionParsing.parseExpression();
+            createForByClause(myBuilder, isLazy).myExpressionParsing.parseExpression();
             delegator.done(DELEGATED_SUPER_TYPE_ENTRY);
         }
         else if (at(LPAR)) {