ktFile has to return script declaration from stub as well

compiler/psi/src/org/jetbrains/kotlin/psi/KtFile.kt

@@ -21,6 +21,7 @@ import com.intellij.openapi.components.ServiceManager
 import com.intellij.openapi.fileTypes.FileType
 import com.intellij.psi.*
 import com.intellij.psi.stubs.StubElement
+import com.intellij.psi.tree.TokenSet
 import com.intellij.psi.util.PsiTreeUtil
 import com.intellij.util.ArrayFactory
 import com.intellij.util.FileContentUtilCore
@@ -138,12 +139,12 @@ open class KtFile(viewProvider: FileViewProvider, val isCompiled: Boolean) :
 
     override fun getFileType(): FileType = KotlinFileType.INSTANCE
 
-    override fun toString(): String = "KtFile: " + name
+    override fun toString(): String = "KtFile: $name"
 
     override fun getDeclarations(): List<KtDeclaration> {
         val stub = stub
-        return stub?.getChildrenByType(KtStubElementTypes.DECLARATION_TYPES, KtDeclaration.ARRAY_FACTORY)?.toList()
+        return stub?.getChildrenByType(FILE_DECLARATION_TYPES, KtDeclaration.ARRAY_FACTORY)?.toList()
-                ?: PsiTreeUtil.getChildrenOfTypeAsList(this, KtDeclaration::class.java)
+            ?: PsiTreeUtil.getChildrenOfTypeAsList(this, KtDeclaration::class.java)
     }
 
     fun <T : KtElementImplStub<out StubElement<*>>> findChildByTypeOrClass(
@@ -264,4 +265,8 @@ open class KtFile(viewProvider: FileViewProvider, val isCompiled: Boolean) :
         // Modification count for Kotlin files is tracked entirely by KotlinCodeBlockModificationListener
         return false
     }
+
+    companion object {
+        val FILE_DECLARATION_TYPES = TokenSet.orSet(KtStubElementTypes.DECLARATION_TYPES, TokenSet.create(KtStubElementTypes.SCRIPT))
+    }
 }

idea/tests/org/jetbrains/kotlin/idea/stubs/KotlinStubsTest.java

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.idea.stubs;
 
+import com.intellij.openapi.application.ApplicationManager;
 import com.intellij.psi.PsiFile;
 import com.intellij.testFramework.LightProjectDescriptor;
 import com.intellij.testFramework.fixtures.LightCodeInsightFixtureTestCase;
@@ -13,6 +14,7 @@ import org.jetbrains.kotlin.idea.test.KotlinLightProjectDescriptor;
 import org.jetbrains.kotlin.psi.KtClass;
 import org.jetbrains.kotlin.psi.KtDeclaration;
 import org.jetbrains.kotlin.psi.KtFile;
+import org.jetbrains.kotlin.psi.KtScript;
 import org.jetbrains.kotlin.psi.stubs.KotlinClassStub;
 import org.jetbrains.kotlin.psi.stubs.elements.KtStubElementTypes;
 import org.jetbrains.kotlin.test.JUnit3WithIdeaConfigurationRunner;
@@ -44,4 +46,22 @@ public class KotlinStubsTest extends LightCodeInsightFixtureTestCase {
         KotlinClassStub stub = KtStubElementTypes.CLASS.createStub(ktClass, null);
         assertEquals(true, stub.isInterface());
     }
+
+    public void testScriptDeclaration() {
+        PsiFile psiFile = myFixture.configureByText("foo.kts", "fun foo() {}");
+        KtFile ktFile = (KtFile) psiFile;
+
+        assertNull("file is parsed from AST", ktFile.getStub());
+        List<KtDeclaration> astDeclarations = ktFile.getDeclarations();
+        assertEquals(1, astDeclarations.size());
+        assertTrue(astDeclarations.get(0) instanceof KtScript);
+
+        // clean up AST
+        ApplicationManager.getApplication().runWriteAction(() -> ktFile.onContentReload());
+
+        assertNotNull("file is parsed from AST", ktFile.getStub());
+        List<KtDeclaration> stubDeclarations = ktFile.getDeclarations();
+        assertEquals(1, stubDeclarations.size());
+        assertTrue(stubDeclarations.get(0) instanceof KtScript);
+    }
 }