[LL FIR] transform script body properly

Problem: transformSingle for expressions can return another expression,
and we should replace the original one with such new expression

^KT-61011 Fixed
^KT-61009 Fixed

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirAbstractBodyTargetResolver.kt

@@ -11,14 +11,17 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.LLFirRetu
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.LLFirLockProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.FirLazyBodiesCalculator
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.isScriptStatement
+import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirElementWithResolveState
 import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.resolve.ResolutionMode
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculator
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirAbstractBodyResolveTransformerDispatcher
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirResolveContextCollector
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.ImplicitBodyResolveComputationSession
+import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.transformSingle
 
 internal abstract class LLFirAbstractBodyTargetResolver(
@@ -27,7 +30,7 @@ internal abstract class LLFirAbstractBodyTargetResolver(
     private val scopeSession: ScopeSession,
     resolvePhase: FirResolvePhase,
     protected val implicitBodyResolveComputationSession: ImplicitBodyResolveComputationSession = ImplicitBodyResolveComputationSession(),
-    isJumpingPhase: Boolean = false
+    isJumpingPhase: Boolean = false,
 ) : LLFirTargetResolver(resolveTarget, lockProvider, resolvePhase, isJumpingPhase) {
     protected fun createReturnTypeCalculator(
         firResolveContextCollector: FirResolveContextCollector?,
@@ -78,12 +81,17 @@ internal abstract class LLFirAbstractBodyTargetResolver(
     protected fun resolveScript(script: FirScript) {
         transformer.declarationsTransformer?.withScript(script) {
             script.parameters.forEach { it.transformSingle(transformer, ResolutionMode.ContextIndependent) }
-            script.statements.forEach {
+            script.transformStatements(
-                if (it.isScriptStatement) {
+                transformer = object : FirTransformer<Any?>() {
-                    transformer.firResolveContextCollector?.addStatementContext(it, transformer.context)
+                    override fun <E : FirElement> transformElement(element: E, data: Any?): E {
-                    it.transformSingle(transformer, ResolutionMode.ContextIndependent)
+                        if (element !is FirStatement || !element.isScriptStatement) return element
-                }
+
-            }
+                        transformer.firResolveContextCollector?.addStatementContext(element, transformer.context)
+                        return element.transformSingle(transformer, ResolutionMode.ContextIndependent)
+                    }
+                },
+                data = null,
+            )
 
             script
         }

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirBodyLazyResolver.kt

@@ -74,6 +74,7 @@ internal object LLFirBodyLazyResolver : LLFirLazyResolver(FirResolvePhase.BODY_R
                 checkBodyIsResolved(target)
             }
             is FirFunction -> checkBodyIsResolved(target)
+            is FirScript -> checkStatementsAreResolved(target)
         }
     }
 }

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/firCheckResolvedUtils.kt

@@ -64,6 +64,16 @@ internal fun checkBodyIsResolved(function: FirFunction) {
     }
 }
 
+internal fun checkStatementsAreResolved(script: FirScript) {
+    for (statement in script.statements) {
+        if (statement.isScriptStatement && statement is FirExpression) {
+            checkTypeRefIsResolved(statement.typeRef, "script statement", script) {
+                withFirEntry("expression", statement)
+            }
+        }
+    }
+}
+
 internal fun checkExpectForActualIsResolved(memberDeclaration: FirMemberDeclaration) {
     if (!memberDeclaration.isActual) return
 

analysis/low-level-api-fir/testdata/getOrBuildFir/calls/compoundAssignOnVarScript.txt

@@ -11,7 +11,7 @@ SCRIPT: [ResolvedTo(BODY_RESOLVE)]
         public [ResolvedTo(IMPLICIT_TYPES_BODY_RESOLVE)] get(): R|kotlin/Int|
         public [ResolvedTo(IMPLICIT_TYPES_BODY_RESOLVE)] set([ResolvedTo(IMPLICIT_TYPES_BODY_RESOLVE)] value: R|kotlin/Int|): R|kotlin/Unit|
 
-    +=(R|/i|, Int(1))
+    R|/i| = R|/i|.R|kotlin/Int.plus|(Int(1))
 
 FIR FILE:
 FILE: [ResolvedTo(IMPORTS)] compoundAssignOnVarScript.kts
@@ -23,4 +23,4 @@ FILE: [ResolvedTo(IMPORTS)] compoundAssignOnVarScript.kts
             public [ResolvedTo(IMPLICIT_TYPES_BODY_RESOLVE)] get(): R|kotlin/Int|
             public [ResolvedTo(IMPLICIT_TYPES_BODY_RESOLVE)] set([ResolvedTo(IMPLICIT_TYPES_BODY_RESOLVE)] value: R|kotlin/Int|): R|kotlin/Unit|
 
-        +=(R|/i|, Int(1))
+        R|/i| = R|/i|.R|kotlin/Int.plus|(Int(1))

analysis/low-level-api-fir/testdata/lazyResolve/statementTransformation.kts

@@ -0,0 +1,9 @@
+// RESOLVE_SCRIPT
+
+val myProp: () -> Unit = {}
+
+myProp {
+    123
+}
+
+fun foo() {}

analysis/low-level-api-fir/testdata/lazyResolve/statementTransformation.txt

@@ -0,0 +1,193 @@
+RAW_FIR:
+FILE: [ResolvedTo(RAW_FIR)] statementTransformation.kts
+    context(<script>@kotlin.script.templates.standard.ScriptTemplateWithArgs)
+    SCRIPT: [ResolvedTo(RAW_FIR)]
+        [ResolvedTo(RAW_FIR)] lval args: R|kotlin/Array<kotlin/String>|
+
+        public? final? [ResolvedTo(RAW_FIR)] val myProp: ( () -> Unit ) = LAZY_EXPRESSION
+            public? [ResolvedTo(RAW_FIR)] get(): ( () -> Unit )
+
+        myProp#(<L> = LAZY_EXPRESSION)
+        public? final? [ResolvedTo(RAW_FIR)] fun foo(): R|kotlin/Unit| { LAZY_BLOCK }
+
+
+IMPORTS:
+FILE: [ResolvedTo(IMPORTS)] statementTransformation.kts
+    context(<script>@kotlin.script.templates.standard.ScriptTemplateWithArgs)
+    SCRIPT: [ResolvedTo(RAW_FIR)]
+        [ResolvedTo(RAW_FIR)] lval args: R|kotlin/Array<kotlin/String>|
+
+        public? final? [ResolvedTo(RAW_FIR)] val myProp: ( () -> Unit ) = LAZY_EXPRESSION
+            public? [ResolvedTo(RAW_FIR)] get(): ( () -> Unit )
+
+        myProp#(<L> = LAZY_EXPRESSION)
+        public? final? [ResolvedTo(RAW_FIR)] fun foo(): R|kotlin/Unit| { LAZY_BLOCK }
+
+
+COMPILER_REQUIRED_ANNOTATIONS:
+FILE: [ResolvedTo(IMPORTS)] statementTransformation.kts
+    context(<script>@kotlin.script.templates.standard.ScriptTemplateWithArgs)
+    SCRIPT: [ResolvedTo(COMPILER_REQUIRED_ANNOTATIONS)]
+        [ResolvedTo(RAW_FIR)] lval args: R|kotlin/Array<kotlin/String>|
+
+        public? final? [ResolvedTo(RAW_FIR)] val myProp: ( () -> Unit ) = LAZY_EXPRESSION
+            public? [ResolvedTo(RAW_FIR)] get(): ( () -> Unit )
+
+        myProp#(<L> = LAZY_EXPRESSION)
+        public? final? [ResolvedTo(RAW_FIR)] fun foo(): R|kotlin/Unit| { LAZY_BLOCK }
+
+
+COMPANION_GENERATION:
+FILE: [ResolvedTo(IMPORTS)] statementTransformation.kts
+    context(<script>@kotlin.script.templates.standard.ScriptTemplateWithArgs)
+    SCRIPT: [ResolvedTo(COMPANION_GENERATION)]
+        [ResolvedTo(RAW_FIR)] lval args: R|kotlin/Array<kotlin/String>|
+
+        public? final? [ResolvedTo(RAW_FIR)] val myProp: ( () -> Unit ) = LAZY_EXPRESSION
+            public? [ResolvedTo(RAW_FIR)] get(): ( () -> Unit )
+
+        myProp#(<L> = LAZY_EXPRESSION)
+        public? final? [ResolvedTo(RAW_FIR)] fun foo(): R|kotlin/Unit| { LAZY_BLOCK }
+
+
+SUPER_TYPES:
+FILE: [ResolvedTo(IMPORTS)] statementTransformation.kts
+    context(<script>@kotlin.script.templates.standard.ScriptTemplateWithArgs)
+    SCRIPT: [ResolvedTo(SUPER_TYPES)]
+        [ResolvedTo(RAW_FIR)] lval args: R|kotlin/Array<kotlin/String>|
+
+        public? final? [ResolvedTo(RAW_FIR)] val myProp: ( () -> Unit ) = LAZY_EXPRESSION
+            public? [ResolvedTo(RAW_FIR)] get(): ( () -> Unit )
+
+        myProp#(<L> = LAZY_EXPRESSION)
+        public? final? [ResolvedTo(RAW_FIR)] fun foo(): R|kotlin/Unit| { LAZY_BLOCK }
+
+
+TYPES:
+FILE: [ResolvedTo(IMPORTS)] statementTransformation.kts
+    context(<script>@R|kotlin/script/templates/standard/ScriptTemplateWithArgs|)
+    SCRIPT: [ResolvedTo(TYPES)]
+        [ResolvedTo(RAW_FIR)] lval args: R|kotlin/Array<kotlin/String>|
+
+        public? final? [ResolvedTo(RAW_FIR)] val myProp: ( () -> Unit ) = LAZY_EXPRESSION
+            public? [ResolvedTo(RAW_FIR)] get(): ( () -> Unit )
+
+        myProp#(<L> = LAZY_EXPRESSION)
+        public? final? [ResolvedTo(RAW_FIR)] fun foo(): R|kotlin/Unit| { LAZY_BLOCK }
+
+
+STATUS:
+FILE: [ResolvedTo(IMPORTS)] statementTransformation.kts
+    context(<script>@R|kotlin/script/templates/standard/ScriptTemplateWithArgs|)
+    SCRIPT: [ResolvedTo(STATUS)]
+        [ResolvedTo(RAW_FIR)] lval args: R|kotlin/Array<kotlin/String>|
+
+        public? final? [ResolvedTo(RAW_FIR)] val myProp: ( () -> Unit ) = LAZY_EXPRESSION
+            public? [ResolvedTo(RAW_FIR)] get(): ( () -> Unit )
+
+        myProp#(<L> = LAZY_EXPRESSION)
+        public? final? [ResolvedTo(RAW_FIR)] fun foo(): R|kotlin/Unit| { LAZY_BLOCK }
+
+
+EXPECT_ACTUAL_MATCHING:
+FILE: [ResolvedTo(IMPORTS)] statementTransformation.kts
+    context(<script>@R|kotlin/script/templates/standard/ScriptTemplateWithArgs|)
+    SCRIPT: [ResolvedTo(EXPECT_ACTUAL_MATCHING)]
+        [ResolvedTo(RAW_FIR)] lval args: R|kotlin/Array<kotlin/String>|
+
+        public? final? [ResolvedTo(RAW_FIR)] val myProp: ( () -> Unit ) = LAZY_EXPRESSION
+            public? [ResolvedTo(RAW_FIR)] get(): ( () -> Unit )
+
+        myProp#(<L> = LAZY_EXPRESSION)
+        public? final? [ResolvedTo(RAW_FIR)] fun foo(): R|kotlin/Unit| { LAZY_BLOCK }
+
+
+ARGUMENTS_OF_ANNOTATIONS:
+FILE: [ResolvedTo(IMPORTS)] statementTransformation.kts
+    context(<script>@R|kotlin/script/templates/standard/ScriptTemplateWithArgs|)
+    SCRIPT: [ResolvedTo(ARGUMENTS_OF_ANNOTATIONS)]
+        [ResolvedTo(RAW_FIR)] lval args: R|kotlin/Array<kotlin/String>|
+
+        public? final? [ResolvedTo(RAW_FIR)] val myProp: ( () -> Unit ) = LAZY_EXPRESSION
+            public? [ResolvedTo(RAW_FIR)] get(): ( () -> Unit )
+
+        myProp#(<L> = LAZY_EXPRESSION)
+        public? final? [ResolvedTo(RAW_FIR)] fun foo(): R|kotlin/Unit| { LAZY_BLOCK }
+
+
+CONTRACTS:
+FILE: [ResolvedTo(IMPORTS)] statementTransformation.kts
+    context(<script>@R|kotlin/script/templates/standard/ScriptTemplateWithArgs|)
+    SCRIPT: [ResolvedTo(CONTRACTS)]
+        [ResolvedTo(RAW_FIR)] lval args: R|kotlin/Array<kotlin/String>|
+
+        public? final? [ResolvedTo(RAW_FIR)] val myProp: ( () -> Unit ) = LAZY_EXPRESSION
+            public? [ResolvedTo(RAW_FIR)] get(): ( () -> Unit )
+
+        myProp#(<L> = LAZY_EXPRESSION)
+        public? final? [ResolvedTo(RAW_FIR)] fun foo(): R|kotlin/Unit| { LAZY_BLOCK }
+
+
+IMPLICIT_TYPES_BODY_RESOLVE:
+FILE: [ResolvedTo(IMPORTS)] statementTransformation.kts
+    context(<script>@R|kotlin/script/templates/standard/ScriptTemplateWithArgs|)
+    SCRIPT: [ResolvedTo(IMPLICIT_TYPES_BODY_RESOLVE)]
+        [ResolvedTo(RAW_FIR)] lval args: R|kotlin/Array<kotlin/String>|
+
+        public? final? [ResolvedTo(RAW_FIR)] val myProp: ( () -> Unit ) = LAZY_EXPRESSION
+            public? [ResolvedTo(RAW_FIR)] get(): ( () -> Unit )
+
+        myProp#(<L> = LAZY_EXPRESSION)
+        public? final? [ResolvedTo(RAW_FIR)] fun foo(): R|kotlin/Unit| { LAZY_BLOCK }
+
+
+ANNOTATIONS_ARGUMENTS_MAPPING:
+FILE: [ResolvedTo(IMPORTS)] statementTransformation.kts
+    context(<script>@R|kotlin/script/templates/standard/ScriptTemplateWithArgs|)
+    SCRIPT: [ResolvedTo(ANNOTATIONS_ARGUMENTS_MAPPING)]
+        [ResolvedTo(RAW_FIR)] lval args: R|kotlin/Array<kotlin/String>|
+
+        public? final? [ResolvedTo(RAW_FIR)] val myProp: ( () -> Unit ) = LAZY_EXPRESSION
+            public? [ResolvedTo(RAW_FIR)] get(): ( () -> Unit )
+
+        myProp#(<L> = LAZY_EXPRESSION)
+        public? final? [ResolvedTo(RAW_FIR)] fun foo(): R|kotlin/Unit| { LAZY_BLOCK }
+
+
+BODY_RESOLVE:
+FILE: [ResolvedTo(IMPORTS)] statementTransformation.kts
+    context(<script>@R|kotlin/script/templates/standard/ScriptTemplateWithArgs|)
+    SCRIPT: [ResolvedTo(BODY_RESOLVE)]
+        [ResolvedTo(BODY_RESOLVE)] lval args: R|kotlin/Array<kotlin/String>|
+
+        public final [ResolvedTo(CONTRACTS)] val myProp: R|() -> kotlin/Unit| = [ResolvedTo(RAW_FIR)] fun <implicit>.<anonymous>(): <implicit> <inline=Unknown>  {
+            ^ Unit
+        }
+
+            public [ResolvedTo(CONTRACTS)] get(): R|() -> kotlin/Unit|
+
+        R|/myProp|.R|SubstitutionOverride<kotlin/Function0.invoke: R|kotlin/Unit|><Inapplicable(INAPPLICABLE_ARGUMENTS_MAPPING_ERROR): kotlin/Function0.invoke>#|(<L> = [ResolvedTo(BODY_RESOLVE)] myProp@fun <anonymous>(): R|kotlin/Int| <inline=Unknown>  {
+            ^ Int(123)
+        }
+        )
+        public? final? [ResolvedTo(RAW_FIR)] fun foo(): R|kotlin/Unit| { LAZY_BLOCK }
+
+
+FILE RAW TO BODY:
+FILE: [ResolvedTo(IMPORTS)] statementTransformation.kts
+    context(<script>@R|kotlin/script/templates/standard/ScriptTemplateWithArgs|)
+    SCRIPT: [ResolvedTo(BODY_RESOLVE)]
+        [ResolvedTo(BODY_RESOLVE)] lval args: R|kotlin/Array<kotlin/String>|
+
+        public final [ResolvedTo(BODY_RESOLVE)] val myProp: R|() -> kotlin/Unit| = [ResolvedTo(BODY_RESOLVE)] fun <anonymous>(): R|kotlin/Unit| <inline=Unknown>  {
+            ^ Unit
+        }
+
+            public [ResolvedTo(BODY_RESOLVE)] get(): R|() -> kotlin/Unit|
+
+        R|/myProp|.R|SubstitutionOverride<kotlin/Function0.invoke: R|kotlin/Unit|><Inapplicable(INAPPLICABLE_ARGUMENTS_MAPPING_ERROR): kotlin/Function0.invoke>#|(<L> = [ResolvedTo(BODY_RESOLVE)] myProp@fun <anonymous>(): R|kotlin/Int| <inline=Unknown>  {
+            ^ Int(123)
+        }
+        )
+        public final [ResolvedTo(BODY_RESOLVE)] fun foo(): R|kotlin/Unit| {
+        }

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/FirScriptLazyDeclarationResolveTestGenerated.java

@@ -468,6 +468,12 @@ public class FirScriptLazyDeclarationResolveTestGenerated extends AbstractFirScr
         runTest("analysis/low-level-api-fir/testdata/lazyResolve/statement2.kts");
     }
 
+    @Test
+    @TestMetadata("statementTransformation.kts")
+    public void testStatementTransformation() throws Exception {
+        runTest("analysis/low-level-api-fir/testdata/lazyResolve/statementTransformation.kts");
+    }
+
     @Test
     @TestMetadata("substitutionFakeOverrideScript.kts")
     public void testSubstitutionFakeOverrideScript() throws Exception {