Optimization: do not load dependencies for scripts not under expected location

idea/idea-analysis/src/org/jetbrains/kotlin/idea/highlighter/KotlinHighlightingUtil.kt

@@ -16,21 +16,17 @@
 
 package org.jetbrains.kotlin.idea.highlighter
 
-import com.intellij.ide.projectView.impl.ProjectRootsUtil.isInTestSource
-import com.intellij.ide.scratch.ScratchFileService
-import com.intellij.openapi.roots.ProjectRootManager
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.idea.caches.project.NotUnderContentRootModuleInfo
 import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
 import org.jetbrains.kotlin.idea.core.script.IdeScriptReportSink
+import org.jetbrains.kotlin.idea.core.script.ScriptDefinitionsManager
 import org.jetbrains.kotlin.idea.core.script.scriptDependencies
 import org.jetbrains.kotlin.idea.util.ProjectRootsUtil
-import org.jetbrains.kotlin.idea.util.ProjectRootsUtil.isInContent
 import org.jetbrains.kotlin.psi.KtCodeFragment
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.script.getScriptDefinition
 import kotlin.script.experimental.dependencies.ScriptReport
-import kotlin.script.experimental.location.ScriptExpectedLocation
 
 object KotlinHighlightingUtil {
     fun shouldHighlight(psiElement: PsiElement): Boolean {
@@ -66,25 +62,12 @@ object KotlinHighlightingUtil {
 
     @Suppress("DEPRECATION")
     private fun shouldHighlightScript(ktFile: KtFile): Boolean {
-        if (ScratchFileService.isInScratchRoot(ktFile.virtualFile)) return true
         if (ktFile.virtualFile.scriptDependencies == null) return false
-        if (ktFile.getUserData(IdeScriptReportSink.Reports)?.any { it.severity == ScriptReport.Severity.FATAL } == true) return false
-
-        val scriptScope = getScriptDefinition(ktFile)?.scriptExpectedLocations ?: return false
-        return when {
-            scriptScope.contains(ScriptExpectedLocation.Everywhere) -> true
-            scriptScope.contains(ScriptExpectedLocation.Project)
-                    && ProjectRootManager.getInstance(ktFile.project).fileIndex.isInContent(ktFile.virtualFile) -> true
-            scriptScope.contains(ScriptExpectedLocation.TestsOnly)
-                    && isInTestSource(ktFile) -> true
-            else -> return isInContent(
-                ktFile,
-                scriptScope.contains(ScriptExpectedLocation.SourcesOnly),
-                scriptScope.contains(ScriptExpectedLocation.Libraries),
-                scriptScope.contains(ScriptExpectedLocation.Libraries),
-                scriptScope.contains(ScriptExpectedLocation.Libraries)
-            )
+        if (ktFile.virtualFile.getUserData(IdeScriptReportSink.Reports)?.any { it.severity == ScriptReport.Severity.FATAL } == true) {
+            return false
         }
 
+        val scriptDefinition = getScriptDefinition(ktFile) ?: return false
+        return ScriptDefinitionsManager.getInstance(ktFile.project).isInExpectedLocation(ktFile, scriptDefinition)
     }
 }

idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/ScriptDefinitionsManager.kt

@@ -16,6 +16,8 @@
 
 package org.jetbrains.kotlin.idea.core.script
 
+import com.intellij.ide.projectView.impl.ProjectRootsUtil.isInTestSource
+import com.intellij.ide.scratch.ScratchFileService
 import com.intellij.openapi.components.ServiceManager
 import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.extensions.ExtensionPointName
@@ -31,6 +33,13 @@ import org.jetbrains.kotlin.idea.caches.project.SdkInfo
 import org.jetbrains.kotlin.idea.caches.project.getScriptRelatedModuleInfo
 import org.jetbrains.kotlin.script.*
 import org.jetbrains.kotlin.scripting.compiler.plugin.KotlinScriptDefinitionAdapterFromNewAPI
+import org.jetbrains.kotlin.script.*
+import org.jetbrains.kotlin.idea.util.ProjectRootsUtil.isInContent
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.script.KotlinScriptDefinition
+import org.jetbrains.kotlin.script.KotlinScriptDefinitionFromAnnotatedTemplate
+import org.jetbrains.kotlin.script.ScriptDefinitionProvider
+import org.jetbrains.kotlin.script.ScriptTemplatesProvider
 import org.jetbrains.kotlin.utils.PathUtil
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 import org.jetbrains.kotlin.utils.addToStdlib.flattenTo
@@ -46,6 +55,7 @@ import kotlin.script.experimental.definitions.ScriptDefinitionFromAnnotatedBaseC
 import kotlin.script.experimental.dependencies.DependenciesResolver
 import kotlin.script.experimental.dependencies.ScriptDependencies
 import kotlin.script.experimental.dependencies.asSuccess
+import kotlin.script.experimental.location.ScriptExpectedLocation
 import kotlin.script.templates.standard.ScriptTemplateWithArgs
 
 class ScriptDefinitionsManager(private val project: Project) : LazyScriptDefinitionProvider() {
@@ -94,6 +104,25 @@ class ScriptDefinitionsManager(private val project: Project) : LazyScriptDefinit
         updateDefinitions()
     }
 
+    fun isInExpectedLocation(ktFile: KtFile, scriptDefinition: KotlinScriptDefinition): Boolean {
+        if (ScratchFileService.isInScratchRoot(ktFile.virtualFile)) return true
+
+        val scriptScope = scriptDefinition.scriptExpectedLocations
+        return when {
+            scriptScope.contains(ScriptExpectedLocation.Everywhere) -> true
+            scriptScope.contains(ScriptExpectedLocation.Project)
+                    && ProjectRootManager.getInstance(ktFile.project).fileIndex.isInContent(ktFile.virtualFile) -> true
+            scriptScope.contains(ScriptExpectedLocation.TestsOnly) && isInTestSource(ktFile) -> true
+            else -> return isInContent(
+                ktFile,
+                scriptScope.contains(ScriptExpectedLocation.SourcesOnly),
+                scriptScope.contains(ScriptExpectedLocation.Libraries),
+                scriptScope.contains(ScriptExpectedLocation.Libraries),
+                scriptScope.contains(ScriptExpectedLocation.Libraries)
+            )
+        }
+    }
+
     private fun updateDefinitions() {
         assert(lock.isWriteLocked) { "updateDefinitions should only be called under the write lock" }
         definitions = definitionsByContributor.values.flattenTo(mutableListOf()).asSequence()

idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/ScriptDependenciesUpdater.kt

@@ -27,10 +27,12 @@ import com.intellij.openapi.fileEditor.FileEditorManagerListener
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.Key
 import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.psi.PsiManager
 import com.intellij.util.Alarm
 import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.idea.core.script.dependencies.FromFileAttributeScriptDependenciesLoader
 import org.jetbrains.kotlin.idea.core.script.dependencies.ScriptDependenciesLoader
+import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.NotNullableUserDataProperty
 import org.jetbrains.kotlin.script.ScriptDefinitionProvider
 import org.jetbrains.kotlin.script.findScriptDefinition
@@ -63,6 +65,9 @@ class ScriptDependenciesUpdater(
         project.messageBus.connect().subscribe(FileEditorManagerListener.FILE_EDITOR_MANAGER, object : FileEditorManagerListener {
             override fun fileOpened(source: FileEditorManager, file: VirtualFile) {
                 val scriptDef = scriptDefinitionProvider.findScriptDefinition(file) ?: return
+                val ktFile = PsiManager.getInstance(project).findFile(file) as? KtFile ?: return
+
+                if (!ScriptDefinitionsManager.getInstance(project).isInExpectedLocation(ktFile, scriptDef)) return
                 ScriptDependenciesLoader.updateDependencies(file, scriptDef, project, shouldNotifyRootsChanged = true)
             }
         })
@@ -86,6 +91,9 @@ class ScriptDependenciesUpdater(
                 }
 
                 val scriptDef = scriptDefinitionProvider.findScriptDefinition(file) ?: return
+                val ktFile = PsiManager.getInstance(project).findFile(file) as? KtFile ?: return
+
+                if (!ScriptDefinitionsManager.getInstance(project).isInExpectedLocation(ktFile, scriptDef)) return
 
                 scriptsQueue.cancelAllRequests()