[LL API] Support FirScript declarations in declaration providers (KTIJ-21108)

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/project/structure/FileBasedKotlinDeclarationProvider.kt

@@ -49,6 +49,13 @@ internal class FileBasedKotlinDeclarationProvider(private val kotlinFile: KtFile
                 val (chunks, element) = tasks.removeFirst()
                 assert(chunks.isNotEmpty())
 
+                if (element is KtScript) {
+                    for (child in element.declarations) {
+                        tasks.addLast(Task(chunks, child))
+                    }
+                    continue
+                }
+
                 if (element !is KtNamedDeclaration || element.nameAsName != chunks[0]) {
                     continue
                 }

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

@@ -7,10 +7,7 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir.util
 
 import org.jetbrains.kotlin.analysis.utils.errors.requireWithAttachmentBuilder
 import org.jetbrains.kotlin.fir.FirElement
-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.declarations.utils.classId
 import org.jetbrains.kotlin.fir.packageFqName
 import org.jetbrains.kotlin.fir.render
@@ -33,11 +30,17 @@ object FirElementFinder {
         val classIdPathSegment = classId.relativeClassName.pathSegments()
         var result: FirClassLikeDeclaration? = null
 
-        fun find(declarations: List<FirDeclaration>, classIdPathIndex: Int) {
+        fun find(declarations: Iterable<FirDeclaration>, classIdPathIndex: Int) {
             if (result != null) return
             val currentClassSegment = classIdPathSegment[classIdPathIndex]
 
             for (subDeclaration in declarations) {
+                if (subDeclaration is FirScript) {
+                    val scriptDeclarations = subDeclaration.statements.asSequence().filterIsInstance<FirDeclaration>()
+                    find(scriptDeclarations.asIterable(), classIdPathIndex)
+                    continue
+                }
+
                 if (subDeclaration is FirClassLikeDeclaration && currentClassSegment == subDeclaration.symbol.name) {
                     if (classIdPathIndex == classIdPathSegment.lastIndex) {
                         result = subDeclaration

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

@@ -61,9 +61,13 @@ internal fun KtElement.findSourceByTraversingWholeTree(
     val firFile = containerFirFile ?: firFileBuilder.buildRawFirFileWithCaching(containingKtFile)
     val originalDeclaration = (this as? KtDeclaration)?.originalDeclaration
     val isDeclaration = this is KtDeclaration
-    return FirElementFinder.findElementIn(firFile, canGoInside = { it is FirRegularClass }) { firDeclaration ->
-        firDeclaration.psi == this || isDeclaration && firDeclaration.psi == originalDeclaration
-    }
+    return FirElementFinder.findElementIn(
+        firFile,
+        canGoInside = { it is FirRegularClass || it is FirScript },
+        predicate = { firDeclaration ->
+            firDeclaration.psi == this || isDeclaration && firDeclaration.psi == originalDeclaration
+        }
+    )
 }
 
 private fun KtDeclaration.findSourceNonLocalFirDeclarationByProvider(

compiler/psi/src/org/jetbrains/kotlin/psi/psiUtil/ClassIdCalculator.kt

@@ -5,7 +5,6 @@
 
 package org.jetbrains.kotlin.psi.psiUtil
 
-import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.SpecialNames
@@ -14,9 +13,9 @@ import org.jetbrains.kotlin.psi.*
 internal object ClassIdCalculator {
     fun calculateClassId(declaration: KtClassLikeDeclaration): ClassId? {
         var ktFile: KtFile? = null
-        var element: PsiElement? = declaration
         val containingClasses = mutableListOf<KtClassLikeDeclaration>()
-        while (element != null) {
+
+        for (element in declaration.parentsWithSelf) {
             when (element) {
                 is KtEnumEntry -> {
                     return null
@@ -24,24 +23,23 @@ internal object ClassIdCalculator {
                 is KtClassLikeDeclaration -> {
                     containingClasses += element
                 }
-                is KtObjectLiteralExpression -> {
-                    return null
-                }
                 is KtFile -> {
                     ktFile = element
                     break
                 }
-                is KtDeclaration -> {
+                is KtScript -> {
+                    // Skip script parent
+                }
+                is KtDeclaration, is KtObjectLiteralExpression -> {
+                    // Local declarations don't have a 'ClassId'
                     return null
                 }
             }
-            element = element.parent
         }
         if (ktFile == null) return null
-        val relativeClassName = FqName.fromSegments(
-            containingClasses.reversed().map { containingClass ->
-                containingClass.name ?: SpecialNames.NO_NAME_PROVIDED.asString()
-            }
+        val relativeClassName = FqName.fromSegments(containingClasses.reversed().map { containingClass ->
+            containingClass.name ?: SpecialNames.NO_NAME_PROVIDED.asString()
+        }
         )
         return ClassId(ktFile.packageFqName, relativeClassName, /*local=*/false)
     }