[LL API] Handle 'KtScript' consistently in file-based provider

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

@@ -16,6 +16,19 @@ import org.jetbrains.kotlin.utils.addIfNotNull
 import org.jetbrains.kotlin.utils.yieldIfNotNull
 
 internal class FileBasedKotlinDeclarationProvider(private val kotlinFile: KtFile) : KotlinDeclarationProvider() {
+    private val topLevelDeclarations: Sequence<KtDeclaration>
+        get() {
+            return sequence {
+                for (child in kotlinFile.declarations) {
+                    if (child is KtScript) {
+                        yieldAll(child.declarations)
+                    } else {
+                        yield(child)
+                    }
+                }
+            }
+        }
+
     override fun getClassLikeDeclarationByClassId(classId: ClassId): KtClassLikeDeclaration? {
         return getClassLikeDeclarationsByClassId(classId).firstOrNull()
     }
@@ -39,7 +52,7 @@ internal class FileBasedKotlinDeclarationProvider(private val kotlinFile: KtFile
             val tasks = ArrayDeque<Task>()
 
             val startingChunks = classId.relativeClassName.pathSegments()
-            for (declaration in kotlinFile.declarations) {
+            for (declaration in topLevelDeclarations) {
                 tasks.addLast(Task(startingChunks, declaration))
             }
 
@@ -49,13 +62,6 @@ 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
                 }
@@ -113,7 +119,7 @@ internal class FileBasedKotlinDeclarationProvider(private val kotlinFile: KtFile
     override fun findFilesForFacade(facadeFqName: FqName): Collection<KtFile> {
         if (kotlinFile.javaFileFacadeFqName != facadeFqName) return emptyList()
 
-        for (declaration in kotlinFile.declarations) {
+        for (declaration in topLevelDeclarations) {
             if (declaration !is KtClassLikeDeclaration) {
                 return listOf(kotlinFile)
             }
@@ -135,7 +141,7 @@ internal class FileBasedKotlinDeclarationProvider(private val kotlinFile: KtFile
         }
 
         return buildList {
-            for (declaration in kotlinFile.declarations) {
+            for (declaration in topLevelDeclarations) {
                 if (declaration is T && declaration.nameAsName == name) {
                     add(declaration)
                 }
@@ -149,7 +155,7 @@ internal class FileBasedKotlinDeclarationProvider(private val kotlinFile: KtFile
         }
 
         return buildSet {
-            for (declaration in kotlinFile.declarations) {
+            for (declaration in topLevelDeclarations) {
                 if (declaration is T) {
                     addIfNotNull(declaration.nameAsName)
                 }

analysis/low-level-api-fir/testdata/fileBasedDeclarationProvider/script.kts

@@ -0,0 +1,22 @@
+// CLASS: test/app/Foo
+// CLASS: test/app/Foo.Bar
+// CLASS: test/app/Foo.Bar.Baz
+// TYPE_ALIAS: test/app/MyString
+// FUNCTION: test/app/#foo
+// PROPERTY: test/app/#bar
+
+package test.app
+
+interface Foo {
+    class Bar {
+        object Baz
+    }
+
+    fun nested() {}
+}
+
+typealias MyString = String
+
+fun foo() {}
+
+val bar: String = "bar"

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/FileBasedKotlinDeclarationProviderTestGenerated.java

@@ -21,7 +21,7 @@ import java.util.regex.Pattern;
 public class FileBasedKotlinDeclarationProviderTestGenerated extends AbstractFileBasedKotlinDeclarationProviderTest {
     @Test
     public void testAllFilesPresentInFileBasedDeclarationProvider() throws Exception {
-        KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("analysis/low-level-api-fir/testdata/fileBasedDeclarationProvider"), Pattern.compile("^(.+)\\.kt$"), null, true);
+        KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("analysis/low-level-api-fir/testdata/fileBasedDeclarationProvider"), Pattern.compile("^(.+)\\.(kt|kts)$"), null, true);
     }
 
     @Test
@@ -48,6 +48,12 @@ public class FileBasedKotlinDeclarationProviderTestGenerated extends AbstractFil
         runTest("analysis/low-level-api-fir/testdata/fileBasedDeclarationProvider/sameNames.kt");
     }
 
+    @Test
+    @TestMetadata("script.kts")
+    public void testScript() throws Exception {
+        runTest("analysis/low-level-api-fir/testdata/fileBasedDeclarationProvider/script.kts");
+    }
+
     @Test
     @TestMetadata("simple.kt")
     public void testSimple() throws Exception {

generators/analysis-api-generator/tests/org/jetbrains/kotlin/generators/tests/analysis/api/firLowLevel.kt

@@ -88,7 +88,7 @@ internal fun TestGroupSuite.generateFirLowLevelApiTests() {
         }
 
         testClass<AbstractFileBasedKotlinDeclarationProviderTest> {
-            model("fileBasedDeclarationProvider")
+            model("fileBasedDeclarationProvider", pattern = TestGeneratorUtil.KT_OR_KTS)
         }
     }