From 6e016ce041842a38fe93377cea923d49250d0f2f Mon Sep 17 00:00:00 2001 From: Vladimir Dolzhenko Date: Wed, 29 Jul 2020 16:51:19 +0200 Subject: [PATCH] ktFile has to return script declaration from stub as well --- .../src/org/jetbrains/kotlin/psi/KtFile.kt | 11 +++++++--- .../kotlin/idea/stubs/KotlinStubsTest.java | 20 +++++++++++++++++++ 2 files changed, 28 insertions(+), 3 deletions(-) diff --git a/compiler/psi/src/org/jetbrains/kotlin/psi/KtFile.kt b/compiler/psi/src/org/jetbrains/kotlin/psi/KtFile.kt index 9ecb453f1e8..fca0f431279 100644 --- a/compiler/psi/src/org/jetbrains/kotlin/psi/KtFile.kt +++ b/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 { val stub = stub - return stub?.getChildrenByType(KtStubElementTypes.DECLARATION_TYPES, KtDeclaration.ARRAY_FACTORY)?.toList() - ?: PsiTreeUtil.getChildrenOfTypeAsList(this, KtDeclaration::class.java) + return stub?.getChildrenByType(FILE_DECLARATION_TYPES, KtDeclaration.ARRAY_FACTORY)?.toList() + ?: PsiTreeUtil.getChildrenOfTypeAsList(this, KtDeclaration::class.java) } fun >> 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)) + } } diff --git a/idea/tests/org/jetbrains/kotlin/idea/stubs/KotlinStubsTest.java b/idea/tests/org/jetbrains/kotlin/idea/stubs/KotlinStubsTest.java index cb49346f43f..f6fe8b0dc78 100644 --- a/idea/tests/org/jetbrains/kotlin/idea/stubs/KotlinStubsTest.java +++ b/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 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 stubDeclarations = ktFile.getDeclarations(); + assertEquals(1, stubDeclarations.size()); + assertTrue(stubDeclarations.get(0) instanceof KtScript); + } }