[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.isAncestor
+import org.jetbrains.kotlin.psi.psiUtil.*
 import org.jetbrains.kotlin.psi2ir.deparenthesize
 
 
@@ -98,28 +96,44 @@ internal class FirElementBuilder(
     }
 }
 
-// TODO: simplify
-internal inline fun PsiElement.getNonLocalContainingOrThisDeclaration(predicate: (KtDeclaration) -> Boolean = { true }): KtDeclaration? {
-    var container: PsiElement? = this
-    while (container != null && container !is KtFile) {
-        if (container is KtNamedDeclaration
-            && (container.isNonAnonymousClassOrObject() || container is KtDeclarationWithBody || container is KtProperty || container is KtTypeAlias)
-            && container !is KtPrimaryConstructor
-            && declarationCanBeLazilyResolved(container)
-            && container !is KtFunctionLiteral
-            && container.containingClassOrObject !is KtEnumEntry
-            && predicate(container)
-        ) {
-            return container
+internal fun PsiElement.getNonLocalContainingOrThisDeclaration(predicate: (KtDeclaration) -> Boolean = { true }): KtDeclaration? {
+    var candidate: KtDeclaration? = null
+
+    fun propose(declaration: KtDeclaration) {
+        if (candidate == null) {
+            candidate = declaration
         }
-        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.FirDeclaration
-import org.jetbrains.kotlin.fir.declarations.FirFile
-import org.jetbrains.kotlin.fir.declarations.FirRegularClass
+import org.jetbrains.kotlin.fir.declarations.*
 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 }