K2 Scripting: wrap script statements into blocks

(reuse anonymous initializers as block wrappers) so the top-level script
elements are all declarations now. Rename the property accordingly (
together with the previous commit).
It makes script more similar to the class and thus simplify e.g.
control flow analysis and resolve code.

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/symbols/pointers/KtFirScriptParameterSymbolPointer.kt

@@ -27,7 +27,7 @@ internal class KtFirScriptParameterSymbolPointer(
             scriptPointer.restoreSymbol()?.firSymbol?.fir as? FirScript
         } ?: return null
 
-        val parameter = script.parameters.find { it.name == parameterName } as? FirProperty ?: return null
+        val parameter = script.parameters.find { it.name == parameterName } ?: return null
         return analysisSession.firSymbolBuilder.variableLikeBuilder.buildLocalVariableSymbol(parameter.symbol)
     }
 

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

@@ -24,7 +24,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.state.LLFirResolvableReso
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.FirElementFinder
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.codeFragment
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.errorWithFirSpecificEntries
-import org.jetbrains.kotlin.analysis.low.level.api.fir.util.isScriptStatement
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.isElementWhichShouldBeResolvedAsPartOfScript
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.originalDeclaration
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.originalKtFile
 import org.jetbrains.kotlin.analysis.utils.printer.parentOfType
@@ -33,7 +33,6 @@ import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.builder.buildFileAnnotationsContainer
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
-import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.realPsi
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirResolveContextCollector
@@ -364,11 +363,11 @@ object LowLevelFirApiFacadeForResolveOnAir {
      * We assume that [newScript] has the same declarations as [originalScript]
      */
     private fun restoreOriginalDeclarationsInScript(originalScript: FirScript, newScript: FirScript) {
-        val updatedStatements = ArrayList<FirStatement>(newScript.declarations.size)
+        val updatedDeclarations = ArrayList<FirDeclaration>(newScript.declarations.size)
         val originalDeclarations = originalScript.declarations.iterator()
-        for (recreatedStatement in newScript.declarations) {
+        for (recreatedDeclaration in newScript.declarations) {
-            updatedStatements += if (recreatedStatement.isScriptStatement) {
+            updatedDeclarations += if (recreatedDeclaration.isElementWhichShouldBeResolvedAsPartOfScript) {
-                recreatedStatement
+                recreatedDeclaration
             } else {
                 originalDeclarations.nextDeclaration() ?: scriptDeclarationInconsistencyError(originalScript, newScript)
             }
@@ -379,12 +378,12 @@ object LowLevelFirApiFacadeForResolveOnAir {
             scriptDeclarationInconsistencyError(originalScript, newScript)
         }
 
-        newScript.replaceDeclarations(updatedStatements)
+        newScript.replaceDeclarations(updatedDeclarations)
     }
 
     private fun scriptDeclarationInconsistencyError(originalScript: FirScript, newScript: FirScript): Nothing {
-        val originalDeclarations = originalScript.declarations.filterNot(FirStatement::isScriptStatement)
+        val originalDeclarations = originalScript.declarations.filterNot(FirDeclaration::isElementWhichShouldBeResolvedAsPartOfScript)
-        val newDeclarations = newScript.declarations.filterNot(FirStatement::isScriptStatement)
+        val newDeclarations = newScript.declarations.filterNot(FirDeclaration::isElementWhichShouldBeResolvedAsPartOfScript)
         errorWithAttachment("New script has ${if (newDeclarations.size > originalDeclarations.size) "more" else "less"} declarations") {
             withFirEntry("originalScript", originalScript)
             withFirEntry("newScript", newScript)
@@ -402,12 +401,12 @@ object LowLevelFirApiFacadeForResolveOnAir {
         }
     }
 
-    private fun Iterator<FirStatement>.nextDeclaration(): FirStatement? {
+    private fun Iterator<FirDeclaration>.nextDeclaration(): FirDeclaration? {
         while (hasNext()) {
-            val statement = next()
+            val declaration = next()
-            if (statement.isScriptStatement) continue
+            if (declaration.isElementWhichShouldBeResolvedAsPartOfScript) continue
 
-            return statement
+            return declaration
         }
 
         return null

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/FileStructureElement.kt

@@ -15,7 +15,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.diagnostics.FileStructure
 import org.jetbrains.kotlin.analysis.low.level.api.fir.diagnostics.ScriptDiagnosticRetriever
 import org.jetbrains.kotlin.analysis.low.level.api.fir.diagnostics.SingleNonLocalDeclarationDiagnosticRetriever
 import org.jetbrains.kotlin.analysis.low.level.api.fir.diagnostics.isImplicitConstructor
-import org.jetbrains.kotlin.analysis.low.level.api.fir.util.isScriptStatement
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.isElementWhichShouldBeResolvedAsPartOfScript
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.correspondingProperty
 import org.jetbrains.kotlin.fir.declarations.*
@@ -142,7 +142,7 @@ internal fun <T, R> visitScriptDependentElements(script: FirScript, visitor: Fir
     script.annotations.forEach { it.accept(visitor, data) }
     script.contextReceivers.forEach { it.accept(visitor, data) }
     script.declarations.forEach {
-        if (it.isScriptStatement) {
+        if (it.isElementWhichShouldBeResolvedAsPartOfScript) {
             it.accept(visitor, data)
         }
     }

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/FirLazyBodiesCalculator.kt

@@ -85,7 +85,7 @@ internal object FirLazyBodiesCalculator {
         return newAnnotationCall.argumentList
     }
 
-    fun createStatementsForScript(script: FirScript): List<FirStatement> {
+    fun createDeclarationsForScript(script: FirScript): List<FirDeclaration> {
         val newScript = revive<FirScript>(FirDesignation(emptyList(), script))
         return newScript.declarations
     }

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/LLFirPhaseUpdater.kt

@@ -6,13 +6,12 @@
 package org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve
 
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.forEachDependentDeclaration
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.isElementWhichShouldBeResolvedAsPartOfScript
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.isScriptDependentDeclaration
-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.declarations.impl.FirDefaultPropertyAccessor
-import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.visitors.FirVisitor
 
 internal object LLFirPhaseUpdater {
@@ -37,9 +36,9 @@ internal object LLFirPhaseUpdater {
                 is FirAnonymousInitializer -> target.body?.accept(PhaseUpdatingTransformer, newPhase)
                 is FirScript -> {
                     target.parameters.forEach { it.accept(PhaseUpdatingTransformer, newPhase) }
-                    for (statement in target.declarations) {
+                    for (declaration in target.declarations) {
-                        if (!statement.isScriptStatement) continue
+                        if (!declaration.isElementWhichShouldBeResolvedAsPartOfScript) continue
-                        statement.accept(PhaseUpdatingTransformer, newPhase)
+                        declaration.accept(PhaseUpdatingTransformer, newPhase)
                     }
                 }
 
@@ -84,7 +83,7 @@ private object PhaseUpdatingTransformer : FirVisitor<Unit, FirResolvePhase>() {
         if (element is FirElementWithResolveState) {
             if (element.resolvePhase >= data &&
                 element !is FirDefaultPropertyAccessor &&
-                !(element is FirStatement && element.isScriptDependentDeclaration)
+                !(element is FirDeclaration && element.isScriptDependentDeclaration)
             ) return
 
             @OptIn(ResolveStateAccess::class)

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

@@ -10,11 +10,10 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.LLFirResolveT
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.LLFirReturnTypeCalculatorWithJump
 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.analysis.low.level.api.fir.util.isElementWhichShouldBeResolvedAsPartOfScript
 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.body.resolve.FirAbstractBodyResolveTransformerDispatcher
@@ -88,9 +87,9 @@ internal abstract class LLFirAbstractBodyTargetResolver(
             script.transformDeclarations(
                 transformer = object : FirTransformer<Any?>() {
                     override fun <E : FirElement> transformElement(element: E, data: Any?): E {
-                        if (element !is FirStatement || !element.isScriptStatement) return element
+                        if (element !is FirDeclaration || !element.isElementWhichShouldBeResolvedAsPartOfScript) return element
 
-                        transformer.firResolveContextCollector?.addStatementContext(element, transformer.context)
+                        transformer.firResolveContextCollector?.addDeclarationContext(element, transformer.context)
                         return element.transformSingle(transformer, ResolutionMode.ContextIndependent)
                     }
                 },

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

@@ -281,7 +281,7 @@ private class LLFirBodyTargetResolver(
 
     override fun rawResolve(target: FirElementWithResolveState) {
         when (target) {
-            is FirScript -> target.takeUnless(FirScript::isCertainlyResolved)?.let(::resolveScript)
+            is FirScript -> target.let(::resolveScript)
             else -> super.rawResolve(target)
         }
 
@@ -291,20 +291,22 @@ private class LLFirBodyTargetResolver(
 
 internal object BodyStateKeepers {
     val SCRIPT: StateKeeper<FirScript, FirDesignationWithFile> = stateKeeper { script, designation ->
-        val oldStatements = script.declarations
+        val oldDeclarations = script.declarations
-        if (oldStatements.none { it.isScriptStatement } || script.isCertainlyResolved) return@stateKeeper
+        if (oldDeclarations.none { it.isElementWhichShouldBeResolvedAsPartOfScript }) return@stateKeeper
 
         add(RESULT_PROPERTY, designation)
         add(FirScript::declarations, FirScript::replaceDeclarations) {
-            val recreatedStatements = FirLazyBodiesCalculator.createStatementsForScript(script)
+            val recreatedDeclarations = FirLazyBodiesCalculator.createDeclarationsForScript(script)
-            requireSameSize(oldStatements, recreatedStatements)
+            requireSameSize(oldDeclarations, recreatedDeclarations)
 
-            ArrayList<FirStatement>(oldStatements.size).apply {
+            ArrayList<FirDeclaration>(oldDeclarations.size).apply {
-                oldStatements.zip(recreatedStatements).mapTo(this) { (old, new) ->
+                oldDeclarations.zip(recreatedDeclarations).mapTo(this) { (old, new) ->
                     when {
-                        !old.isScriptStatement -> old
+                        !old.isElementWhichShouldBeResolvedAsPartOfScript -> old
                         old is FirProperty && old.origin == FirDeclarationOrigin.ScriptCustomization.ResultProperty -> {
-                            old.replaceInitializer((new as FirProperty).initializer)
+                            if (old.bodyResolveState != FirPropertyBodyResolveState.ALL_BODIES_RESOLVED) {
+                                old.replaceInitializer((new as FirProperty).initializer)
+                            }
                             old
                         }
 
@@ -451,14 +453,6 @@ private fun FirScript.findResultProperty(): FirProperty? = declarations.findIsIn
     it.origin == FirDeclarationOrigin.ScriptCustomization.ResultProperty
 }
 
-private val FirScript.isCertainlyResolved: Boolean
-    get() {
-        val dependentProperty = findResultProperty() ?: return false
-
-        // This meant that we already resolve the entire script on implicit body phase
-        return dependentProperty.bodyResolveState == FirPropertyBodyResolveState.ALL_BODIES_RESOLVED
-    }
-
 private val FirFunction.isCertainlyResolved: Boolean
     get() {
         if (this is FirPropertyAccessor) {
@@ -528,10 +522,10 @@ private fun delegatedConstructorCallGuard(fir: FirDelegatedConstructorCall): Fir
     }
 }
 
-private fun requireSameSize(old: List<FirStatement>, new: List<FirStatement>) {
+private fun requireSameSize(old: List<FirDeclaration>, new: List<FirDeclaration>) {
     requireWithAttachment(
         condition = old.size == new.size,
-        message = { "The number of statements are different" }
+        message = { "The number of declarations are different" }
     ) {
         withEntryGroup("originalStatements") {
             old.forEachIndexed { index, statement -> withFirEntry("statement$index", statement) }

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

@@ -97,7 +97,7 @@ internal object FirElementFinder {
                     }
 
                     subDeclaration is FirScript -> {
-                        val scriptDeclarations = subDeclaration.declarations.asSequence().filterIsInstance<FirDeclaration>()
+                        val scriptDeclarations = subDeclaration.declarations.asSequence().filterNot { it is FirAnonymousInitializer }
                         if (find(scriptDeclarations.asIterable(), classIdPathIndex)) {
                             return true
                         }

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

@@ -16,7 +16,6 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.getNonLoc
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.LLFirFileBuilder
 import org.jetbrains.kotlin.analysis.low.level.api.fir.providers.LLFirProvider
 import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.realPsi
 import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
@@ -60,7 +59,7 @@ internal fun KtDeclaration.findSourceNonLocalFirDeclaration(firFile: FirFile, pr
                                 return@findSourceNonLocalFirDeclarationByProvider firScript?.takeIf { it.psi == declaration }
                             }
 
-                            firScript?.declarations?.filterIsInstance<FirDeclaration>()
+                            firScript?.declarations?.filterNot { it is FirAnonymousInitializer }
                         } else {
                             firFile.declarations
                         }
@@ -212,8 +211,8 @@ val FirDeclaration.isGeneratedDeclaration
 
 internal inline fun FirScript.forEachDeclaration(action: (FirDeclaration) -> Unit) {
     for (statement in declarations) {
-        if (statement.isScriptStatement) continue
+        if (statement.isElementWhichShouldBeResolvedAsPartOfScript) continue
-        action(statement as FirDeclaration)
+        action(statement)
     }
 }
 
@@ -236,15 +235,15 @@ internal inline fun FirDeclaration.forEachDeclaration(action: (FirDeclaration) -
     }
 }
 
-internal val FirStatement.isScriptStatement: Boolean get() = this !is FirDeclaration || isScriptDependentDeclaration
+internal val FirDeclaration.isElementWhichShouldBeResolvedAsPartOfScript: Boolean get() = this is FirAnonymousInitializer || isScriptDependentDeclaration
 
-internal val FirStatement.isScriptDependentDeclaration: Boolean
+internal val FirDeclaration.isScriptDependentDeclaration: Boolean
-    get() = this is FirDeclaration && origin == FirDeclarationOrigin.ScriptCustomization.ResultProperty
+    get() = origin == FirDeclarationOrigin.ScriptCustomization.ResultProperty
 
 internal inline fun FirScript.forEachDependentDeclaration(action: (FirDeclaration) -> Unit) {
-    for (statement in declarations) {
+    for (declaration in declarations) {
-        if (statement !is FirDeclaration || !statement.isScriptDependentDeclaration) continue
+        if (declaration is FirAnonymousInitializer || !declaration.isScriptDependentDeclaration) continue
-        action(statement)
+        action(declaration)
     }
 }
 

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

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.fir.contracts.FirLegacyRawContractDescription
 import org.jetbrains.kotlin.fir.contracts.FirResolvedContractDescription
 import org.jetbrains.kotlin.fir.contracts.impl.FirEmptyContractDescription
 import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.declarations.utils.isActual
 import org.jetbrains.kotlin.fir.declarations.utils.isExpect
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
@@ -108,11 +109,16 @@ internal fun checkBodyIsResolved(function: FirFunction) {
 }
 
 internal fun checkStatementsAreResolved(script: FirScript) {
-    for (statement in script.declarations) {
+    fun checkExpression(expression: FirExpression) {
-        if (statement.isScriptStatement && statement is FirExpression) {
+        checkExpressionTypeIsResolved(expression.coneTypeOrNull, "script statement", script) {
-            checkExpressionTypeIsResolved(statement.coneTypeOrNull, "script statement", script) {
+            withFirEntry("expression", expression)
-                withFirEntry("expression", statement)
+        }
-            }
+    }
+    for (declaration in script.declarations) {
+        if (declaration.isScriptDependentDeclaration) {
+            (declaration as? FirProperty)?.initializer?.let(::checkExpression)
+        } else if (declaration is FirAnonymousInitializer) {
+            declaration.body?.statements?.filterIsInstance<FirExpression>()?.forEach(::checkExpression)
         }
     }
 }

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/AbstractFirLazyDeclarationResolveTestCase.kt

@@ -91,7 +91,7 @@ abstract class AbstractFirLazyDeclarationResolveTestCase : AbstractLowLevelApiLa
         val (classSymbol, declarations) = when (symbol) {
             is FirClassSymbol -> symbol to symbol.declarationSymbols
             is FirScriptSymbol -> {
-                symbol to symbol.fir.let { it.parameters + it.declarations }.mapNotNull { (it as? FirDeclaration)?.symbol }
+                symbol to symbol.fir.let { it.parameters + it.declarations }.map { it.symbol }
             }
 
             else -> error("Unknown container: ${symbol::class.simpleName}")

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrConverter.kt

@@ -518,19 +518,19 @@ class Fir2IrConverter(
                 addDeclarationToParentIfNeeded(irScript)
                 declarationStorage.withScope(irScript.symbol) {
                     irScript.parent = parent
-                    for (scriptStatement in declaration.declarations) {
+                    for (scriptDeclaration in declaration.declarations) {
-                        when (scriptStatement) {
+                        when (scriptDeclaration) {
                             is FirRegularClass -> {
-                                registerClassAndNestedClasses(scriptStatement, irScript)
+                                registerClassAndNestedClasses(scriptDeclaration, irScript)
-                                processClassAndNestedClassHeaders(scriptStatement)
+                                processClassAndNestedClassHeaders(scriptDeclaration)
                             }
-                            is FirTypeAlias -> classifierStorage.createAndCacheIrTypeAlias(scriptStatement, irScript)
+                            is FirTypeAlias -> classifierStorage.createAndCacheIrTypeAlias(scriptDeclaration, irScript)
                             else -> {}
                         }
                     }
-                    for (scriptStatement in declaration.declarations) {
+                    for (scriptDeclaration in declaration.declarations) {
-                        if (scriptStatement is FirDeclaration) {
+                        if (scriptDeclaration !is FirAnonymousInitializer) {
-                            processMemberDeclaration(scriptStatement, containingClass = null, irScript, delegateFieldToPropertyMap = null)
+                            processMemberDeclaration(scriptDeclaration, containingClass = null, irScript, delegateFieldToPropertyMap = null)
                         }
                     }
                 }

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrVisitor.kt

@@ -255,8 +255,8 @@ class Fir2IrVisitor(
             conversionScope.withParent(irScript) {
                 val destructComposites = mutableMapOf<FirVariableSymbol<*>, IrComposite>()
                 for (statement in script.declarations) {
-                    val irStatement = if (statement is FirDeclaration) {
+                    if (statement !is FirAnonymousInitializer) {
-                        when {
+                        val irStatement = when {
                             statement is FirProperty && statement.name == SpecialNames.UNDERSCORE_FOR_UNUSED_VAR -> {
                                 continue
                             }
@@ -311,10 +311,12 @@ class Fir2IrVisitor(
                                 statement.accept(this@Fir2IrVisitor, null) as? IrDeclaration
                             }
                         }
+                        irScript.statements.add(irStatement!!)
                     } else {
-                        statement.toIrStatement()
+                        statement.body?.statements?.mapNotNull { it.toIrStatement() }?.let {
+                            irScript.statements.addAll(it)
+                        }
                     }
-                    irScript.statements.add(irStatement!!)
                 }
             }
             for (configurator in session.extensionService.fir2IrScriptConfigurators) {

compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/PsiRawFirBuilder.kt

@@ -1276,50 +1276,81 @@ open class PsiRawFirBuilder(
             fileName: String,
             sourceFile: KtSourceFile?,
             setup: FirScriptBuilder.() -> Unit = {},
-        ): FirScript = buildScript {
+        ): FirScript {
-            source = script.toFirSourceElement()
+            val scriptName = Name.special("<script-$fileName>")
-            moduleData = baseModuleData
+            val scriptSymbol = FirScriptSymbol(context.packageFqName.child(scriptName))
-            origin = FirDeclarationOrigin.Source
-            name = Name.special("<script-$fileName>")
-            symbol = FirScriptSymbol(context.packageFqName.child(name))
-            withContainerSymbol(symbol) {
-                for (declaration in script.declarations) {
-                    when (declaration) {
-                        is KtScriptInitializer -> {
-                            declaration.body?.let { declarations.add(it.toFirStatement()) }
-                        }
-                        is KtDestructuringDeclaration -> {
-                            val destructuringContainerVar = generateTemporaryVariable(
-                                baseModuleData,
-                                declaration.toFirSourceElement(),
-                                "destruct",
-                                declaration.initializer.toFirExpression { ConeSyntaxDiagnostic("Initializer required for destructuring declaration") },
-                                origin = FirDeclarationOrigin.Synthetic.ScriptTopLevelDestructuringDeclarationContainer,
-                                extractAnnotationsTo = { extractAnnotationsTo(it) }
-                            ).apply {
-                                isDestructuringDeclarationContainerVariable = true
-                            }
-                            declarations.add(destructuringContainerVar)
 
-                            declarations.addDestructuringVariables(
+            return buildScript {
-                                moduleData,
+                source = script.toFirSourceElement()
-                                declaration,
+                moduleData = baseModuleData
-                                destructuringContainerVar,
+                origin = FirDeclarationOrigin.Source
-                                tmpVariable = false,
+                name = scriptName
-                                localEntries = false,
+                symbol = scriptSymbol
-                            ) {
+
-                                (it as FirProperty).destructuringDeclarationContainerVariable = destructuringContainerVar.symbol
+                val scriptDeclarationsIter = script.declarations.listIterator()
+                withContainerSymbol(symbol) {
+                    while (scriptDeclarationsIter.hasNext()) {
+                        val declaration = scriptDeclarationsIter.next()
+                        val isLast = !scriptDeclarationsIter.hasNext()
+                        val declarationSource = declaration.toFirSourceElement()
+                        when (declaration) {
+                            is KtScriptInitializer -> {
+                                val firBlock =
+                                    if (isLast) {
+                                        // the last one need to be analyzed in script configurator to decide on result property
+                                        // therefore no lazy conversion in this case
+                                        withForcedLocalContext { declaration.body.toFirBlock() }
+                                    } else {
+                                        buildOrLazyBlock { withForcedLocalContext { declaration.body.toFirBlock() } }
+                                    }
+                                declarations.add(
+                                    buildAnonymousInitializer {
+                                        moduleData = baseModuleData
+                                        origin = FirDeclarationOrigin.Source
+                                        source = declarationSource
+                                        body = firBlock
+                                        declaration.extractAnnotationsTo(this)
+                                    }
+                                )
+                            }
+                            is KtDestructuringDeclaration -> {
+                                val destructuringContainerVar = generateTemporaryVariable(
+                                    baseModuleData,
+                                    declarationSource,
+                                    "destruct",
+                                    declaration.initializer.toFirExpression { ConeSyntaxDiagnostic("Initializer required for destructuring declaration") },
+                                    origin = FirDeclarationOrigin.Synthetic.ScriptTopLevelDestructuringDeclarationContainer,
+                                    extractAnnotationsTo = { extractAnnotationsTo(it) }
+                                ).apply {
+                                    isDestructuringDeclarationContainerVariable = true
+                                }
+                                declarations.add(destructuringContainerVar)
+
+                                declarations.addDestructuringVariables(
+                                    moduleData,
+                                    declaration,
+                                    destructuringContainerVar,
+                                    tmpVariable = false,
+                                    localEntries = false,
+                                ) {
+                                    (it as FirProperty).destructuringDeclarationContainerVariable = destructuringContainerVar.symbol
+                                }
+                            }
+                            else -> {
+                                val firStatement = declaration.toFirStatement()
+                                if (firStatement is FirDeclaration) {
+                                    declarations.add(firStatement)
+                                } else {
+                                    error("unexpected declaration type in script")
+                                }
                             }
-                        }
-                        else -> {
-                            declarations.add(declaration.toFirStatement())
                         }
                     }
-                }
+                    setup()
-                setup()
+                    if (sourceFile != null) {
-                if (sourceFile != null) {
+                        for (configurator in baseSession.extensionService.scriptConfigurators) {
-                    for (configurator in baseSession.extensionService.scriptConfigurators) {
+                            with(configurator) { configure(sourceFile) }
-                        with(configurator) { configure(sourceFile) }
+                        }
                     }
                 }
             }

compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/script.txt

@@ -37,28 +37,43 @@ FILE: script.kts
             ^build2 Builder#().apply#(action#)
         }
 
-        build#(<L> = build@fun <implicit>.<anonymous>(): <implicit> <inline=Unknown>  {
+        init {
-            version# = String(123)
+            build#(<L> = build@fun <implicit>.<anonymous>(): <implicit> <inline=Unknown>  {
-            local final? class A : R|kotlin/Any| {
+                version# = String(123)
-                public? [ContainingClassKey=A] constructor(): R|<local>/A| {
+                local final? class A : R|kotlin/Any| {
-                    super<R|kotlin/Any|>()
+                    public? [ContainingClassKey=A] constructor(): R|<local>/A| {
-                }
+                        super<R|kotlin/Any|>()
-
+                    }
-                public? final? fun doo(): R|kotlin/Unit| {
+
+                    public? final? fun doo(): R|kotlin/Unit| {
+                    }
+
                 }
 
+                execute#()
             }
-
+            )
-            execute#()
         }
-        )
+
-        +=(version#, IntegerLiteral(123))
+        init {
+            +=(version#, IntegerLiteral(123))
+        }
+
         public? final? val builder: <implicit> = build#(<L> = build@fun <implicit>.<anonymous>(): <implicit> <inline=Unknown>  {
             version# = String(321)
         }
         )
             public? get(): <implicit>
 
-        println#()
+        init {
-        builder#.version# = String()
+            println#()
-        builder#.execute#()
+        }
+
+        init {
+            builder#.version# = String()
+        }
+
+        init {
+            builder#.execute#()
+        }
+

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/declarations/FirScript.kt

@@ -12,7 +12,6 @@ import org.jetbrains.kotlin.KtSourceElement
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirModuleData
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
-import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.references.FirControlFlowGraphReference
 import org.jetbrains.kotlin.fir.symbols.impl.FirScriptSymbol
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
@@ -30,9 +29,9 @@ abstract class FirScript : FirDeclaration(), FirControlFlowGraphOwner {
     abstract override val attributes: FirDeclarationAttributes
     abstract override val controlFlowGraphReference: FirControlFlowGraphReference?
     abstract val name: Name
-    abstract val declarations: List<FirStatement>
+    abstract val declarations: List<FirDeclaration>
     abstract override val symbol: FirScriptSymbol
-    abstract val parameters: List<FirVariable>
+    abstract val parameters: List<FirProperty>
     abstract val contextReceivers: List<FirContextReceiver>
     abstract val resultPropertyName: Name?
 
@@ -47,9 +46,13 @@ abstract class FirScript : FirDeclaration(), FirControlFlowGraphOwner {
 
     abstract override fun replaceControlFlowGraphReference(newControlFlowGraphReference: FirControlFlowGraphReference?)
 
-    abstract fun replaceDeclarations(newDeclarations: List<FirStatement>)
+    abstract fun replaceDeclarations(newDeclarations: List<FirDeclaration>)
 
     abstract override fun <D> transformAnnotations(transformer: FirTransformer<D>, data: D): FirScript
 
     abstract fun <D> transformDeclarations(transformer: FirTransformer<D>, data: D): FirScript
+
+    abstract fun <D> transformParameters(transformer: FirTransformer<D>, data: D): FirScript
+
+    abstract fun <D> transformContextReceivers(transformer: FirTransformer<D>, data: D): FirScript
 }

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/declarations/builder/FirScriptBuilder.kt

@@ -19,7 +19,6 @@ import org.jetbrains.kotlin.fir.builder.toMutableOrEmpty
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.impl.FirScriptImpl
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
-import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.symbols.impl.FirScriptSymbol
 import org.jetbrains.kotlin.name.Name
 
@@ -32,9 +31,9 @@ class FirScriptBuilder : FirAnnotationContainerBuilder {
     lateinit var origin: FirDeclarationOrigin
     var attributes: FirDeclarationAttributes = FirDeclarationAttributes()
     lateinit var name: Name
-    val declarations: MutableList<FirStatement> = mutableListOf()
+    val declarations: MutableList<FirDeclaration> = mutableListOf()
     lateinit var symbol: FirScriptSymbol
-    val parameters: MutableList<FirVariable> = mutableListOf()
+    val parameters: MutableList<FirProperty> = mutableListOf()
     val contextReceivers: MutableList<FirContextReceiver> = mutableListOf()
     var resultPropertyName: Name? = null
 
@@ -47,7 +46,7 @@ class FirScriptBuilder : FirAnnotationContainerBuilder {
             origin,
             attributes,
             name,
-            declarations.toMutableOrEmpty(),
+            declarations,
             symbol,
             parameters,
             contextReceivers.toMutableOrEmpty(),

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/declarations/impl/FirScriptImpl.kt

@@ -16,7 +16,6 @@ import org.jetbrains.kotlin.fir.MutableOrEmptyList
 import org.jetbrains.kotlin.fir.builder.toMutableOrEmpty
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
-import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.references.FirControlFlowGraphReference
 import org.jetbrains.kotlin.fir.symbols.impl.FirScriptSymbol
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
@@ -33,9 +32,9 @@ internal class FirScriptImpl(
     override val origin: FirDeclarationOrigin,
     override val attributes: FirDeclarationAttributes,
     override val name: Name,
-    override var declarations: MutableOrEmptyList<FirStatement>,
+    override val declarations: MutableList<FirDeclaration>,
     override val symbol: FirScriptSymbol,
-    override val parameters: MutableList<FirVariable>,
+    override val parameters: MutableList<FirProperty>,
     override var contextReceivers: MutableOrEmptyList<FirContextReceiver>,
     override val resultPropertyName: Name?,
 ) : FirScript() {
@@ -58,8 +57,8 @@ internal class FirScriptImpl(
         transformAnnotations(transformer, data)
         controlFlowGraphReference = controlFlowGraphReference?.transform(transformer, data)
         transformDeclarations(transformer, data)
-        parameters.transformInplace(transformer, data)
+        transformParameters(transformer, data)
-        contextReceivers.transformInplace(transformer, data)
+        transformContextReceivers(transformer, data)
         return this
     }
 
@@ -73,6 +72,16 @@ internal class FirScriptImpl(
         return this
     }
 
+    override fun <D> transformParameters(transformer: FirTransformer<D>, data: D): FirScriptImpl {
+        parameters.transformInplace(transformer, data)
+        return this
+    }
+
+    override fun <D> transformContextReceivers(transformer: FirTransformer<D>, data: D): FirScriptImpl {
+        contextReceivers.transformInplace(transformer, data)
+        return this
+    }
+
     override fun replaceAnnotations(newAnnotations: List<FirAnnotation>) {
         annotations = newAnnotations.toMutableOrEmpty()
     }
@@ -81,7 +90,8 @@ internal class FirScriptImpl(
         controlFlowGraphReference = newControlFlowGraphReference
     }
 
-    override fun replaceDeclarations(newDeclarations: List<FirStatement>) {
+    override fun replaceDeclarations(newDeclarations: List<FirDeclaration>) {
-        declarations = newDeclarations.toMutableOrEmpty()
+        declarations.clear()
+        declarations.addAll(newDeclarations)
     }
 }

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/NodeConfigurator.kt

@@ -494,10 +494,10 @@ object NodeConfigurator : AbstractFieldConfigurator<FirTreeBuilder>(FirTreeBuild
 
         script.configure {
             +name
-            +FieldList("declarations", statement, withReplace = true, useMutableOrEmpty = true).withTransform()
+            +declarations.withTransform().withReplace()
             +symbol("FirScriptSymbol")
-            +fieldList("parameters", variable, withReplace = false)
+            +fieldList("parameters", property).withTransform()
-            +fieldList(contextReceiver, useMutableOrEmpty = true)
+            +fieldList(contextReceiver, useMutableOrEmpty = true).withTransform()
             +field("resultPropertyName", nameType, nullable = true)
         }
 

plugins/scripting/scripting-compiler/src/org/jetbrains/kotlin/scripting/compiler/plugin/services/FirScriptConfigurationExtensionImpl.kt

@@ -12,14 +12,16 @@ import org.jetbrains.kotlin.descriptors.Visibilities
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.builder.FirScriptConfiguratorExtension
 import org.jetbrains.kotlin.fir.builder.FirScriptConfiguratorExtension.Factory
+import org.jetbrains.kotlin.fir.declarations.FirAnonymousInitializer
 import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
 import org.jetbrains.kotlin.fir.declarations.builder.*
 import org.jetbrains.kotlin.fir.declarations.impl.FirDeclarationStatusImpl
 import org.jetbrains.kotlin.fir.declarations.impl.FirDefaultPropertyGetter
 import org.jetbrains.kotlin.fir.declarations.primaryConstructorIfAny
 import org.jetbrains.kotlin.fir.declarations.utils.SCRIPT_SPECIAL_NAME_STRING
-import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.expressions.*
-import org.jetbrains.kotlin.fir.expressions.UnresolvedExpressionTypeAccess
+import org.jetbrains.kotlin.fir.expressions.builder.buildLazyExpression
+import org.jetbrains.kotlin.fir.expressions.impl.FirSingleExpressionBlock
 import org.jetbrains.kotlin.fir.moduleData
 import org.jetbrains.kotlin.fir.resolve.providers.dependenciesSymbolProvider
 import org.jetbrains.kotlin.fir.symbols.SymbolInternals
@@ -137,11 +139,21 @@ class FirScriptConfiguratorExtensionImpl(
         }
 
         configuration[ScriptCompilationConfiguration.resultField]?.takeIf { it.isNotBlank() }?.let { resultFieldName ->
-            val lastExpression = declarations.lastOrNull()
+            val lastScriptBlock = declarations.lastOrNull() as? FirAnonymousInitializer
-            if (lastExpression != null && lastExpression is FirExpression) {
+            val lastExpression =
-                declarations.removeAt(declarations.size - 1)
+                when (val lastScriptBlockBody = lastScriptBlock?.body) {
+                    is FirLazyBlock -> null
+                    is FirSingleExpressionBlock -> lastScriptBlockBody.statement as? FirExpression
+                    else -> lastScriptBlockBody?.statements?.single()?.takeIf { it is FirExpression } as? FirExpression
+                }?.takeUnless { it is FirErrorExpression }
+
+            if (lastExpression != null) {
+                declarations.removeLast()
+                @OptIn(UnresolvedExpressionTypeAccess::class)
+                val lastExpressionTypeRef =
+                    lastExpression.takeUnless { it is FirLazyExpression }?.coneTypeOrNull?.toFirResolvedTypeRef()
+                        ?: FirImplicitTypeRefImplWithoutSource
                 declarations.add(
-                    @OptIn(UnresolvedExpressionTypeAccess::class)
                     buildProperty {
                         this.name = Name.identifier(resultFieldName)
                         this.symbol = FirPropertySymbol(this.name)
@@ -149,12 +161,12 @@ class FirScriptConfiguratorExtensionImpl(
                         moduleData = session.moduleData
                         origin = FirDeclarationOrigin.ScriptCustomization.ResultProperty
                         initializer = lastExpression
-                        returnTypeRef = lastExpression.coneTypeOrNull?.toFirResolvedTypeRef() ?: FirImplicitTypeRefImplWithoutSource
+                        returnTypeRef = lastExpressionTypeRef
                         getter = FirDefaultPropertyGetter(
                             lastExpression.source,
                             session.moduleData,
                             FirDeclarationOrigin.ScriptCustomization.ResultProperty,
-                            lastExpression.coneTypeOrNull?.toFirResolvedTypeRef() ?: FirImplicitTypeRefImplWithoutSource,
+                            lastExpressionTypeRef,
                             Visibilities.Public,
                             this.symbol,
                         )