[LL API] Fix on-air analysis for scripts (KTIJ-21108)

The implementation is rather limited. Advanced cases, such as
replacing a script with another one with different import directives,
won't work because of non-trivial relationship between 'FirFile' and
'FirScript'.

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

@@ -30,6 +30,15 @@ internal object DeclarationCopyBuilder {
         resolvePhase = minOf(this.resolvePhase, FirResolvePhase.IMPORTS) //TODO move into initDeclaration?
     }
 
+    fun FirScript.withBodyFrom(
+        scriptWithBody: FirScript
+    ): FirScript = buildScriptCopy(this) {
+        statements.clear()
+        statements.addAll(scriptWithBody.statements)
+        symbol = scriptWithBody.symbol
+        initDeclaration(this@withBodyFrom, scriptWithBody)
+    }
+
     fun FirProperty.withBodyFrom(propertyWithBody: FirProperty): FirProperty {
         val newSetter = getAccessorToUse(this, propertyWithBody) { it.setter }
         val newGetter = getAccessorToUse(this, propertyWithBody) { it.getter }
@@ -88,6 +97,15 @@ internal object DeclarationCopyBuilder {
         moduleData = originalDeclaration.moduleData
     }
 
+    private fun FirScriptBuilder.initDeclaration(
+        originalDeclaration: FirDeclaration,
+        builtDeclaration: FirDeclaration,
+    ) {
+        resolvePhase = minOf(originalDeclaration.resolvePhase, FirResolvePhase.DECLARATIONS)
+        source = builtDeclaration.source
+        moduleData = originalDeclaration.moduleData
+    }
+
     private fun FirFunction.reassignAllReturnTargets(from: FirFunction) {
         this.accept(object : FirVisitorVoid() {
             override fun visitElement(element: FirElement) {

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

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirResolveSessionDepend
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.FileTowerProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.FirTowerContextProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.FirTowerDataContextAllElementsCollector
+import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.getNonLocalContainingOrThisDeclaration
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.FirElementsRecorder
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.KtToFirMapping
 import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.RawFirNonLocalDeclarationBuilder
@@ -43,20 +44,10 @@ import org.jetbrains.kotlin.fir.symbols.lazyResolveToPhase
 import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
 import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
 import org.jetbrains.kotlin.psi.psiUtil.isAncestor
-import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf
 
 object LowLevelFirApiFacadeForResolveOnAir {
-    private fun findNonLocalParentMaybeSelf(position: KtElement): KtNamedDeclaration? {
+    private fun recordOriginalDeclaration(targetDeclaration: KtDeclaration, originalDeclaration: KtDeclaration) {
-        return position.parentsWithSelf
-            .filterIsInstance<KtNamedDeclaration>()
-            .filter { it is KtNamedFunction || it is KtProperty || (it is KtClassOrObject && it !is KtEnumEntry) || it is KtTypeAlias }
-            .filter { !KtPsiUtil.isLocal(it) && it.containingClassOrObject !is KtEnumEntry }
-            .firstOrNull()
-    }
-
-    private fun recordOriginalDeclaration(targetDeclaration: KtNamedDeclaration, originalDeclaration: KtNamedDeclaration) {
         require(originalDeclaration.containingKtFile !== targetDeclaration.containingKtFile)
         val originalDeclarationParents = originalDeclaration.parentsOfType<KtDeclaration>().toList()
         val fakeDeclarationParents = targetDeclaration.parentsOfType<KtDeclaration>().toList()
@@ -145,13 +136,13 @@ object LowLevelFirApiFacadeForResolveOnAir {
         require(originalFirResolveSession is LLFirResolvableResolveSession)
         require(elementToAnalyze !is KtFile) { "KtFile for dependency element not supported" }
 
-        val dependencyNonLocalDeclaration = findNonLocalParentMaybeSelf(elementToAnalyze)
+        val dependencyNonLocalDeclaration = elementToAnalyze.getNonLocalContainingOrThisDeclaration() as? KtNamedDeclaration
-            ?: return LLFirResolveSessionDepended(
-                originalFirResolveSession,
-                FileTowerProvider(elementToAnalyze.containingKtFile, onAirGetTowerContextForFile(originalFirResolveSession, originalKtFile)),
-                ktToFirMapping = null
-            )
 
+        if (dependencyNonLocalDeclaration == null) {
+            val towerDataContext = onAirGetTowerContextForFile(originalFirResolveSession, originalKtFile)
+            val fileTowerProvider = FileTowerProvider(elementToAnalyze.containingKtFile, towerDataContext)
+            return LLFirResolveSessionDepended(originalFirResolveSession, fileTowerProvider, ktToFirMapping = null)
+        }
 
         val sameDeclarationInOriginalFile = PsiTreeUtil.findSameElementInCopy(dependencyNonLocalDeclaration, originalKtFile)
             ?: buildErrorWithAttachment("Cannot find original function matching") {
@@ -214,7 +205,7 @@ object LowLevelFirApiFacadeForResolveOnAir {
         onAirCreatedDeclaration: Boolean,
         collector: FirTowerDataContextCollector? = null,
     ): FirElement {
-        val nonLocalDeclaration = findNonLocalParentMaybeSelf(replacement.from)
+        val nonLocalDeclaration = replacement.from.getNonLocalContainingOrThisDeclaration()
         val originalFirFile = firResolveSession.getOrBuildFirFile(replacement.from.containingKtFile)
 
         if (nonLocalDeclaration == null) {
@@ -255,6 +246,8 @@ object LowLevelFirApiFacadeForResolveOnAir {
                         originalDeclaration.withBodyFrom(newDeclarationWithReplacement as FirProperty)
                     is FirRegularClass ->
                         originalDeclaration.withBodyFrom(newDeclarationWithReplacement as FirRegularClass)
+                    is FirScript ->
+                        originalDeclaration.withBodyFrom(newDeclarationWithReplacement as FirScript)
                     is FirTypeAlias -> newDeclarationWithReplacement
                     else -> error("Not supported type ${originalDeclaration::class.simpleName}")
                 }
@@ -278,26 +271,23 @@ object LowLevelFirApiFacadeForResolveOnAir {
 
     }
 
-    private fun isInBodyReplacement(ktDeclaration: KtDeclaration, replacement: RawFirReplacement): Boolean = when (ktDeclaration) {
+    private fun isInBodyReplacement(ktDeclaration: KtDeclaration, replacement: RawFirReplacement): Boolean {
-        is KtNamedFunction ->
+        fun check(container: KtElement?): Boolean {
-            ktDeclaration.bodyBlockExpression?.let { it.isAncestor(replacement.from, true) } ?: false
+            return container != null && container.isAncestor(replacement.from, true)
-        is KtProperty -> {
+        }
-            val insideGetterBody = ktDeclaration.getter?.bodyBlockExpression?.let {
+
-                it.isAncestor(replacement.from, true)
+        return when (ktDeclaration) {
-            } ?: false
+            is KtNamedFunction -> check(ktDeclaration.bodyBlockExpression)
-
+            is KtProperty -> {
-            val insideGetterOrSetterBody = insideGetterBody || ktDeclaration.setter?.bodyBlockExpression?.let {
+                check(ktDeclaration.getter?.bodyBlockExpression)
-                it.isAncestor(replacement.from, true)
+                        || check(ktDeclaration.setter?.bodyBlockExpression)
-            } ?: false
+                        || check(ktDeclaration.initializer)
-
+            }
-            insideGetterOrSetterBody || ktDeclaration.initializer?.let {
+            is KtClassOrObject -> check(ktDeclaration.body)
-                it.isAncestor(replacement.from, true)
+            is KtScript -> check(ktDeclaration.blockExpression)
-            } ?: false
+            is KtTypeAlias -> false
+            else -> error("Not supported type ${ktDeclaration::class.simpleName}")
         }
-        is KtClassOrObject ->
-            ktDeclaration.body?.let { it.isAncestor(replacement.from, true) } ?: false
-        is KtTypeAlias -> false
-        else -> error("Not supported type ${ktDeclaration::class.simpleName}")
     }
 
     fun onAirResolveTypeInPlace(

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

@@ -10,7 +10,7 @@ import org.jetbrains.kotlin.psi.*
 internal data class RawFirReplacement(val from: KtElement, val to: KtElement) {
     companion object {
         fun isApplicableForReplacement(element: KtElement) = when (element) {
-            is KtFile, is KtClassInitializer, is KtClassOrObject, is KtObjectLiteralExpression, is KtTypeAlias,
+            is KtFile, is KtScript, is KtClassInitializer, is KtClassOrObject, is KtObjectLiteralExpression, is KtTypeAlias,
             is KtNamedFunction, is KtLambdaExpression, is KtAnonymousInitializer, is KtProperty, is KtTypeReference,
             is KtAnnotationEntry, is KtTypeParameter, is KtTypeProjection, is KtParameter, is KtBlockExpression,
             is KtSimpleNameExpression, is KtConstantExpression, is KtStringTemplateExpression, is KtReturnExpression,

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

@@ -1049,7 +1049,11 @@ open class RawFirBuilder(
                     declarations += when (declaration) {
                         is KtScript -> {
                             require(file.declarations.size == 1) { "Expect the script to be the only declaration in the file $name" }
-                            convertScript(declaration, this)
+                            convertScript(declaration, this.name) {
+                                for (configurator in baseSession.extensionService.scriptConfigurators) {
+                                    with(configurator) { configure(this@buildFile) }
+                                }
+                            }
                         }
                         is KtDestructuringDeclaration -> buildErrorTopLevelDestructuringDeclaration(declaration.toFirSourceElement())
                         else -> declaration.convert()
@@ -1062,12 +1066,12 @@ open class RawFirBuilder(
             }
         }
 
-        private fun convertScript(script: KtScript, containingFile: FirFileBuilder): FirScript {
+        private fun convertScript(script: KtScript, fileName: String, setup: FirScriptBuilder.() -> Unit = {}): FirScript {
             return buildScript {
                 source = script.toFirSourceElement()
                 moduleData = baseModuleData
                 origin = FirDeclarationOrigin.Source
-                name = Name.special("<script-${containingFile.name}>")
+                name = Name.special("<script-$fileName>")
                 symbol = FirScriptSymbol(context.packageFqName.child(name))
                 for (declaration in script.declarations) {
                     when (declaration) {
@@ -1079,12 +1083,13 @@ open class RawFirBuilder(
                         }
                     }
                 }
-                baseSession.extensionService.scriptConfigurators.forEach { with(it) { configure(containingFile) } }
+                setup()
             }
         }
 
         override fun visitScript(script: KtScript, data: Unit?): FirElement {
-            error("should not be here")
+            val fileName = script.containingKtFile.name
+            return convertScript(script, fileName)
         }
 
         protected fun KtEnumEntry.toFirEnumEntry(

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

@@ -70,3 +70,23 @@ inline fun buildScript(init: FirScriptBuilder.() -> Unit): FirScript {
     }
     return FirScriptBuilder().apply(init).build()
 }
+
+@OptIn(ExperimentalContracts::class)
+inline fun buildScriptCopy(original: FirScript, init: FirScriptBuilder.() -> Unit): FirScript {
+    contract {
+        callsInPlace(init, kotlin.contracts.InvocationKind.EXACTLY_ONCE)
+    }
+    val copyBuilder = FirScriptBuilder()
+    copyBuilder.source = original.source
+    copyBuilder.resolvePhase = original.resolvePhase
+    copyBuilder.annotations.addAll(original.annotations)
+    copyBuilder.moduleData = original.moduleData
+    copyBuilder.origin = original.origin
+    copyBuilder.attributes = original.attributes.copy()
+    copyBuilder.name = original.name
+    copyBuilder.statements.addAll(original.statements)
+    copyBuilder.symbol = original.symbol
+    copyBuilder.parameters.addAll(original.parameters)
+    copyBuilder.contextReceivers.addAll(original.contextReceivers)
+    return copyBuilder.apply(init).build()
+}

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

@@ -357,6 +357,10 @@ object BuilderConfigurator : AbstractBuilderConfigurator<FirTreeBuilder>(FirTree
             fields from resolvedQualifier
         }
 
+        builder(script) {
+            withCopy()
+        }
+
         builder(resolvedQualifier) {
             parents += abstractResolvedQualifierBuilder
             defaultFalse("isNullableLHSForCallableReference", "isFullyQualified")