[LL API] Recognize scripts in Analysis API

Although the proper support for script resolution is yet to arrive
in K2, some initial steps are required to make LL API understand
such a concept.

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/element/builder/FirElementBuilder.kt

@@ -12,15 +12,13 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirModuleResolveCompone
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LLFirResolveSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.FileStructureElement
 import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.declarationCanBeLazilyResolved
-import org.jetbrains.kotlin.analysis.low.level.api.fir.util.isNonAnonymousClassOrObject
 import org.jetbrains.kotlin.analysis.utils.printer.getElementTextInContext
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.declarations.FirFile
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.symbols.lazyResolveToPhase
 import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
+import org.jetbrains.kotlin.psi.psiUtil.*
-import org.jetbrains.kotlin.psi.psiUtil.isAncestor
 import org.jetbrains.kotlin.psi2ir.deparenthesize
 
 
@@ -98,28 +96,44 @@ internal class FirElementBuilder(
     }
 }
 
-// TODO: simplify
+internal fun PsiElement.getNonLocalContainingOrThisDeclaration(predicate: (KtDeclaration) -> Boolean = { true }): KtDeclaration? {
-internal inline fun PsiElement.getNonLocalContainingOrThisDeclaration(predicate: (KtDeclaration) -> Boolean = { true }): KtDeclaration? {
+    var candidate: KtDeclaration? = null
-    var container: PsiElement? = this
+
-    while (container != null && container !is KtFile) {
+    fun propose(declaration: KtDeclaration) {
-        if (container is KtNamedDeclaration
+        if (candidate == null) {
-            && (container.isNonAnonymousClassOrObject() || container is KtDeclarationWithBody || container is KtProperty || container is KtTypeAlias)
+            candidate = declaration
-            && container !is KtPrimaryConstructor
-            && declarationCanBeLazilyResolved(container)
-            && container !is KtFunctionLiteral
-            && container.containingClassOrObject !is KtEnumEntry
-            && predicate(container)
-        ) {
-            return container
         }
-        if (container is KtDestructuringDeclaration && container.parent is KtFile) {
-            return container
-        }
-        container = container.parent
     }
-    return null
+
+    for (parent in parentsWithSelf) {
+        if (candidate != null) {
+            if (parent is KtEnumEntry || parent is KtCallableDeclaration) {
+                // Candidate turned to be local. Let's find another one
+                candidate = null
+            }
+        }
+
+        when (parent) {
+            is KtScript -> propose(parent)
+            is KtDestructuringDeclaration -> propose(parent)
+            is KtNamedDeclaration -> {
+                val isKindApplicable = when (parent) {
+                    is KtClassOrObject -> !parent.isObjectLiteral()
+                    is KtDeclarationWithBody, is KtProperty, is KtTypeAlias -> true
+                    else -> false
+                }
+
+                if (isKindApplicable && declarationCanBeLazilyResolved(parent) && predicate(parent)) {
+                    propose(parent)
+                }
+            }
+        }
+    }
+
+    return candidate
 }
 
+@Suppress("unused") // Used in the IDE plugin
 fun PsiElement.getNonLocalContainingInBodyDeclarationWith(): KtDeclaration? =
     getNonLocalContainingOrThisDeclaration { declaration ->
         when (declaration) {

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/builder/LLFirFileBuilder.kt

@@ -10,8 +10,6 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirModuleResolveCompone
 import org.jetbrains.kotlin.fir.builder.RawFirBuilder
 import org.jetbrains.kotlin.fir.builder.BodyBuildingMode
 import org.jetbrains.kotlin.fir.declarations.FirFile
-import org.jetbrains.kotlin.fir.scopes.FirScopeProvider
-import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirPhaseRunner
 import org.jetbrains.kotlin.psi.KtFile
 
 /**
@@ -22,10 +20,20 @@ internal class LLFirFileBuilder(
     val moduleComponents: LLFirModuleResolveComponents,
 ) {
     fun buildRawFirFileWithCaching(ktFile: KtFile): FirFile = moduleComponents.cache.fileCached(ktFile) {
+        val bodyBuildingMode = when {
+            ktFile.isScript() -> {
+                // As 'FirScript' content is never transformed, lazy bodies are not replaced with calculated ones even on BODY_RESOLVE.
+                // Such behavior breaks file structure mapping computation.
+                // TODO: remove this clause when proper support for scripts is implemented in K2.
+                BodyBuildingMode.NORMAL
+            }
+            else -> BodyBuildingMode.LAZY_BODIES
+        }
+
         RawFirBuilder(
             moduleComponents.session,
             moduleComponents.scopeProvider,
-            bodyBuildingMode = BodyBuildingMode.LAZY_BODIES
+            bodyBuildingMode = bodyBuildingMode
         ).buildFirFile(ktFile)
     }
 }

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

@@ -17,6 +17,7 @@ fun FirElementWithResolvePhase.getContainingFile(): FirFile? {
     val provider = moduleData.session.firProvider
     return when (this) {
         is FirFile -> this
+        is FirScript -> containingFileSymbol.fir
         is FirFileAnnotationsContainer -> containingFileSymbol.fir
         is FirTypeParameter -> containingDeclarationSymbol.fir.getContainingFile()
         is FirPropertyAccessor -> propertySymbol.fir.getContainingFile()

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

@@ -9,10 +9,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.throwUnexpectedFirEle
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.getNonLocalContainingOrThisDeclaration
 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.FirClassLikeDeclaration
+import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.declarations.FirDeclaration
-import org.jetbrains.kotlin.fir.declarations.FirFile
-import org.jetbrains.kotlin.fir.declarations.FirRegularClass
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.realPsi
 import org.jetbrains.kotlin.fir.resolve.providers.FirProvider
@@ -105,6 +102,7 @@ private fun KtDeclaration.findSourceNonLocalFirDeclarationByProvider(
             val firFile = containerFirFile ?: firFileBuilder.buildRawFirFileWithCaching(containingKtFile)
             firFile.declarations.firstOrNull { it.psi == this }
         }
+        this is KtScript -> containerFirFile?.declarations?.singleOrNull { it is FirScript }
         else -> errorWithFirSpecificEntries("Invalid container", psi = this)
     }
     return candidate?.takeIf { it.realPsi == this }

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

@@ -1079,6 +1079,7 @@ open class RawFirBuilder(
                 origin = FirDeclarationOrigin.Source
                 name = Name.special("<script-${containingFile.name}>")
                 symbol = FirScriptSymbol(context.packageFqName.child(name))
+                containingFileSymbol = containingFile.symbol
                 for (declaration in script.declarations) {
                     when (declaration) {
                         is KtScriptInitializer -> {

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

@@ -10,6 +10,7 @@ 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.symbols.impl.FirFileSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirScriptSymbol
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.fir.visitors.*
@@ -29,6 +30,7 @@ abstract class FirScript : FirDeclaration() {
     abstract val name: Name
     abstract val statements: List<FirStatement>
     abstract override val symbol: FirScriptSymbol
+    abstract val containingFileSymbol: FirFileSymbol
     abstract val parameters: List<FirVariable>
     abstract val contextReceivers: List<FirContextReceiver>
 

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

@@ -22,6 +22,7 @@ import org.jetbrains.kotlin.fir.declarations.FirVariable
 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.FirFileSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirScriptSymbol
 import org.jetbrains.kotlin.fir.visitors.*
 import org.jetbrains.kotlin.name.Name
@@ -42,6 +43,7 @@ class FirScriptBuilder : FirAnnotationContainerBuilder {
     lateinit var name: Name
     val statements: MutableList<FirStatement> = mutableListOf()
     lateinit var symbol: FirScriptSymbol
+    lateinit var containingFileSymbol: FirFileSymbol
     val parameters: MutableList<FirVariable> = mutableListOf()
     val contextReceivers: MutableList<FirContextReceiver> = mutableListOf()
 
@@ -56,6 +58,7 @@ class FirScriptBuilder : FirAnnotationContainerBuilder {
             name,
             statements,
             symbol,
+            containingFileSymbol,
             parameters,
             contextReceivers.toMutableOrEmpty(),
         )

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

@@ -17,6 +17,7 @@ import org.jetbrains.kotlin.fir.declarations.FirScript
 import org.jetbrains.kotlin.fir.declarations.FirVariable
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
 import org.jetbrains.kotlin.fir.expressions.FirStatement
+import org.jetbrains.kotlin.fir.symbols.impl.FirFileSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirScriptSymbol
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.fir.visitors.*
@@ -39,6 +40,7 @@ internal class FirScriptImpl(
     override val name: Name,
     override val statements: MutableList<FirStatement>,
     override val symbol: FirScriptSymbol,
+    override val containingFileSymbol: FirFileSymbol,
     override val parameters: MutableList<FirVariable>,
     override var contextReceivers: MutableOrEmptyList<FirContextReceiver>,
 ) : FirScript() {

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

@@ -475,6 +475,9 @@ object NodeConfigurator : AbstractFieldConfigurator<FirTreeBuilder>(FirTreeBuild
             +name
             +fieldList(statement).withTransform()
             +symbol("FirScriptSymbol")
+            +field("containingFileSymbol", type("fir.symbols.impl", "FirFileSymbol"), argument = null).apply {
+                withBindThis = false
+            }
             +FieldList("parameters", variable, withReplace = false)
             +fieldList(contextReceiver, useMutableOrEmpty = true)
         }