From 2ecddc996ab792ef9d5311a58c3305b36b9707b5 Mon Sep 17 00:00:00 2001
From: Pavel Talanov <pavel.talanov@jebrains.com>
Date: Wed, 27 Jun 2018 11:01:04 +0200
Subject: [PATCH] Implement support for static resolving to scripts in source
 roots in IDEA

---
 .../kotlin/idea/caches/project/getModuleInfo.kt       | 11 ++++++-----
 .../idea/caches/resolve/KotlinCacheServiceImpl.kt     |  4 ++--
 .../KotlinPackageSourcesMemberNamesIndex.kt           |  9 +++++++--
 3 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/project/getModuleInfo.kt b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/project/getModuleInfo.kt
index 3cf9ea8344c..73c5856da8e 100644
--- a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/project/getModuleInfo.kt
+++ b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/project/getModuleInfo.kt
@@ -205,11 +205,6 @@ private inline fun <T> collectInfosByVirtualFile(
     isScript: Boolean = getScriptDefinition(virtualFile, project) != null,
     onOccurrence: (IdeaModuleInfo?) -> T
 ): T {
-    if (isScript) {
-        getScriptDefinition(virtualFile, project)?.let {
-            onOccurrence(ScriptModuleInfo(project, virtualFile, it))
-        }
-    }
 
     val projectFileIndex = ProjectFileIndex.SERVICE.getInstance(project)
 
@@ -218,6 +213,12 @@ private inline fun <T> collectInfosByVirtualFile(
         onOccurrence(moduleRelatedModuleInfo)
     }
 
+    if (moduleRelatedModuleInfo == null && isScript) {
+        getScriptDefinition(virtualFile, project)?.let {
+            onOccurrence(ScriptModuleInfo(project, virtualFile, it))
+        }
+    }
+
     projectFileIndex.getOrderEntriesForFile(virtualFile).forEach {
         it.toIdeaModuleInfo(project, virtualFile, treatAsLibrarySource)?.let(onOccurrence)
     }
diff --git a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/resolve/KotlinCacheServiceImpl.kt b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/resolve/KotlinCacheServiceImpl.kt
index 6dfa1a6a26f..7ddab78704d 100644
--- a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/resolve/KotlinCacheServiceImpl.kt
+++ b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/resolve/KotlinCacheServiceImpl.kt
@@ -391,13 +391,13 @@ class KotlinCacheServiceImpl(val project: Project) : KotlinCacheService {
     private fun getFacadeToAnalyzeFiles(files: Collection<KtFile>): ResolutionFacade {
         val file = files.first()
         val moduleInfo = file.getModuleInfo()
-        val scripts = files.filterScripts()
+        val specialFiles = files.filterNotInProjectSource(moduleInfo)
+        val scripts = specialFiles.filterScripts()
         if (scripts.isNotEmpty()) {
             val projectFacade = getFacadeForScripts(scripts)
             return ModuleResolutionFacadeImpl(projectFacade, moduleInfo)
         }
 
-        val specialFiles = files.filterNotInProjectSource(moduleInfo)
         if (specialFiles.isNotEmpty()) {
             val projectFacade = getFacadeForSpecialFiles(specialFiles)
             return ModuleResolutionFacadeImpl(projectFacade, moduleInfo)
diff --git a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/vfilefinder/KotlinPackageSourcesMemberNamesIndex.kt b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/vfilefinder/KotlinPackageSourcesMemberNamesIndex.kt
index 30ae31432df..39468b3ccc6 100644
--- a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/vfilefinder/KotlinPackageSourcesMemberNamesIndex.kt
+++ b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/vfilefinder/KotlinPackageSourcesMemberNamesIndex.kt
@@ -9,6 +9,7 @@ import com.intellij.util.indexing.*
 import com.intellij.util.io.DataExternalizer
 import com.intellij.util.io.EnumeratorStringDescriptor
 import org.jetbrains.kotlin.idea.KotlinFileType
+import org.jetbrains.kotlin.parsing.KotlinParserDefinition
 import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtFile
 import java.io.DataInput
@@ -28,7 +29,7 @@ object KotlinPackageSourcesMemberNamesIndex : FileBasedIndexExtension<String, Co
     override fun getValueExternalizer() = StringSetExternalizer
 
     override fun getInputFilter(): FileBasedIndex.InputFilter =
-            FileBasedIndex.InputFilter { file -> file.extension == KotlinFileType.EXTENSION }
+        FileBasedIndex.InputFilter { file -> file.extension == KotlinFileType.EXTENSION || file.extension == KotlinParserDefinition.STD_SCRIPT_SUFFIX }
 
     override fun getVersion(): Int = 2
 
@@ -40,7 +41,11 @@ object KotlinPackageSourcesMemberNamesIndex : FileBasedIndexExtension<String, Co
             val ktFile = inputData.psiFile as? KtFile ?: return@DataIndexer emptyMap()
             val packageName = ktFile.packageDirective?.fqName?.asString() ?: ""
 
-            mapOf(packageName to ktFile.declarations.mapNotNullTo(hashSetOf(), KtDeclaration::getName))
+            if (!ktFile.isScript()) {
+                mapOf(packageName to ktFile.declarations.mapNotNullTo(hashSetOf(), KtDeclaration::getName))
+            } else {
+                mapOf(packageName to listOfNotNull(ktFile.script?.name).toHashSet())
+            }
         }
 }