From 37a95765cb5ccce29c2223c35b775134cc5070b6 Mon Sep 17 00:00:00 2001
From: Natalia Selezneva <natalia.selezneva@jetbrains.com>
Date: Mon, 9 Apr 2018 09:57:35 +0300
Subject: [PATCH] Add scriptDependencies moduleInfo for scripts located in
 module source

---
 .../idea/caches/project/ScriptModuleInfos.kt  |   9 +-
 .../idea/caches/project/getModuleInfo.kt      |  80 ++++++++----
 .../caches/resolve/KotlinCacheServiceImpl.kt  | 116 ++++++++++--------
 3 files changed, 123 insertions(+), 82 deletions(-)

diff --git a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/project/ScriptModuleInfos.kt b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/project/ScriptModuleInfos.kt
index 61ea58c83d0..94d20ce1295 100644
--- a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/project/ScriptModuleInfos.kt
+++ b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/project/ScriptModuleInfos.kt
@@ -37,14 +37,17 @@ data class ScriptModuleInfo(
     val externalDependencies: ScriptDependencies
         get() = ScriptDependenciesManager.getInstance(project).getScriptDependencies(scriptFile)
 
+    val relatedModuleSourceInfo: ModuleSourceInfo? = getScriptRelatedModuleInfo(project, scriptFile)
+
     override val name: Name = Name.special("<script ${scriptFile.name} ${scriptDefinition.name}>")
 
     override fun contentScope() = GlobalSearchScope.fileScope(project, scriptFile)
 
     override fun dependencies(): List<IdeaModuleInfo> {
-        return listOf(
-            this, ScriptDependenciesModuleInfo(project, this)
-        ) + sdkDependencies(externalDependencies, project)
+        val result = arrayListOf(this, ScriptDependenciesModuleInfo(project, this))
+        relatedModuleSourceInfo?.let { result.add(it) }
+        result.addAll(sdkDependencies(externalDependencies, project))
+        return result
     }
 }
 
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 cc1c111a491..13e5a97f831 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
@@ -90,11 +90,27 @@ fun collectAllModuleInfosFromIdeaModel(project: Project): List<IdeaModuleInfo> {
     return modulesSourcesInfos + librariesInfos + sdksInfos
 }
 
+fun getScriptRelatedModuleInfo(project: Project, virtualFile: VirtualFile): ModuleSourceInfo? {
+    val projectFileIndex = ProjectFileIndex.SERVICE.getInstance(project)
+
+    val moduleRelatedModuleInfo = getModuleRelatedModuleInfo(projectFileIndex, virtualFile)
+    if (moduleRelatedModuleInfo != null) {
+        return moduleRelatedModuleInfo
+    }
+
+    if (ScratchFileService.isInScratchRoot(virtualFile)) {
+        val scratchModule = virtualFile.scriptRelatedModuleName?.let { ModuleManager.getInstance(project).findModuleByName(it) }
+        return scratchModule?.testSourceInfo() ?: scratchModule?.productionSourceInfo()
+    }
+
+    return null
+}
+
 internal fun getAllProjectSdks(): Collection<Sdk> {
     return ProjectJdkTable.getInstance().allJdks.toList()
 }
 
-private typealias VirtualFileProcessor<T> = (Project, VirtualFile, Boolean) -> T
+private typealias VirtualFileProcessor<T> = (Project, VirtualFile, Boolean, Boolean) -> T
 
 private sealed class ModuleInfoCollector<out T>(
     val onResult: (IdeaModuleInfo?) -> T,
@@ -107,11 +123,12 @@ private sealed class ModuleInfoCollector<out T>(
             LOG.error("Could not find correct module information.\nReason: $reason")
             NotUnderContentRootModuleInfo
         },
-        virtualFileProcessor = processor@ { project, virtualFile, isLibrarySource ->
+        virtualFileProcessor = processor@ { project, virtualFile, isLibrarySource, isScript ->
             collectInfosByVirtualFile(
                 project,
                 virtualFile,
                 isLibrarySource,
+                isScript,
                 {
                     return@processor it ?: NotUnderContentRootModuleInfo
                 })
@@ -124,11 +141,12 @@ private sealed class ModuleInfoCollector<out T>(
             LOG.warn("Could not find correct module information.\nReason: $reason")
             null
         },
-        virtualFileProcessor = processor@ { project, virtualFile, isLibrarySource ->
+        virtualFileProcessor = processor@ { project, virtualFile, isLibrarySource, isScript ->
             collectInfosByVirtualFile(
                 project,
                 virtualFile,
                 isLibrarySource,
+                isScript,
                 { return@processor it })
         }
     )
@@ -139,12 +157,13 @@ private sealed class ModuleInfoCollector<out T>(
             LOG.warn("Could not find correct module information.\nReason: $reason")
             emptySequence()
         },
-        virtualFileProcessor = { project, virtualFile, isLibrarySource ->
+        virtualFileProcessor = { project, virtualFile, isLibrarySource, isScript ->
             buildSequence {
                 collectInfosByVirtualFile(
                     project,
                     virtualFile,
                     isLibrarySource,
+                    isScript,
                     { yieldIfNotNull(it) })
             }
         }
@@ -189,7 +208,8 @@ private fun <T> PsiElement.collectInfos(c: ModuleInfoCollector<T>): T {
     return c.virtualFileProcessor(
         project,
         virtualFile,
-        (containingFile as? KtFile)?.isCompiled ?: false
+        (containingFile as? KtFile)?.isCompiled ?: false,
+        containingKtFile?.isScript() ?: false
     )
 }
 
@@ -199,7 +219,8 @@ private fun <T> KtLightElement<*, *>.processLightElement(c: ModuleInfoCollector<
         return c.virtualFileProcessor(
             project,
             containingFile.virtualFile.sure { "Decompiled class should be build from physical file" },
-            false
+            false,
+            (containingFile as? KtFile)?.isScript() ?: false
         )
     }
 
@@ -213,35 +234,29 @@ private fun <T> KtLightElement<*, *>.processLightElement(c: ModuleInfoCollector<
 }
 
 private inline fun <T> collectInfosByVirtualFile(
-    project: Project, virtualFile: VirtualFile,
-    treatAsLibrarySource: Boolean, onOccurrence: (IdeaModuleInfo?) -> T
+    project: Project,
+    virtualFile: VirtualFile,
+    treatAsLibrarySource: Boolean,
+    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)
 
-    val module = projectFileIndex.getModuleForFile(virtualFile)
-    if (module != null && !module.isDisposed) {
-        val moduleFileIndex = ModuleRootManager.getInstance(module).fileIndex
-        if (moduleFileIndex.isInTestSourceContent(virtualFile)) {
-            onOccurrence(module.testSourceInfo())
-        } else if (moduleFileIndex.isInSourceContentWithoutInjected(virtualFile)) {
-            onOccurrence(module.productionSourceInfo())
-        }
+    val moduleRelatedModuleInfo = getModuleRelatedModuleInfo(projectFileIndex, virtualFile)
+    if (moduleRelatedModuleInfo != null) {
+        onOccurrence(moduleRelatedModuleInfo)
     }
 
     projectFileIndex.getOrderEntriesForFile(virtualFile).forEach {
         it.toIdeaModuleInfo(project, virtualFile, treatAsLibrarySource)?.let(onOccurrence)
     }
 
-    if (ScratchFileService.isInScratchRoot(virtualFile)) {
-        val scratchModule = virtualFile.scriptRelatedModuleName?.let { ModuleManager.getInstance(project).findModuleByName(it) }
-        onOccurrence(scratchModule?.testSourceInfo() ?: scratchModule?.productionSourceInfo())
-    }
-
-    val scriptDefinition = getScriptDefinition(virtualFile, project)
-    if (scriptDefinition != null) {
-        onOccurrence(ScriptModuleInfo(project, virtualFile, scriptDefinition))
-    }
-
     val isBinary = virtualFile.fileType.isKotlinBinary()
     val scriptConfigurationManager = ScriptDependenciesManager.getInstance(project)
     if (isBinary && virtualFile in scriptConfigurationManager.getAllScriptsClasspathScope()) {
@@ -258,6 +273,19 @@ private inline fun <T> collectInfosByVirtualFile(
     return onOccurrence(null)
 }
 
+private fun getModuleRelatedModuleInfo(projectFileIndex: ProjectFileIndex, virtualFile: VirtualFile): ModuleSourceInfo? {
+    val module = projectFileIndex.getModuleForFile(virtualFile)
+    if (module != null && !module.isDisposed) {
+        val moduleFileIndex = ModuleRootManager.getInstance(module).fileIndex
+        if (moduleFileIndex.isInTestSourceContent(virtualFile)) {
+            return module.testSourceInfo()
+        } else if (moduleFileIndex.isInSourceContentWithoutInjected(virtualFile)) {
+            return module.productionSourceInfo()
+        }
+    }
+    return null
+}
+
 private inline fun <reified T : IdeaModuleInfo> collectModuleInfosByType(project: Project, virtualFile: VirtualFile): Collection<T> {
     val result = linkedSetOf<T>()
     collectInfosByVirtualFile(project, virtualFile, treatAsLibrarySource = false) {
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 2c23469f8b3..6b8027e0178 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
@@ -102,22 +102,32 @@ class KotlinCacheServiceImpl(val project: Project) : KotlinCacheService {
         dependenciesModuleInfo: ScriptDependenciesModuleInfo,
         syntheticFiles: Collection<KtFile> = listOf()
     ): ProjectResolutionFacade {
-        val sdk = findJdk(dependenciesModuleInfo.scriptModuleInfo?.externalDependencies, project)
+        val scriptModuleInfo = dependenciesModuleInfo.scriptModuleInfo
+        val sdk = findJdk(scriptModuleInfo?.externalDependencies, project)
         val platform = JvmPlatform // TODO: Js scripts?
         val settings = PlatformAnalysisSettings(platform, sdk, true)
-        val sdkFacade = GlobalFacade(settings).facadeForSdk
-        val globalContext = sdkFacade.globalContext.contextWithNewLockAndCompositeExceptionTracker()
+
+        val dependenciesForScriptDependencies = listOf(
+            LibraryModificationTracker.getInstance(project),
+            ProjectRootModificationTracker.getInstance(project),
+            ScriptDependenciesModificationTracker.getInstance(project)
+        )
+
+        val globalFacade = if (scriptModuleInfo?.relatedModuleSourceInfo != null) {
+            globalFacade(settings)
+        } else {
+            getOrBuildGlobalFacade(settings).facadeForSdk
+        }
+
+        val globalContext = globalFacade.globalContext.contextWithNewLockAndCompositeExceptionTracker()
         return ProjectResolutionFacade(
-            "facadeForScriptDependencies", "dependencies of scripts",
+            "facadeForScriptDependencies",
+            "dependencies of scripts",
             project, globalContext, settings,
-            reuseDataFrom = sdkFacade,
+            reuseDataFrom = globalFacade,
             allModules = dependenciesModuleInfo.dependencies(),
             //TODO: provide correct trackers
-            dependencies = listOf(
-                LibraryModificationTracker.getInstance(project),
-                ProjectRootModificationTracker.getInstance(project),
-                ScriptDependenciesModificationTracker.getInstance(project)
-            ),
+            dependencies = dependenciesForScriptDependencies,
             moduleFilter = { it == dependenciesModuleInfo },
             invalidateOnOOCB = true,
             syntheticFiles = syntheticFiles
@@ -184,12 +194,12 @@ class KotlinCacheServiceImpl(val project: Project) : KotlinCacheService {
 
     private val IdeaModuleInfo.sdk: Sdk? get() = dependencies().firstIsInstanceOrNull<SdkInfo>()?.sdk
 
-    private fun createFacadeForSyntheticFiles(files: Set<KtFile>): ProjectResolutionFacade {
+    private fun createFacadeForFilesWithSpecialModuleInfo(files: Set<KtFile>): ProjectResolutionFacade {
         // we assume that all files come from the same module
         val targetPlatform = files.map { TargetPlatformDetector.getPlatform(it) }.toSet().single()
-        val syntheticFileModule = files.map(KtFile::getModuleInfo).toSet().single()
-        val sdk = syntheticFileModule.sdk
-        val settings = PlatformAnalysisSettings(targetPlatform, sdk, syntheticFileModule.supportsAdditionalBuiltInsMembers())
+        val specialModuleInfo = files.map(KtFile::getModuleInfo).toSet().single()
+        val sdk = specialModuleInfo.sdk
+        val settings = PlatformAnalysisSettings(targetPlatform, sdk, specialModuleInfo.supportsAdditionalBuiltInsMembers())
         // File copies are created during completion and receive correct modification events through POM.
         // Dummy files created e.g. by J2K do not receive events.
         val filesModificationTracker = if (files.all { it.originalFile != it }) {
@@ -204,13 +214,13 @@ class KotlinCacheServiceImpl(val project: Project) : KotlinCacheService {
 
         val dependenciesForSyntheticFileCache =
             listOf(
-            PsiModificationTracker.OUT_OF_CODE_BLOCK_MODIFICATION_COUNT,
-            filesModificationTracker,
-            ScriptDependenciesModificationTracker.getInstance(project)
-        )
+                PsiModificationTracker.OUT_OF_CODE_BLOCK_MODIFICATION_COUNT,
+                filesModificationTracker,
+                ScriptDependenciesModificationTracker.getInstance(project)
+            )
 
         val resolverDebugName =
-            "completion/highlighting in $syntheticFileModule for files ${files.joinToString { it.name }} for platform $targetPlatform"
+            "completion/highlighting in $specialModuleInfo for files ${files.joinToString { it.name }} for platform $targetPlatform"
 
         fun makeProjectResolutionFacade(
             debugName: String,
@@ -235,69 +245,69 @@ class KotlinCacheServiceImpl(val project: Project) : KotlinCacheService {
         }
 
         return when {
-            syntheticFileModule is ModuleSourceInfo -> {
-                val dependentModules = syntheticFileModule.getDependentModules()
+            specialModuleInfo is ModuleSourceInfo -> {
+                val dependentModules = specialModuleInfo.getDependentModules()
                 val modulesFacade = globalFacade(settings)
                 val globalContext = modulesFacade.globalContext.contextWithNewLockAndCompositeExceptionTracker()
                 makeProjectResolutionFacade(
-                    "facadeForSynthetic in ModuleSourceInfo",
+                    "facadeForSpecialModuleInfo (ModuleSourceInfo)",
                     globalContext,
                     reuseDataFrom = modulesFacade,
                     moduleFilter = { it in dependentModules }
                 )
             }
 
-            syntheticFileModule is ScriptModuleInfo -> {
-                val facadeForScriptDependencies = getFacadeForScriptDependencies(syntheticFileModule)
+            specialModuleInfo is ScriptModuleInfo -> {
+                val facadeForScriptDependencies = getFacadeForScriptDependencies(specialModuleInfo)
                 val globalContext = facadeForScriptDependencies.globalContext.contextWithNewLockAndCompositeExceptionTracker()
                 makeProjectResolutionFacade(
-                    "facadeForSynthetic in ScriptModuleInfo",
+                    "facadeForSpecialModuleInfo (ScriptModuleInfo)",
                     globalContext,
                     reuseDataFrom = facadeForScriptDependencies,
-                    allModules = syntheticFileModule.dependencies(),
-                    moduleFilter = { it == syntheticFileModule }
+                    allModules = specialModuleInfo.dependencies(),
+                    moduleFilter = { it == specialModuleInfo }
                 )
             }
-            syntheticFileModule is ScriptDependenciesModuleInfo -> {
-                createFacadeForScriptDependencies(syntheticFileModule, files)
+            specialModuleInfo is ScriptDependenciesModuleInfo -> {
+                createFacadeForScriptDependencies(specialModuleInfo, files)
             }
-            syntheticFileModule is ScriptDependenciesSourceModuleInfo -> {
+            specialModuleInfo is ScriptDependenciesSourceModuleInfo -> {
                 // TODO: can be optimized by caching facadeForScriptDependencies
-                val facadeForScriptDependencies = createFacadeForScriptDependencies(syntheticFileModule.binariesModuleInfo, files)
+                val facadeForScriptDependencies = createFacadeForScriptDependencies(specialModuleInfo.binariesModuleInfo, files)
                 val globalContext = facadeForScriptDependencies.globalContext.contextWithNewLockAndCompositeExceptionTracker()
                 makeProjectResolutionFacade(
-                    "facadeForSynthetic in ScriptDependenciesSourceModuleInfo",
+                    "facadeForSpecialModuleInfo (ScriptDependenciesSourceModuleInfo)",
                     globalContext,
                     reuseDataFrom = facadeForScriptDependencies,
-                    allModules = syntheticFileModule.dependencies(),
-                    moduleFilter = { it == syntheticFileModule }
+                    allModules = specialModuleInfo.dependencies(),
+                    moduleFilter = { it == specialModuleInfo }
                 )
             }
 
-            syntheticFileModule is LibrarySourceInfo || syntheticFileModule is NotUnderContentRootModuleInfo -> {
+            specialModuleInfo is LibrarySourceInfo || specialModuleInfo is NotUnderContentRootModuleInfo -> {
                 val librariesFacade = librariesFacade(settings)
                 val globalContext = librariesFacade.globalContext.contextWithNewLockAndCompositeExceptionTracker()
                 makeProjectResolutionFacade(
-                    "facadeForSynthetic in LibrarySourceInfo or NotUnderContentRootModuleInfo",
+                    "facadeForSpecialModuleInfo (LibrarySourceInfo or NotUnderContentRootModuleInfo)",
                     globalContext,
                     reuseDataFrom = librariesFacade,
-                    moduleFilter = { it == syntheticFileModule }
+                    moduleFilter = { it == specialModuleInfo }
                 )
             }
 
-            syntheticFileModule.isLibraryClasses() -> {
+            specialModuleInfo.isLibraryClasses() -> {
                 //NOTE: this code should not be called for sdk or library classes
                 // currently the only known scenario is when we cannot determine that file is a library source
                 // (file under both classes and sources root)
-                LOG.warn("Creating cache with synthetic files ($files) in classes of library $syntheticFileModule")
+                LOG.warn("Creating cache with synthetic files ($files) in classes of library $specialModuleInfo")
                 val globalContext = GlobalContext()
                 makeProjectResolutionFacade(
-                    "facadeForSynthetic for file under both classes and root",
+                    "facadeForSpecialModuleInfo for file under both classes and root",
                     globalContext
                 )
             }
 
-            else -> throw IllegalStateException("Unknown IdeaModuleInfo ${syntheticFileModule::class.java}")
+            else -> throw IllegalStateException("Unknown IdeaModuleInfo ${specialModuleInfo::class.java}")
         }
     }
 
@@ -351,19 +361,19 @@ class KotlinCacheServiceImpl(val project: Project) : KotlinCacheService {
         false
     )
 
-    private val syntheticFileCachesLock = Any()
+    private val specialFileCachesLock = Any()
 
-    private val syntheticFilesCacheProvider = CachedValueProvider {
+    private val specialFilesCacheProvider = CachedValueProvider {
         CachedValueProvider.Result(object : SLRUCache<Set<KtFile>, ProjectResolutionFacade>(2, 3) {
-            override fun createValue(files: Set<KtFile>) = createFacadeForSyntheticFiles(files)
+            override fun createValue(files: Set<KtFile>) = createFacadeForFilesWithSpecialModuleInfo(files)
         }, LibraryModificationTracker.getInstance(project), ProjectRootModificationTracker.getInstance(project))
     }
 
-    private fun getFacadeForSyntheticFiles(files: Set<KtFile>): ProjectResolutionFacade {
-        val cachedValue = synchronized(syntheticFileCachesLock) {
-            //NOTE: computations inside createCacheForSyntheticFiles depend on project root structure
+    private fun getFacadeForSpecialFiles(files: Set<KtFile>): ProjectResolutionFacade {
+        val cachedValue = synchronized(specialFileCachesLock) {
+            //NOTE: computations inside createFacadeForFilesWithSpecialModuleInfo depend on project root structure
             // so we additionally drop the whole slru cache on change
-            CachedValuesManager.getManager(project).getCachedValue(project, syntheticFilesCacheProvider)
+            CachedValuesManager.getManager(project).getCachedValue(project, specialFilesCacheProvider)
         }
         // In Upsource, we create multiple instances of KotlinCacheService, which all access the same CachedValue instance (UP-8046)
         // To avoid race conditions, we can't use the local lock to access the cached value contents.
@@ -375,9 +385,9 @@ class KotlinCacheServiceImpl(val project: Project) : KotlinCacheService {
     private fun getFacadeToAnalyzeFiles(files: Collection<KtFile>): ResolutionFacade {
         val file = files.first()
         val moduleInfo = file.getModuleInfo()
-        val notInSourceFiles = files.filterNotInProjectSource(moduleInfo)
-        return if (notInSourceFiles.isNotEmpty()) {
-            val projectFacade = getFacadeForSyntheticFiles(notInSourceFiles)
+        val specialFiles = files.filterNotInProjectSourceOrScript(moduleInfo)
+        return if (specialFiles.isNotEmpty()) {
+            val projectFacade = getFacadeForSpecialFiles(specialFiles)
             ModuleResolutionFacadeImpl(projectFacade, moduleInfo)
         } else {
             val platform = TargetPlatformDetector.getPlatform(file)
@@ -401,11 +411,11 @@ class KotlinCacheServiceImpl(val project: Project) : KotlinCacheService {
     override fun getResolutionFacadeByModuleInfo(moduleInfo: ModuleInfo, platform: TargetPlatform): ResolutionFacade? =
         (moduleInfo as? IdeaModuleInfo)?.let { getResolutionFacadeByModuleInfo(it, platform) }
 
-    private fun Collection<KtFile>.filterNotInProjectSource(moduleInfo: IdeaModuleInfo): Set<KtFile> {
+    private fun Collection<KtFile>.filterNotInProjectSourceOrScript(moduleInfo: IdeaModuleInfo): Set<KtFile> {
         return mapNotNull {
             if (it is KtCodeFragment) it.getContextFile() else it
         }.filter {
-            !ProjectRootsUtil.isInProjectSource(it) || !moduleInfo.contentScope().contains(it)
+            !ProjectRootsUtil.isInProjectSource(it) || !moduleInfo.contentScope().contains(it) || it.isScript()
         }.toSet()
     }