Do not highlight kts files outside source roots if not all definitions are initialized

^KT-20521 Fixed
^KT-21420 Fixed

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

@@ -19,6 +19,7 @@ package org.jetbrains.kotlin.idea.highlighter
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.idea.caches.resolve.NotUnderContentRootModuleInfo
 import org.jetbrains.kotlin.idea.caches.resolve.getModuleInfo
+import org.jetbrains.kotlin.idea.core.script.scriptDependencies
 import org.jetbrains.kotlin.idea.util.ProjectRootsUtil
 import org.jetbrains.kotlin.psi.KtCodeFragment
 import org.jetbrains.kotlin.psi.KtFile
@@ -26,9 +27,16 @@ import org.jetbrains.kotlin.psi.KtFile
 object KotlinHighlightingUtil {
     fun shouldHighlight(psiElement: PsiElement): Boolean {
         val ktFile = psiElement.containingFile as? KtFile ?: return false
-        return ktFile is KtCodeFragment && ktFile.context != null ||
+
-               ktFile.isScript() ||
+        if (ktFile is KtCodeFragment && ktFile.context != null) {
-               ProjectRootsUtil.isInProjectOrLibraryContent(ktFile) && ktFile.getModuleInfo() !is NotUnderContentRootModuleInfo
+            return true
+        }
+
+        if (ktFile.isScript()) {
+            return ktFile.virtualFile.scriptDependencies != null
+        }
+
+        return ProjectRootsUtil.isInProjectOrLibraryContent(ktFile) && ktFile.getModuleInfo() !is NotUnderContentRootModuleInfo
     }
 
     fun shouldHighlightErrors(psiElement: PsiElement): Boolean {
@@ -36,6 +44,14 @@ object KotlinHighlightingUtil {
         if (ktFile.isCompiled) {
             return false
         }
-        return (ktFile is KtCodeFragment && ktFile.context != null) || ktFile.isScript() || ProjectRootsUtil.isInProjectSource(ktFile)
+        if (ktFile is KtCodeFragment && ktFile.context != null) {
+            return true
+        }
+
+        if (ktFile.isScript()) {
+            return ktFile.virtualFile.scriptDependencies != null
+        }
+
+        return ProjectRootsUtil.isInProjectSource(ktFile)
     }
 }

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

@@ -46,6 +46,9 @@ class ScriptDefinitionsManager(private val project: Project): ScriptDefinitionPr
     private var definitionsByContributor = mutableMapOf<ScriptDefinitionContributor, List<KotlinScriptDefinition>>()
     private var definitions: List<KotlinScriptDefinition> = emptyList()
 
+    var hasFailedDefinitions = false
+        private set
+
     fun reloadDefinitionsBy(contributor: ScriptDefinitionContributor) = lock.write {
         val notLoadedYet = definitions.isEmpty()
         if (notLoadedYet) return
@@ -53,6 +56,9 @@ class ScriptDefinitionsManager(private val project: Project): ScriptDefinitionPr
         if (contributor !in definitionsByContributor) error("Unknown contributor: ${contributor.id}")
 
         definitionsByContributor[contributor] = contributor.safeGetDefinitions()
+
+        hasFailedDefinitions = getContributors().any { it.isError() }
+
         updateDefinitions()
     }
 
@@ -84,7 +90,13 @@ class ScriptDefinitionsManager(private val project: Project): ScriptDefinitionPr
     }
 
     fun reloadScriptDefinitions() = lock.write {
-        definitionsByContributor = getContributors().associateByTo(mutableMapOf(), { it }, { it.safeGetDefinitions() })
+        for (contributor in getContributors()) {
+            val definitions = contributor.safeGetDefinitions()
+            definitionsByContributor[contributor] = definitions
+        }
+
+        hasFailedDefinitions = getContributors().any { it.isError() }
+
         updateDefinitions()
     }
 
@@ -148,6 +160,8 @@ interface ScriptDefinitionContributor {
 
     fun getDefinitions(): List<KotlinScriptDefinition>
 
+    fun isError(): Boolean = false
+
     companion object {
         val EP_NAME: ExtensionPointName<ScriptDefinitionContributor> =
                 ExtensionPointName.create<ScriptDefinitionContributor>("org.jetbrains.kotlin.scriptDefinitionContributor")

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

@@ -16,72 +16,58 @@
 
 package org.jetbrains.kotlin.idea.core.script
 
-import com.intellij.openapi.components.ServiceManager
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.openapi.vfs.VirtualFileWithId
+import com.intellij.openapi.vfs.newvfs.FileAttribute
 import com.intellij.util.io.DataInputOutputUtil.readSeq
 import com.intellij.util.io.DataInputOutputUtil.writeSeq
 import com.intellij.util.io.IOUtil.readUTF
 import com.intellij.util.io.IOUtil.writeUTF
-import org.jetbrains.kotlin.idea.caches.FileAttributeService
 import java.io.DataInput
 import java.io.DataOutput
 import java.io.File
+import kotlin.reflect.KProperty
 import kotlin.script.experimental.dependencies.ScriptDependencies
 
-object ScriptDependenciesFileAttribute {
+var VirtualFile.scriptDependencies: ScriptDependencies? by ScriptDependenciesProperty()
-    private val VERSION = 1
+private val scriptDependencies = FileAttribute("kotlin-script-dependencies", 2, false)
-    private val ID = "kotlin-script-dependencies"
 
-    private val fileAttributeService = ServiceManager.getService(FileAttributeService::class.java)
+private class ScriptDependenciesProperty {
 
-    init {
+    operator fun setValue(file: VirtualFile, property: KProperty<*>, newValue: ScriptDependencies?) {
-        fileAttributeService.register(ID, VERSION, false)
+        if (file !is VirtualFileWithId) return
-    }
 
-    fun write(virtualFile: VirtualFile, dependencies: ScriptDependencies) {
+        if (newValue != null) {
-        if (virtualFile !is VirtualFileWithId) return
+            val output = scriptDependencies.writeAttribute(file)
-        fileAttributeService.write(virtualFile, ID, dependencies) { output, dep ->
+            output.use {
-            with(dep) {
+                with(newValue) {
-                output.writeInt(VERSION)
+                    with(output) {
+                        writeFileList(classpath)
+                        writeStringList(imports)
+                        writeNullable(javaHome, DataOutput::writeFile)
+                        writeFileList(scripts)
+                        writeFileList(sources)
 
-                writeDependencies(this, output)
+                    }
+                }
             }
         }
     }
 
-    fun read(virtualFile: VirtualFile): ScriptDependencies? {
+    operator fun getValue(file: VirtualFile, property: KProperty<*>): ScriptDependencies? {
-        if (virtualFile !is VirtualFileWithId) return null
+        if (file !is VirtualFileWithId) return null
 
-        return fileAttributeService.read(virtualFile, ID) { input ->
+        val input = scriptDependencies.readAttribute(file)
-            val version = input.readInt()
+        return input?.use {
-            if (version != VERSION) null
+            with(input) {
-            else readDependencies(input)
+                return ScriptDependencies(
-        }?.value
-    }
-
-    private fun writeDependencies(scriptDependencies: ScriptDependencies, output: DataOutput) {
-        with(scriptDependencies) {
-            with(output) {
-                writeFileList(classpath)
-                writeStringList(imports)
-                writeNullable(javaHome, DataOutput::writeFile)
-                writeFileList(scripts)
-                writeFileList(sources)
-
-            }
-        }
-    }
-
-    private fun readDependencies(input: DataInput): ScriptDependencies {
-        with(input) {
-            return ScriptDependencies(
                     classpath = readFileList(),
                     imports = readStringList(),
                     javaHome = readNullable(DataInput::readFile),
                     scripts = readFileList(),
                     sources = readFileList()
-            )
+                )
+            }
         }
     }
 }

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

@@ -37,6 +37,7 @@ import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.extensions.ProjectExtensionDescriptor
 import org.jetbrains.kotlin.idea.core.util.EDT
 import org.jetbrains.kotlin.idea.core.util.cancelOnDisposal
+import org.jetbrains.kotlin.idea.util.ProjectRootsUtil
 import org.jetbrains.kotlin.idea.util.application.runWriteAction
 import org.jetbrains.kotlin.psi.NotNullableUserDataProperty
 import org.jetbrains.kotlin.script.*
@@ -93,7 +94,7 @@ class ScriptDependenciesUpdater(
     }
 
     private fun tryLoadingFromDisk(file: VirtualFile): Boolean {
-        val deserializedDependencies = ScriptDependenciesFileAttribute.read(file) ?: return false
+        val deserializedDependencies = file.scriptDependencies ?: return false
         saveToCache(deserializedDependencies, file)
         return true
     }
@@ -117,6 +118,10 @@ class ScriptDependenciesUpdater(
     }
 
     private fun performUpdate(file: VirtualFile) {
+        if (ScriptDefinitionsManager.getInstance(project).hasFailedDefinitions && !ProjectRootsUtil.isProjectSourceFile(project, file)) {
+            return
+        }
+
         val scriptDef = scriptDefinitionProvider.findScriptDefinition(file) ?: return
         when (scriptDef.dependencyResolver) {
             is AsyncDependenciesResolver, is LegacyResolverWrapper -> {
@@ -221,7 +226,7 @@ class ScriptDependenciesUpdater(
     ): Boolean {
         val rootsChanged = cache.hasNotCachedRoots(new)
         if (cache.save(file, new)) {
-            ScriptDependenciesFileAttribute.write(file, new)
+            file.scriptDependencies = new
         }
         return rootsChanged
     }

idea/idea-gradle/src/org/jetbrains/kotlin/idea/core/script/GradleScriptTemplateProvider.kt

@@ -18,7 +18,6 @@ package org.jetbrains.kotlin.idea.core.script
 
 import com.intellij.execution.configurations.CommandLineTokenizer
 import com.intellij.openapi.components.ServiceManager
-import com.intellij.openapi.components.service
 import com.intellij.openapi.externalSystem.model.task.ExternalSystemTaskId
 import com.intellij.openapi.externalSystem.model.task.ExternalSystemTaskNotificationListenerAdapter
 import com.intellij.openapi.externalSystem.model.task.ExternalSystemTaskType
@@ -34,6 +33,7 @@ import org.jetbrains.plugins.gradle.config.GradleSettingsListenerAdapter
 import org.jetbrains.plugins.gradle.service.execution.GradleExecutionHelper
 import org.jetbrains.plugins.gradle.settings.DistributionType
 import org.jetbrains.plugins.gradle.settings.GradleExecutionSettings
+import org.jetbrains.plugins.gradle.settings.GradleProjectSettings
 import org.jetbrains.plugins.gradle.settings.GradleSettingsListener
 import org.jetbrains.plugins.gradle.util.GradleConstants
 import java.io.File
@@ -74,17 +74,19 @@ class GradleScriptDefinitionsContributor(private val project: Project): ScriptDe
     }
 
     override fun getDefinitions(): List<KotlinScriptDefinition> {
-        val definitions = loadDefinitions()
+        return loadDefinitions()
-        failedToLoad.set(definitions.isEmpty())
-        return definitions
     }
 
+    override fun isError() = failedToLoad.get()
+
     // NOTE: control flow here depends on suppressing exceptions from loadGradleTemplates calls
     // TODO: possibly combine exceptions from every loadGradleTemplates call, be mindful of KT-19276
     private fun loadDefinitions(): List<KotlinScriptDefinition> {
         val kotlinDslDependencySelector = Regex("^gradle-(?:kotlin-dsl|core).*\\.jar\$")
         val kotlinDslAdditionalResolverCp = ::kotlinStdlibAndCompiler
 
+        failedToLoad.set(false)
+
         // KotlinBuildScript should be last because it has wide scriptFilePattern
         val kotlinDslTemplates = loadGradleTemplates(
             templateClass = "org.gradle.kotlin.dsl.KotlinInitScript",
@@ -104,12 +106,18 @@ class GradleScriptDefinitionsContributor(private val project: Project): ScriptDe
         if (kotlinDslTemplates.isNotEmpty()) {
             return kotlinDslTemplates
         }
-        val gradleScriptKotlinLegacyTemplates = loadGradleTemplates(
+
-                templateClass = "org.gradle.script.lang.kotlin.KotlinBuildScript",
+        return tryToLoadOldBuildScriptDefinition()
-                dependencySelector = Regex("^gradle-(?:script-kotlin|core).*\\.jar\$"),
+    }
-                additionalResolverClasspath = { emptyList() }
+
+    private fun tryToLoadOldBuildScriptDefinition(): List<KotlinScriptDefinition> {
+        failedToLoad.set(false)
+
+        return loadGradleTemplates(
+            templateClass = "org.gradle.script.lang.kotlin.KotlinBuildScript",
+            dependencySelector = Regex("^gradle-(?:script-kotlin|core).*\\.jar\$"),
+            additionalResolverClasspath = { emptyList() }
         )
-        return gradleScriptKotlinLegacyTemplates
     }
 
     // TODO: check this against kotlin-dsl branch that uses daemon
@@ -127,6 +135,7 @@ class GradleScriptDefinitionsContributor(private val project: Project): ScriptDe
     }
     catch (t: Throwable) {
         // TODO: review exception handling
+        failedToLoad.set(true)
         emptyList()
     }
 
@@ -157,9 +166,14 @@ class GradleScriptDefinitionsContributor(private val project: Project): ScriptDe
 
         }
 
+        val gradleSettings = ExternalSystemApiUtil.getSettings(project, GradleConstants.SYSTEM_ID)
+        if (gradleSettings.getLinkedProjectsSettings().isEmpty()) return emptyList()
+
+        val projectSettings = gradleSettings.getLinkedProjectsSettings().filterIsInstance<GradleProjectSettings>().firstOrNull() ?: return emptyList()
+
         val gradleExeSettings = ExternalSystemApiUtil.getExecutionSettings<GradleExecutionSettings>(
                 project,
-                ExternalSystemApiUtil.toCanonicalPath(project.basePath!!),
+                projectSettings.externalProjectPath,
                 GradleConstants.SYSTEM_ID)
 
         val gradleHome = gradleExeSettings.gradleHome ?: error("Unable to get Gradle home directory")
@@ -200,7 +214,7 @@ class ReloadGradleTemplatesOnSync : ExternalSystemTaskNotificationListenerAdapte
             val gradleDefinitionsContributor = ScriptDefinitionContributor.find<GradleScriptDefinitionsContributor>(project)
             gradleDefinitionsContributor?.reloadIfNeccessary()
 
-            project.service<ScriptDependenciesUpdater>().reloadModifiedScripts()
+            ServiceManager.getService(project, ScriptDependenciesUpdater::class.java).reloadModifiedScripts()
         }
     }
 }