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