[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>>