From 320c5f6f3afdaf346d8f33fe3d614e1cd87851aa Mon Sep 17 00:00:00 2001
From: "Pavel V. Talanov" <talanov.pavel@gmail.com>
Date: Thu, 12 Oct 2017 19:06:03 +0300
Subject: [PATCH] Try to approximate "build.gradle.kts" script dependencies on
 first access

This is a hacky solution to prevent whole file being highlighted as unresolved while we wait for gradle response
---
 .../idea/caches/resolve/ScriptModuleInfos.kt  |  2 +-
 .../core/script/ScriptDependenciesCache.kt    |  2 +
 .../core/script/ScriptDependenciesManager.kt  |  2 +-
 .../core/script/ScriptDependenciesUpdater.kt  | 44 ++++++++++++++-----
 .../script/GradleScriptTemplateProvider.kt    | 38 +++++++++++++++-
 idea/src/META-INF/extensions/common.xml       |  3 ++
 idea/src/META-INF/gradle.xml                  |  1 +
 7 files changed, 79 insertions(+), 13 deletions(-)

diff --git a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/resolve/ScriptModuleInfos.kt b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/resolve/ScriptModuleInfos.kt
index 5c2cd6bccaf..08e013c0693 100644
--- a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/resolve/ScriptModuleInfos.kt
+++ b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/resolve/ScriptModuleInfos.kt
@@ -45,7 +45,7 @@ data class ScriptModuleInfo(
     override val moduleOrigin: ModuleOrigin
         get() = ModuleOrigin.OTHER
 
-    val externalDependencies: ScriptDependencies?
+    val externalDependencies: ScriptDependencies
         get() = ScriptDependenciesManager.getInstance(project).getScriptDependencies(scriptFile)
 
     override val name: Name = Name.special("<script ${scriptFile.name} ${scriptDefinition.name}>")
diff --git a/idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/ScriptDependenciesCache.kt b/idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/ScriptDependenciesCache.kt
index 7937fd895d5..e06a4660865 100644
--- a/idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/ScriptDependenciesCache.kt
+++ b/idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/ScriptDependenciesCache.kt
@@ -58,6 +58,8 @@ class ScriptDependenciesCache(private val project: Project) {
         NonClasspathDirectoriesScope(allLibrarySources)
     }
 
+    fun <T> inspectCache(body: (Map<String, ScriptDependencies>) -> T): T = cacheLock.read { body(cache) }
+
     private fun onChange(file: VirtualFile?) {
         this::allScriptsClasspath.clearValue()
         this::allScriptsClasspathScope.clearValue()
diff --git a/idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/ScriptDependenciesManager.kt b/idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/ScriptDependenciesManager.kt
index 70787534888..4a8f7a046c8 100644
--- a/idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/ScriptDependenciesManager.kt
+++ b/idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/ScriptDependenciesManager.kt
@@ -51,7 +51,7 @@ class ScriptDependenciesManager internal constructor(
     }
 
     fun getScriptClasspath(file: VirtualFile): List<VirtualFile> = toVfsRoots(cacheUpdater.getCurrentDependencies(file).classpath)
-    fun getScriptDependencies(file: VirtualFile) = cacheUpdater.getCurrentDependencies(file)
+    fun getScriptDependencies(file: VirtualFile): ScriptDependencies = cacheUpdater.getCurrentDependencies(file)
 
     private fun reloadScriptDefinitions() {
         val def = makeScriptDefsFromTemplatesProviderExtensions(
diff --git a/idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/ScriptDependenciesUpdater.kt b/idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/ScriptDependenciesUpdater.kt
index 583cc0a72c2..47559619e70 100644
--- a/idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/ScriptDependenciesUpdater.kt
+++ b/idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/ScriptDependenciesUpdater.kt
@@ -18,7 +18,6 @@ package org.jetbrains.kotlin.idea.core.script
 
 import com.intellij.openapi.application.Application
 import com.intellij.openapi.application.ApplicationManager
-import com.intellij.openapi.application.ModalityState
 import com.intellij.openapi.components.ServiceManager
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.project.isProjectOrWorkspaceFile
@@ -35,11 +34,13 @@ import kotlinx.coroutines.experimental.Job
 import kotlinx.coroutines.experimental.asCoroutineDispatcher
 import kotlinx.coroutines.experimental.launch
 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.application.runWriteAction
 import org.jetbrains.kotlin.psi.NotNullableUserDataProperty
 import org.jetbrains.kotlin.script.*
+import org.jetbrains.kotlin.utils.addToStdlib.firstNotNullResult
 import java.util.concurrent.ConcurrentHashMap
 import java.util.concurrent.Executors
 import kotlin.script.experimental.dependencies.AsyncDependenciesResolver
@@ -77,20 +78,33 @@ internal class ScriptDependenciesUpdater(
     fun getCurrentDependencies(file: VirtualFile): ScriptDependencies {
         cache[file]?.let { return it }
 
-        tryLoadingFromDisk(file)
+        val loaded = tryLoadingFromDisk(file)
+
+        if (!loaded) {
+            tryUsingDefault(file)
+        }
 
         updateCache(listOf(file))
 
         return cache[file] ?: ScriptDependencies.Empty
     }
 
-    private fun tryLoadingFromDisk(file: VirtualFile) {
-        ScriptDependenciesFileAttribute.read(file)?.let { deserialized ->
-            val rootsChanged = cache.hasNotCachedRoots(deserialized)
-            cache.save(file, deserialized)
-            if (rootsChanged) {
-                notifyRootsChanged()
-            }
+    private fun tryUsingDefault(file: VirtualFile) {
+        val defaults = DefaultScriptDependenciesProvider.getInstances(project).firstNotNullResult { it.defaultDependenciesFor(file) } ?: return
+        saveToCache(defaults, file)
+    }
+
+    private fun tryLoadingFromDisk(file: VirtualFile): Boolean {
+        val deserializedDependencies = ScriptDependenciesFileAttribute.read(file) ?: return false
+        saveToCache(deserializedDependencies, file)
+        return true
+    }
+
+    private fun saveToCache(deserialized: ScriptDependencies, file: VirtualFile) {
+        val rootsChanged = cache.hasNotCachedRoots(deserialized)
+        cache.save(file, deserialized)
+        if (rootsChanged) {
+            notifyRootsChanged()
         }
     }
 
@@ -280,4 +294,14 @@ private object TimeStamps {
 }
 
 @set: TestOnly
-var Application.isScriptDependenciesUpdaterDisabled by NotNullableUserDataProperty(Key.create("SCRIPT_DEPENDENCIES_UPDATER_DISABLED"), false)
\ No newline at end of file
+var Application.isScriptDependenciesUpdaterDisabled by NotNullableUserDataProperty(Key.create("SCRIPT_DEPENDENCIES_UPDATER_DISABLED"), false)
+
+interface DefaultScriptDependenciesProvider {
+    fun defaultDependenciesFor(scriptFile: VirtualFile): ScriptDependencies?
+
+    companion object : ProjectExtensionDescriptor<DefaultScriptDependenciesProvider>(
+            "org.jetbrains.kotlin.defaultScriptDependenciesProvider",
+            DefaultScriptDependenciesProvider::class.java
+    )
+
+}
\ No newline at end of file
diff --git a/idea/idea-gradle/src/org/jetbrains/kotlin/idea/core/script/GradleScriptTemplateProvider.kt b/idea/idea-gradle/src/org/jetbrains/kotlin/idea/core/script/GradleScriptTemplateProvider.kt
index d76fa09b618..1ca71fb8047 100644
--- a/idea/idea-gradle/src/org/jetbrains/kotlin/idea/core/script/GradleScriptTemplateProvider.kt
+++ b/idea/idea-gradle/src/org/jetbrains/kotlin/idea/core/script/GradleScriptTemplateProvider.kt
@@ -20,7 +20,9 @@ import com.intellij.execution.configurations.CommandLineTokenizer
 import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.externalSystem.util.ExternalSystemApiUtil
 import com.intellij.openapi.project.Project
+import com.intellij.openapi.vfs.VirtualFile
 import org.gradle.tooling.ProjectConnection
+import org.jetbrains.kotlin.gradle.kdsl.DEFAULT_SCRIPT_NAME
 import org.jetbrains.kotlin.lexer.KotlinLexer
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.script.ScriptTemplatesProvider
@@ -28,6 +30,7 @@ import org.jetbrains.plugins.gradle.service.execution.GradleExecutionHelper
 import org.jetbrains.plugins.gradle.settings.GradleExecutionSettings
 import java.io.File
 import java.util.*
+import kotlin.script.experimental.dependencies.ScriptDependencies
 
 abstract class AbstractGradleScriptTemplatesProvider(
         project: Project, override val id: String, private val templateClass: String, private val dependencySelector: Regex
@@ -164,4 +167,37 @@ class TopLevelSectionTokensEnumerator(script: CharSequence, identifier: String)
 fun topLevelSectionCodeTextTokens(script: CharSequence, sectionIdentifier: String): Sequence<CharSequence> =
         TopLevelSectionTokensEnumerator(script, sectionIdentifier).asSequence()
                 .filter { it.tokenType !in KtTokens.WHITE_SPACE_OR_COMMENT_BIT_SET }
-                .map { it.tokenSequence }
\ No newline at end of file
+                .map { it.tokenSequence }
+
+
+private const val KOTLIN_BUILD_FILE_NAME = DEFAULT_SCRIPT_NAME
+
+class GradleScriptDefaultDependenciesProvider(
+        private val scriptDependenciesCache: ScriptDependenciesCache
+) : DefaultScriptDependenciesProvider {
+    override fun defaultDependenciesFor(scriptFile: VirtualFile): ScriptDependencies? {
+        if (scriptFile.name != KOTLIN_BUILD_FILE_NAME) return null
+
+        return previouslyAnalyzedScriptsCombinedDependencies().takeUnless { it.classpath.isEmpty() }
+    }
+
+    private fun previouslyAnalyzedScriptsCombinedDependencies(): ScriptDependencies {
+        val sources = mutableListOf<File>()
+        val binaries = mutableListOf<File>()
+        val imports = mutableListOf<String>()
+
+        scriptDependenciesCache.inspectCache { cache ->
+            cache.entries.filter { it.key.endsWith(KOTLIN_BUILD_FILE_NAME) }
+                    .map { it.value }.forEach {
+                sources += it.sources
+                binaries += it.classpath
+                imports += it.imports
+            }
+        }
+        return ScriptDependencies(
+                classpath = binaries.distinct(),
+                sources = sources.distinct(),
+                imports = imports.distinct()
+        )
+    }
+}
diff --git a/idea/src/META-INF/extensions/common.xml b/idea/src/META-INF/extensions/common.xml
index f90a599ff4f..eea74ce50e6 100644
--- a/idea/src/META-INF/extensions/common.xml
+++ b/idea/src/META-INF/extensions/common.xml
@@ -51,5 +51,8 @@
                         interface="com.intellij.ide.util.projectWizard.ModuleBuilder"/>
         <extensionPoint name="gradleFrameworkSupport"
                         interface="org.jetbrains.kotlin.gradle.kdsl.frameworkSupport.GradleFrameworkSupportProvider"/>
+        <extensionPoint name="defaultScriptDependenciesProvider"
+                        interface="org.jetbrains.kotlin.idea.core.script.DefaultScriptDependenciesProvider"
+                        area="IDEA_PROJECT"/>
     </extensionPoints>
 </idea-plugin>
diff --git a/idea/src/META-INF/gradle.xml b/idea/src/META-INF/gradle.xml
index 9e4672a0988..5d852f55359 100644
--- a/idea/src/META-INF/gradle.xml
+++ b/idea/src/META-INF/gradle.xml
@@ -49,6 +49,7 @@
         <gradleProjectImportHandler implementation="org.jetbrains.kotlin.allopen.ide.AllOpenGradleProjectImportHandler"/>
         <gradleProjectImportHandler implementation="org.jetbrains.kotlin.noarg.ide.NoArgGradleProjectImportHandler"/>
         <gradleProjectImportHandler implementation="org.jetbrains.kotlin.samWithReceiver.ide.SamWithReceiverGradleProjectImportHandler"/>
+        <defaultScriptDependenciesProvider implementation="org.jetbrains.kotlin.idea.core.script.GradleScriptDefaultDependenciesProvider"/>
 
         <projectConfigurator implementation="org.jetbrains.kotlin.idea.configuration.KotlinGradleModuleConfigurator"/>
         <projectConfigurator implementation="org.jetbrains.kotlin.idea.configuration.KotlinJsGradleModuleConfigurator"/>