[LL FIR] do not rebuild entire script for lazy calculation of result$$ property

After the previous commit, it is possible to rebuild initializer for
`result$$` property without a full script creation.
Also, this commit fixes anchor symbol for annotations
inside script statements

^KT-65344

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

@@ -11,6 +11,7 @@ import kotlinx.collections.immutable.persistentListOf
 import kotlinx.collections.immutable.toPersistentList
 import org.jetbrains.kotlin.KtFakeSourceElementKind
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirDesignation
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.withFirDesignationEntry
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.forEachDependentDeclaration
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.builder.PsiRawFirBuilder
@@ -85,11 +86,31 @@ internal object FirLazyBodiesCalculator {
         return newAnnotationCall.argumentList
     }
 
-    fun createDeclarationsForScript(designation: FirDesignation): List<FirDeclaration> {
-        designation.target as FirScript
+    fun createResultPropertyInitializer(scriptDesignation: FirDesignation, property: FirProperty): FirExpression {
+        val script = scriptDesignation.target
+        requireWithAttachment(script is FirScript, { "${FirScript::class.simpleName} is not found" }) {
+            withFirDesignationEntry("designation", scriptDesignation)
+            withFirEntry("property", property)
+        }
 
-        val newScript = revive<FirScript>(designation)
-        return newScript.declarations
+        val propertyDesignation = FirDesignation(path = scriptDesignation.path + script, target = property)
+        val newInitializer = revive<FirAnonymousInitializer>(propertyDesignation)
+        val body = newInitializer.body
+        requireWithAttachment(body != null, { "${FirAnonymousInitializer::class.simpleName} without body" }) {
+            withFirDesignationEntry("designation", propertyDesignation)
+            withFirEntry("initializer", newInitializer)
+        }
+
+        val singleStatement = body.statements.singleOrNull()
+        requireWithAttachment(singleStatement is FirExpression, { "Unexpected body content" }) {
+            withFirDesignationEntry("designation", propertyDesignation)
+            withFirEntry("initializer", newInitializer)
+            singleStatement?.let {
+                withFirEntry("statement", it)
+            }
+        }
+
+        return singleStatement
     }
 
     fun needCalculatingAnnotationCall(firAnnotationCall: FirAnnotationCall): Boolean =

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

@@ -152,6 +152,12 @@ internal class RawFirNonLocalDeclarationBuilder private constructor(
             return buildErrorTopLevelDestructuringDeclaration(replacementDeclaration.toFirSourceElement())
         }
 
+        fun convertAnonymousInitializer(element: KtAnonymousInitializer, containingDeclaration: FirDeclaration?): FirAnonymousInitializer {
+            val replacementDeclaration = replacementApplier?.tryReplace(element) ?: element
+            requireIsInstance<KtAnonymousInitializer>(replacementDeclaration)
+            return buildAnonymousInitializer(replacementDeclaration, containingDeclaration?.symbol)
+        }
+
         override fun convertElement(element: KtElement, original: FirElement?): FirElement? =
             super.convertElement(replacementApplier?.tryReplace(element) ?: element, original)
 
@@ -347,6 +353,7 @@ internal class RawFirNonLocalDeclarationBuilder private constructor(
                     val firFile = visitor.convertElement(declarationToBuild, originalDeclaration) as FirFile
                     firFile.codeFragment
                 }
+                is KtAnonymousInitializer -> visitor.convertAnonymousInitializer(declarationToBuild, containingDeclaration)
                 else -> visitor.convertElement(declarationToBuild, originalDeclaration)
             } as FirDeclaration
         }

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

@@ -51,7 +51,6 @@ import org.jetbrains.kotlin.psi.KtCodeFragment
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.utils.exceptions.checkWithAttachment
 import org.jetbrains.kotlin.utils.exceptions.errorWithAttachment
-import org.jetbrains.kotlin.utils.exceptions.requireWithAttachment
 import org.jetbrains.kotlin.utils.exceptions.withPsiEntry
 import org.jetbrains.kotlin.utils.findIsInstanceAnd
 
@@ -352,37 +351,22 @@ internal object BodyStateKeepers {
         val oldDeclarations = script.declarations
         if (oldDeclarations.none { it.isElementWhichShouldBeResolvedAsPartOfScript }) return@stateKeeper
 
-        val lastProperty = oldDeclarations.lastOrNull()
-        if (lastProperty is FirProperty &&
-            lastProperty.origin == FirDeclarationOrigin.ScriptCustomization.ResultProperty &&
-            lastProperty.bodyResolveState != FirPropertyBodyResolveState.ALL_BODIES_RESOLVED
-        ) {
-            add(RESULT_PROPERTY, designation)
-            val recreatedDeclarations = FirLazyBodiesCalculator.createDeclarationsForScript(designation)
-            requireSameSize(oldDeclarations, recreatedDeclarations)
-            lastProperty.replaceInitializer((recreatedDeclarations.last() as FirProperty).initializer)
+        val resultProperty = script.findResultProperty()
+        if (resultProperty != null && resultProperty.bodyResolveState != FirPropertyBodyResolveState.ALL_BODIES_RESOLVED) {
+            entity(resultProperty, RESULT_PROPERTY, designation)
         }
 
         entityList(oldDeclarations.mapNotNull { it as? FirAnonymousInitializer }, ANONYMOUS_INITIALIZER, designation)
         add(FirScript::controlFlowGraphReference, FirScript::replaceControlFlowGraphReference)
     }
 
-    private val RESULT_PROPERTY: StateKeeper<FirScript, FirDesignation> = stateKeeper { script, _ ->
-        val resultedProperty = script.findResultProperty() ?: return@stateKeeper
-        add(
-            provider = { resultedProperty.bodyResolveState },
-            mutator = { _, value -> resultedProperty.replaceBodyResolveState(value) },
-        )
-
-        add(
-            provider = { resultedProperty.returnTypeRef },
-            mutator = { _, value -> resultedProperty.replaceReturnTypeRef(value) },
-        )
-
-        add(
-            provider = { resultedProperty.controlFlowGraphReference },
-            mutator = { _, value -> resultedProperty.replaceControlFlowGraphReference(value) },
-        )
+    private val RESULT_PROPERTY: StateKeeper<FirProperty, FirDesignation> = stateKeeper { property, scriptDesignation ->
+        add(FirProperty::bodyResolveState, FirProperty::replaceBodyResolveState)
+        add(FirProperty::returnTypeRef, FirProperty::replaceReturnTypeRef)
+        add(FirProperty::controlFlowGraphReference, FirProperty::replaceControlFlowGraphReference)
+        add(FirProperty::initializer, FirProperty::replaceInitializer) { _ ->
+            FirLazyBodiesCalculator.createResultPropertyInitializer(scriptDesignation, property)
+        }
     }
 
     val CODE_FRAGMENT: StateKeeper<FirCodeFragment, FirDesignation> = stateKeeper { _, _ ->
@@ -572,21 +556,6 @@ private fun delegatedConstructorCallGuard(fir: FirDelegatedConstructorCall): Fir
     }
 }
 
-private fun requireSameSize(old: List<FirDeclaration>, new: List<FirDeclaration>) {
-    requireWithAttachment(
-        condition = old.size == new.size,
-        message = { "The number of declarations are different" }
-    ) {
-        withEntryGroup("originalStatements") {
-            old.forEachIndexed { index, statement -> withFirEntry("statement$index", statement) }
-        }
-
-        withEntryGroup("newStatements") {
-            new.forEachIndexed { index, statement -> withFirEntry("statement$index", statement) }
-        }
-    }
-}
-
 private class LLFirCodeFragmentContext(
     override val towerDataContext: FirTowerDataContext,
     override val smartCasts: Map<RealVariable, Set<ConeKotlinType>>

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2024 JetBrains s.r.o. and Kotlin Programming Language contributors.
  * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
  */
 
@@ -1304,24 +1304,15 @@ open class PsiRawFirBuilder(
                         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)
-                                        containingDeclarationSymbol = scriptSymbol
-                                    }
+                                val initializer = buildAnonymousInitializer(
+                                    initializer = declaration,
+                                    containingDeclarationSymbol = scriptSymbol,
+                                    // the last one need to be analyzed in script configurator to decide on result property
+                                    // therefore no lazy conversion in this case
+                                    allowLazyBody = !isLast,
                                 )
+
+                                declarations.add(initializer)
                             }
                             is KtDestructuringDeclaration -> {
                                 val destructuringContainerVar = generateTemporaryVariable(
@@ -2237,24 +2228,34 @@ open class PsiRawFirBuilder(
         }
 
         override fun visitAnonymousInitializer(initializer: KtAnonymousInitializer, data: FirElement?): FirElement {
-            return buildAnonymousInitializer(initializer, null)
+            return buildAnonymousInitializer(initializer, containingDeclarationSymbol = null)
         }
 
-        private fun buildAnonymousInitializer(initializer: KtAnonymousInitializer, containingDeclarationSymbol: FirBasedSymbol<*>?) =
-            buildAnonymousInitializer {
-                withContainerSymbol(symbol) {
-                    source = initializer.toFirSourceElement()
-                    moduleData = baseModuleData
-                    origin = FirDeclarationOrigin.Source
-                    body = buildOrLazyBlock {
+        protected fun buildAnonymousInitializer(
+            initializer: KtAnonymousInitializer,
+            containingDeclarationSymbol: FirBasedSymbol<*>?,
+            allowLazyBody: Boolean = true,
+        ) = buildAnonymousInitializer {
+            withContainerSymbol(symbol) {
+                source = initializer.toFirSourceElement()
+                moduleData = baseModuleData
+                origin = FirDeclarationOrigin.Source
+                body = if (allowLazyBody) {
+                    buildOrLazyBlock {
                         withForcedLocalContext {
                             initializer.body.toFirBlock()
                         }
                     }
-                    this.containingDeclarationSymbol = containingDeclarationSymbol
-                    initializer.extractAnnotationsTo(this)
+                } else {
+                    withForcedLocalContext {
+                        initializer.body.toFirBlock()
+                    }
                 }
+
+                this.containingDeclarationSymbol = containingDeclarationSymbol
+                initializer.extractAnnotationsTo(this)
             }
+        }
 
         override fun visitProperty(property: KtProperty, data: FirElement?): FirElement {
             return property.toFirProperty(

compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/scriptStatementLevelDestructuringWithAnnotation.annotationOwners.txt

@@ -1,23 +1,23 @@
 @DestrAnno(<strcat>(String(destr 1 ), prop#))
-owner -> FirScriptSymbol util.<script-scriptStatementLevelDestructuringWithAnnotation.kts>
+owner -> FirAnonymousInitializerSymbol <init>
 context -> FirPropertySymbol <local>/<destruct> from FirAnonymousInitializerSymbol <init> from FirScriptSymbol util.<script-scriptStatementLevelDestructuringWithAnnotation.kts> from FirFileSymbol scriptStatementLevelDestructuringWithAnnotation.kts
 
 @LeftAnno(<strcat>(String(a ), prop#))
-owner -> FirScriptSymbol util.<script-scriptStatementLevelDestructuringWithAnnotation.kts>
+owner -> FirAnonymousInitializerSymbol <init>
 context -> FirPropertySymbol <local>/a from FirAnonymousInitializerSymbol <init> from FirScriptSymbol util.<script-scriptStatementLevelDestructuringWithAnnotation.kts> from FirFileSymbol scriptStatementLevelDestructuringWithAnnotation.kts
 
 @RightAnno(<strcat>(String(b ), prop#))
-owner -> FirScriptSymbol util.<script-scriptStatementLevelDestructuringWithAnnotation.kts>
+owner -> FirAnonymousInitializerSymbol <init>
 context -> FirPropertySymbol <local>/b from FirAnonymousInitializerSymbol <init> from FirScriptSymbol util.<script-scriptStatementLevelDestructuringWithAnnotation.kts> from FirFileSymbol scriptStatementLevelDestructuringWithAnnotation.kts
 
 @Destr2Anno(<strcat>(String(destr 1 ), prop#))
-owner -> FirScriptSymbol util.<script-scriptStatementLevelDestructuringWithAnnotation.kts>
+owner -> FirAnonymousInitializerSymbol <init>
 context -> FirPropertySymbol <local>/<destruct> from FirAnonymousInitializerSymbol <init> from FirScriptSymbol util.<script-scriptStatementLevelDestructuringWithAnnotation.kts> from FirFileSymbol scriptStatementLevelDestructuringWithAnnotation.kts
 
 @SecondLeftAnno(<strcat>(String(c ), prop#))
-owner -> FirScriptSymbol util.<script-scriptStatementLevelDestructuringWithAnnotation.kts>
+owner -> FirAnonymousInitializerSymbol <init>
 context -> FirPropertySymbol <local>/c from FirAnonymousInitializerSymbol <init> from FirScriptSymbol util.<script-scriptStatementLevelDestructuringWithAnnotation.kts> from FirFileSymbol scriptStatementLevelDestructuringWithAnnotation.kts
 
 @SecondRightAnno(<strcat>(String(d ), prop#))
-owner -> FirScriptSymbol util.<script-scriptStatementLevelDestructuringWithAnnotation.kts>
+owner -> FirAnonymousInitializerSymbol <init>
 context -> FirPropertySymbol <local>/d from FirAnonymousInitializerSymbol <init> from FirScriptSymbol util.<script-scriptStatementLevelDestructuringWithAnnotation.kts> from FirFileSymbol scriptStatementLevelDestructuringWithAnnotation.kts

compiler/fir/raw-fir/psi2fir/testData/rawBuilder/declarations/scriptStatementLevelDestructuringWithAnnotationAsLastStatement.annotationOwners.txt

@@ -1,23 +1,23 @@
 @DestrAnno(<strcat>(String(destr 1 ), prop#))
-owner -> FirScriptSymbol util.<script-scriptStatementLevelDestructuringWithAnnotationAsLastStatement.kts>
+owner -> FirAnonymousInitializerSymbol <init>
 context -> FirPropertySymbol <local>/<destruct> from FirAnonymousInitializerSymbol <init> from FirScriptSymbol util.<script-scriptStatementLevelDestructuringWithAnnotationAsLastStatement.kts> from FirFileSymbol scriptStatementLevelDestructuringWithAnnotationAsLastStatement.kts
 
 @LeftAnno(<strcat>(String(a ), prop#))
-owner -> FirScriptSymbol util.<script-scriptStatementLevelDestructuringWithAnnotationAsLastStatement.kts>
+owner -> FirAnonymousInitializerSymbol <init>
 context -> FirPropertySymbol <local>/a from FirAnonymousInitializerSymbol <init> from FirScriptSymbol util.<script-scriptStatementLevelDestructuringWithAnnotationAsLastStatement.kts> from FirFileSymbol scriptStatementLevelDestructuringWithAnnotationAsLastStatement.kts
 
 @RightAnno(<strcat>(String(b ), prop#))
-owner -> FirScriptSymbol util.<script-scriptStatementLevelDestructuringWithAnnotationAsLastStatement.kts>
+owner -> FirAnonymousInitializerSymbol <init>
 context -> FirPropertySymbol <local>/b from FirAnonymousInitializerSymbol <init> from FirScriptSymbol util.<script-scriptStatementLevelDestructuringWithAnnotationAsLastStatement.kts> from FirFileSymbol scriptStatementLevelDestructuringWithAnnotationAsLastStatement.kts
 
 @Destr2Anno(<strcat>(String(destr 1 ), prop#))
-owner -> FirScriptSymbol util.<script-scriptStatementLevelDestructuringWithAnnotationAsLastStatement.kts>
+owner -> FirAnonymousInitializerSymbol <init>
 context -> FirPropertySymbol <local>/<destruct> from FirAnonymousInitializerSymbol <init> from FirScriptSymbol util.<script-scriptStatementLevelDestructuringWithAnnotationAsLastStatement.kts> from FirFileSymbol scriptStatementLevelDestructuringWithAnnotationAsLastStatement.kts
 
 @SecondLeftAnno(<strcat>(String(c ), prop#))
-owner -> FirScriptSymbol util.<script-scriptStatementLevelDestructuringWithAnnotationAsLastStatement.kts>
+owner -> FirAnonymousInitializerSymbol <init>
 context -> FirPropertySymbol <local>/c from FirAnonymousInitializerSymbol <init> from FirScriptSymbol util.<script-scriptStatementLevelDestructuringWithAnnotationAsLastStatement.kts> from FirFileSymbol scriptStatementLevelDestructuringWithAnnotationAsLastStatement.kts
 
 @SecondRightAnno(<strcat>(String(d ), prop#))
-owner -> FirScriptSymbol util.<script-scriptStatementLevelDestructuringWithAnnotationAsLastStatement.kts>
+owner -> FirAnonymousInitializerSymbol <init>
 context -> FirPropertySymbol <local>/d from FirAnonymousInitializerSymbol <init> from FirScriptSymbol util.<script-scriptStatementLevelDestructuringWithAnnotationAsLastStatement.kts> from FirFileSymbol scriptStatementLevelDestructuringWithAnnotationAsLastStatement.kts