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) {
-            updatedStatements += if (recreatedStatement.isScriptStatement) {
-                recreatedStatement
+        for (recreatedDeclaration in newScript.declarations) {
+            updatedDeclarations += if (recreatedDeclaration.isElementWhichShouldBeResolvedAsPartOfScript) {
+                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 newDeclarations = newScript.declarations.filterNot(FirStatement::isScriptStatement)
+        val originalDeclarations = originalScript.declarations.filterNot(FirDeclaration::isElementWhichShouldBeResolvedAsPartOfScript)
+        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()
-            if (statement.isScriptStatement) continue
+            val declaration = next()
+            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) {
-                        if (!statement.isScriptStatement) continue
-                        statement.accept(PhaseUpdatingTransformer, newPhase)
+                    for (declaration in target.declarations) {
+                        if (!declaration.isElementWhichShouldBeResolvedAsPartOfScript) continue
+                        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
-        if (oldStatements.none { it.isScriptStatement } || script.isCertainlyResolved) return@stateKeeper
+        val oldDeclarations = script.declarations
+        if (oldDeclarations.none { it.isElementWhichShouldBeResolvedAsPartOfScript }) return@stateKeeper
 
         add(RESULT_PROPERTY, designation)
         add(FirScript::declarations, FirScript::replaceDeclarations) {
-            val recreatedStatements = FirLazyBodiesCalculator.createStatementsForScript(script)
-            requireSameSize(oldStatements, recreatedStatements)
+            val recreatedDeclarations = FirLazyBodiesCalculator.createDeclarationsForScript(script)
+            requireSameSize(oldDeclarations, recreatedDeclarations)
 
-            ArrayList<FirStatement>(oldStatements.size).apply {
-                oldStatements.zip(recreatedStatements).mapTo(this) { (old, new) ->
+            ArrayList<FirDeclaration>(oldDeclarations.size).apply {
+                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
-        action(statement as FirDeclaration)
+        if (statement.isElementWhichShouldBeResolvedAsPartOfScript) continue
+        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
-    get() = this is FirDeclaration && origin == FirDeclarationOrigin.ScriptCustomization.ResultProperty
+internal val FirDeclaration.isScriptDependentDeclaration: Boolean
+    get() = origin == FirDeclarationOrigin.ScriptCustomization.ResultProperty
 
 internal inline fun FirScript.forEachDependentDeclaration(action: (FirDeclaration) -> Unit) {
-    for (statement in declarations) {
-        if (statement !is FirDeclaration || !statement.isScriptDependentDeclaration) continue
-        action(statement)
+    for (declaration in declarations) {
+        if (declaration is FirAnonymousInitializer || !declaration.isScriptDependentDeclaration) continue
+        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) {
-        if (statement.isScriptStatement && statement is FirExpression) {
-            checkExpressionTypeIsResolved(statement.coneTypeOrNull, "script statement", script) {
-                withFirEntry("expression", statement)
-            }
+    fun checkExpression(expression: FirExpression) {
+        checkExpressionTypeIsResolved(expression.coneTypeOrNull, "script statement", script) {
+            withFirEntry("expression", expression)
+        }
+    }
+    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}")