Reformat LightTree2Fir class to be able to parse block of code

compiler/fir/lightTree/src/org/jetbrains/kotlin/fir/lightTree/LightTree2Fir.kt

@@ -9,21 +9,45 @@ import com.intellij.lang.LighterASTNode
 import com.intellij.lang.ParserDefinition
 import com.intellij.lang.impl.PsiBuilderFactoryImpl
 import com.intellij.lexer.Lexer
+import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.io.FileUtil
 import com.intellij.openapi.vfs.CharsetToolkit
+import com.intellij.psi.DummyHolderViewProvider
+import com.intellij.psi.PsiManager
+import com.intellij.psi.SingleRootFileViewProvider
+import com.intellij.psi.impl.PsiManagerEx
+import com.intellij.testFramework.LightVirtualFile
 import com.intellij.util.diff.FlyweightCapableTreeStructure
 import org.jetbrains.kotlin.fir.FirSessionBase
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.lightTree.converter.DeclarationsConverter
+import org.jetbrains.kotlin.idea.KotlinFileType
+import org.jetbrains.kotlin.lexer.KotlinLexer
+import org.jetbrains.kotlin.parsing.KotlinParser
+import org.jetbrains.kotlin.parsing.KotlinParserDefinition
 import org.jetbrains.kotlin.parsing.MyKotlinParser
+import org.jetbrains.kotlin.psi.KtFile
 import java.io.File
 import java.nio.file.Path
 
 class LightTree2Fir(
     private val stubMode: Boolean,
-    private val parserDefinition: ParserDefinition,
-    private val lexer: Lexer
+    private val project: Project
 ) {
+    private val ktDummyFile = KtFile(SingleRootFileViewProvider(PsiManager.getInstance(project), LightVirtualFile()), false)
+
+    companion object {
+        private val parserDefinition = KotlinParserDefinition()
+        private val lexer = KotlinLexer()
+
+        fun buildLightTreeBlockExpression(code: String): FlyweightCapableTreeStructure<LighterASTNode> {
+            val builder = PsiBuilderFactoryImpl().createBuilder(parserDefinition, lexer, code)
+            //KotlinParser.parseBlockExpression(builder)
+            MyKotlinParser.parseBlockExpression(builder)
+            return builder.lightTree
+        }
+    }
+
     fun buildFirFile(path: Path): FirFile {
         return buildFirFile(path.toFile())
     }
@@ -35,6 +59,7 @@ class LightTree2Fir(
 
     fun buildLightTree(code: String): FlyweightCapableTreeStructure<LighterASTNode> {
         val builder = PsiBuilderFactoryImpl().createBuilder(parserDefinition, lexer, code)
+        //KotlinParser(project).parse(null, builder, ktDummyFile)
         MyKotlinParser.parse(builder)
         return builder.lightTree
     }

compiler/fir/lightTree/tests/org/jetbrains/kotlin/fir/lightTree/LightTree2FirConverterTestCases.kt

@@ -21,9 +21,7 @@ class LightTree2FirConverterTestCases : AbstractRawFirBuilderTestCase() {
     private val testDirPath = "compiler/fir/psi2fir/testData/rawBuilder/declarations"
 
     private fun executeTest(filePath: String) {
-        val parserDefinition = KotlinParserDefinition()
-        val lexer = parserDefinition.createLexer(myProject)
-        val lightTree2Fir = LightTree2Fir(true, parserDefinition, lexer).buildFirFile(Paths.get(filePath)).render()
+        val lightTree2Fir = LightTree2Fir(true, myProject).buildFirFile(Paths.get(filePath)).render()
 
         val file = createKtFile(filePath)
         val firFile = file.toFirFile(stubMode = true)

compiler/fir/lightTree/tests/org/jetbrains/kotlin/fir/lightTree/TotalKotlinTest.kt

@@ -47,9 +47,7 @@ class TotalKotlinTest : AbstractRawFirBuilderTestCase() {
         var counter = 0
         var time = 0L
 
-        val parserDefinition = KotlinParserDefinition()
-        val lexer = parserDefinition.createLexer(myProject)
-        val lightTreeConverter = LightTree2Fir(true, parserDefinition, lexer)
+        val lightTreeConverter = LightTree2Fir(true, myProject)
 
         if (onlyLightTree) println("LightTree generation") else println("Fir from LightTree converter")
         println("BASE PATH: $path")

compiler/fir/lightTree/tests/org/jetbrains/kotlin/fir/lightTree/benchmark/generators/LightTree2FirGenerator.kt

@@ -29,9 +29,7 @@ open class LightTree2FirGenerator : TreeGenerator, AbstractRawFirBuilderTestCase
     }
 
     private fun createConverter() {
-        val parserDefinition = KotlinParserDefinition()
-        val lexer = parserDefinition.createLexer(project)
-        lightTreeConverter = LightTree2Fir(true, parserDefinition, lexer)
+        lightTreeConverter = LightTree2Fir(true, myProject)
     }
 
     override fun setUp() {

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

@@ -15,9 +15,11 @@ import org.jetbrains.kotlin.fir.resolve.transformers.FirAbstractTreeTransformer
 import org.jetbrains.kotlin.fir.visitors.CompositeTransformResult
 import org.jetbrains.kotlin.fir.visitors.compose
 import org.jetbrains.kotlin.psi.KtForExpression
+import org.jetbrains.kotlin.psi.KtLambdaExpression
 
 class FirPartialTransformer(
     private val visitAnonymousFunction: Boolean = false,
+    private val visitLambdaExpression: Boolean = false,
     private val visitNamedFunction: Boolean = false,
     private val visitMemberDeclaration: Boolean = false,
     private val visitVariable: Boolean = false,
@@ -40,6 +42,7 @@ class FirPartialTransformer(
     private val visitTryExpression: Boolean = false,
     private val visitWhenExpression: Boolean = false,
     private val visitNamedArgumentExpression: Boolean = false,
+    private val visitLambdaArgumentExpression: Boolean = false,
     private val visitSpreadArgumentExpression: Boolean = false,
     private val visitAnonymousObject: Boolean = false,
     private val visitDoWhileLoop: Boolean = false,
@@ -52,6 +55,8 @@ class FirPartialTransformer(
     ): CompositeTransformResult<FirDeclaration> {
         return if (visitAnonymousFunction) {
             (anonymousFunction.transformChildren(this, data) as FirAnonymousFunction).compose()
+        } else if (visitLambdaExpression && anonymousFunction.psi is KtLambdaExpression) {
+            (anonymousFunction.transformChildren(this, data) as FirAnonymousFunction).compose()
         } else {
             DummyFirAnonymousFunction().compose()
         }
@@ -265,6 +270,17 @@ class FirPartialTransformer(
         }
     }
 
+    override fun transformLambdaArgumentExpression(
+        lambdaArgumentExpression: FirLambdaArgumentExpression,
+        data: Nothing?
+    ): CompositeTransformResult<FirStatement> {
+        return if (visitLambdaArgumentExpression) {
+            (lambdaArgumentExpression.transformChildren(this, data) as FirLambdaArgumentExpression).compose()
+        } else {
+            DummyFirStatement().compose()
+        }
+    }
+
     override fun transformSpreadArgumentExpression(
         spreadArgumentExpression: FirSpreadArgumentExpression,
         data: Nothing?

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

@@ -23,18 +23,14 @@ import java.io.File
 @TestDataPath("\$PROJECT_ROOT")
 @RunWith(JUnit3RunnerWithInners::class)
 class TreesCompareTest : AbstractRawFirBuilderTestCase() {
-    private fun compareBase(stubMode: Boolean, compareFir: (File, LightTree2Fir) -> Boolean) {
+    private fun compareBase(stubMode: Boolean, compareFir: (File) -> Boolean) {
         val path = System.getProperty("user.dir")
         var counter = 0
         var errorCounter = 0
 
-        val parserDefinition = KotlinParserDefinition()
-        val lexer = parserDefinition.createLexer(myProject)
-        val lightTreeConverter = LightTree2Fir(stubMode, parserDefinition, lexer)
-
         println("BASE PATH: $path")
         path.walkTopDown { file ->
-            if (!compareFir(file, lightTreeConverter)) errorCounter++
+            if (!compareFir(file)) errorCounter++
             counter++
         }
         println("All scanned files: $counter")
@@ -43,7 +39,8 @@ class TreesCompareTest : AbstractRawFirBuilderTestCase() {
     }
 
     private fun compareAll(stubMode: Boolean) {
-        compareBase(stubMode) { file, lightTreeConverter ->
+        val lightTreeConverter = LightTree2Fir(stubMode, myProject)
+        compareBase(stubMode) { file ->
             val text = FileUtil.loadFile(file, CharsetToolkit.UTF8, true).trim()
 
             //light tree
@@ -62,6 +59,7 @@ class TreesCompareTest : AbstractRawFirBuilderTestCase() {
     private fun compare(
         stubMode: Boolean,
         visitAnonymousFunction: Boolean = false,
+        visitLambdaExpression: Boolean = false,
         visitNamedFunction: Boolean = false,
         visitMemberDeclaration: Boolean = false,
         visitVariable: Boolean = false,
@@ -84,6 +82,7 @@ class TreesCompareTest : AbstractRawFirBuilderTestCase() {
         visitTryExpression: Boolean = false,
         visitWhenExpression: Boolean = false,
         visitNamedArgumentExpression: Boolean = false,
+        visitLambdaArgumentExpression: Boolean = false,
         visitSpreadArgumentExpression: Boolean = false,
         visitAnonymousObject: Boolean = false,
         visitDoWhileLoop: Boolean = false,
@@ -92,6 +91,7 @@ class TreesCompareTest : AbstractRawFirBuilderTestCase() {
     ) {
         val firVisitor = FirPartialTransformer(
             visitAnonymousFunction = visitAnonymousFunction,
+            visitLambdaExpression = visitLambdaExpression,
             visitNamedFunction = visitNamedFunction,
             visitMemberDeclaration = visitMemberDeclaration,
             visitVariable = visitVariable,
@@ -114,13 +114,15 @@ class TreesCompareTest : AbstractRawFirBuilderTestCase() {
             visitTryExpression = visitTryExpression,
             visitWhenExpression = visitWhenExpression,
             visitNamedArgumentExpression = visitNamedArgumentExpression,
+            visitLambdaArgumentExpression = visitLambdaArgumentExpression,
             visitSpreadArgumentExpression = visitSpreadArgumentExpression,
             visitAnonymousObject = visitAnonymousObject,
             visitDoWhileLoop = visitDoWhileLoop,
             visitWhileLoop = visitWhileLoop,
             visitAssignment = visitAssignment
         )
-        compareBase(stubMode) { file, lightTreeConverter ->
+        val lightTreeConverter = LightTree2Fir(stubMode, myProject)
+        compareBase(stubMode) { file ->
             val text = FileUtil.loadFile(file, CharsetToolkit.UTF8, true).trim()
 
             //light tree
@@ -153,4 +155,8 @@ class TreesCompareTest : AbstractRawFirBuilderTestCase() {
     fun testConstExpression() {
         compare(stubMode = false, visitConstExpression = true, visitAnnotation = true)
     }
+
+    fun testCallExpression() {
+        compare(stubMode = false, visitFunctionCall = true, visitLambdaExpression = false)
+    }
 }

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

@@ -16,4 +16,10 @@ public class MyKotlinParser {
 
         return builder.getLightTree();
     }
+
+    public static FlyweightCapableTreeStructure<LighterASTNode> parseBlockExpression(PsiBuilder builder) {
+        KotlinParsing ktParsing = KotlinParsing.createForTopLevel(new SemanticWhitespaceAwarePsiBuilderImpl(builder));
+        ktParsing.parseBlockExpression();
+        return builder.getLightTree();
+    }
 }