From b1ac45195976ed2893fca058f7fb58e08a91c0ed Mon Sep 17 00:00:00 2001
From: "Pavel V. Talanov" <talanov.pavel@gmail.com>
Date: Fri, 14 Jul 2017 20:56:18 +0300
Subject: [PATCH] Make sure template classpath is always among script
 dependencies

---
 ...inScriptDefinitionFromAnnotatedTemplate.kt |  4 ++-
 .../kotlin/script/ScriptContentLoader.kt      | 10 ++++++-
 .../scriptTemplateProviderExtensionPoint.kt   |  5 +++-
 .../core/script/ScriptDependenciesUpdater.kt  | 30 +++++++++----------
 .../emptyAsyncResolver/script.kts             |  3 ++
 .../emptyAsyncResolver/template/template.kt   | 23 ++++++++++++++
 .../highlighting/noResolver/script.kts        |  3 ++
 .../noResolver/template/template.kt           | 13 ++++++++
 ...onfigurationHighlightingTestGenerated.java | 12 ++++++++
 9 files changed, 85 insertions(+), 18 deletions(-)
 create mode 100644 idea/testData/script/definition/highlighting/emptyAsyncResolver/script.kts
 create mode 100644 idea/testData/script/definition/highlighting/emptyAsyncResolver/template/template.kt
 create mode 100644 idea/testData/script/definition/highlighting/noResolver/script.kts
 create mode 100644 idea/testData/script/definition/highlighting/noResolver/template/template.kt

diff --git a/compiler/frontend.script/src/org/jetbrains/kotlin/script/KotlinScriptDefinitionFromAnnotatedTemplate.kt b/compiler/frontend.script/src/org/jetbrains/kotlin/script/KotlinScriptDefinitionFromAnnotatedTemplate.kt
index ca3f08afd47..db4694529cd 100644
--- a/compiler/frontend.script/src/org/jetbrains/kotlin/script/KotlinScriptDefinitionFromAnnotatedTemplate.kt
+++ b/compiler/frontend.script/src/org/jetbrains/kotlin/script/KotlinScriptDefinitionFromAnnotatedTemplate.kt
@@ -23,6 +23,7 @@ import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.name.NameUtils
 import org.jetbrains.kotlin.psi.KtScript
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
+import java.io.File
 import kotlin.reflect.KClass
 import kotlin.reflect.KFunction
 import kotlin.reflect.KParameter
@@ -36,7 +37,8 @@ open class KotlinScriptDefinitionFromAnnotatedTemplate(
         template: KClass<out Any>,
         providedResolver: DependenciesResolver? = null,
         providedScriptFilePattern: String? = null,
-        val environment: Map<String, Any?>? = null
+        val environment: Map<String, Any?>? = null,
+        val templateClasspath: List<File> = emptyList()
 ) : KotlinScriptDefinition(template) {
 
     val scriptFilePattern by lazy {
diff --git a/compiler/frontend.script/src/org/jetbrains/kotlin/script/ScriptContentLoader.kt b/compiler/frontend.script/src/org/jetbrains/kotlin/script/ScriptContentLoader.kt
index 690e56dc307..dbfaaec0f9a 100644
--- a/compiler/frontend.script/src/org/jetbrains/kotlin/script/ScriptContentLoader.kt
+++ b/compiler/frontend.script/src/org/jetbrains/kotlin/script/ScriptContentLoader.kt
@@ -67,10 +67,18 @@ class ScriptContentLoader(private val project: Project) {
                 getEnvironment(scriptDef)
         )
         ServiceManager.getService(project, ScriptReportSink::class.java)?.attachReports(file, result.reports)
-        return result.dependencies
+        return result.dependencies?.adjustByDefinition(scriptDef)
     }
 
     fun getEnvironment(scriptDef: KotlinScriptDefinition) =
             (scriptDef as? KotlinScriptDefinitionFromAnnotatedTemplate)?.environment.orEmpty()
 }
 
+fun ScriptDependencies.adjustByDefinition(
+        scriptDef: KotlinScriptDefinition
+): ScriptDependencies {
+    val additionalClasspath = (scriptDef as? KotlinScriptDefinitionFromAnnotatedTemplate)?.templateClasspath ?: return this
+    if (additionalClasspath.isEmpty()) return this
+    return copy(classpath = additionalClasspath + classpath)
+}
+
diff --git a/compiler/frontend.script/src/org/jetbrains/kotlin/script/scriptTemplateProviderExtensionPoint.kt b/compiler/frontend.script/src/org/jetbrains/kotlin/script/scriptTemplateProviderExtensionPoint.kt
index 4e6c7984c5c..4473bbfec91 100644
--- a/compiler/frontend.script/src/org/jetbrains/kotlin/script/scriptTemplateProviderExtensionPoint.kt
+++ b/compiler/frontend.script/src/org/jetbrains/kotlin/script/scriptTemplateProviderExtensionPoint.kt
@@ -72,7 +72,10 @@ fun makeScriptDefsFromTemplatesProviders(providers: Iterable<ScriptTemplatesProv
     try {
         val loader = createClassLoader(provider)
         provider.templateClassNames.map {
-            KotlinScriptDefinitionFromAnnotatedTemplate(loader.loadClass(it).kotlin, provider.resolver, provider.filePattern, provider.environment)
+            KotlinScriptDefinitionFromAnnotatedTemplate(
+                    loader.loadClass(it).kotlin, provider.resolver,
+                    provider.filePattern, provider.environment, provider.templateClasspath
+            )
         }
     }
     catch (ex: Throwable) {
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 d2b467d2eca..e959dd9ce81 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
@@ -139,11 +139,23 @@ internal class ScriptDependenciesUpdater(
             scriptDef: KotlinScriptDefinition
     ): TimeStampedJob {
         val currentTimeStamp = TimeStamps.next()
+
+        fun process(result: DependenciesResolver.ResolveResult) {
+            val lastTimeStamp = requests[file.path]?.job?.timeStamp
+            val isLastSentRequest = lastTimeStamp == null || lastTimeStamp == currentTimeStamp
+            if (isLastSentRequest) {
+                ServiceManager.getService(project, ScriptReportSink::class.java)?.attachReports(file, result.reports)
+                val resultingDependencies = (result.dependencies ?: ScriptDependencies.Empty).adjustByDefinition(scriptDef)
+                if (cache(resultingDependencies, file)) {
+                    onChange()
+                }
+            }
+        }
+
         val dependenciesResolver = scriptDef.dependencyResolver
         val newJob = if (dependenciesResolver is AsyncDependenciesResolver) {
             launch(asyncUpdatesDispatcher) {
-                processResult(
-                        file, currentTimeStamp,
+                process(
                         dependenciesResolver.resolveAsync(
                                 contentLoader.getScriptContents(scriptDef, file),
                                 contentLoader.getEnvironment(scriptDef)
@@ -154,8 +166,7 @@ internal class ScriptDependenciesUpdater(
         else {
             assert(dependenciesResolver is LegacyResolverWrapper)
             launch(legacyUpdatesDispatcher) {
-                processResult(
-                        file, currentTimeStamp,
+                process(
                         dependenciesResolver.resolve(
                                 contentLoader.getScriptContents(scriptDef, file),
                                 contentLoader.getEnvironment(scriptDef)
@@ -166,17 +177,6 @@ internal class ScriptDependenciesUpdater(
         return TimeStampedJob(newJob, currentTimeStamp)
     }
 
-    private fun processResult(file: VirtualFile, currentTimeStamp: TimeStamp, result: DependenciesResolver.ResolveResult) {
-        val lastTimeStamp = requests[file.path]?.job?.timeStamp
-        val isLastSentRequest = lastTimeStamp == null || lastTimeStamp == currentTimeStamp
-        if (isLastSentRequest) {
-            ServiceManager.getService(project, ScriptReportSink::class.java)?.attachReports(file, result.reports)
-            if (cache(result.dependencies ?: ScriptDependencies.Empty, file)) {
-                onChange()
-            }
-        }
-    }
-
     fun updateSync(file: VirtualFile, scriptDef: KotlinScriptDefinition): Boolean {
         val newDeps = contentLoader.loadContentsAndResolveDependencies(scriptDef, file) ?: ScriptDependencies.Empty
         return cache(newDeps, file)
diff --git a/idea/testData/script/definition/highlighting/emptyAsyncResolver/script.kts b/idea/testData/script/definition/highlighting/emptyAsyncResolver/script.kts
new file mode 100644
index 00000000000..328e15eb5f1
--- /dev/null
+++ b/idea/testData/script/definition/highlighting/emptyAsyncResolver/script.kts
@@ -0,0 +1,3 @@
+val j: Int = i
+
+val s: String = str
\ No newline at end of file
diff --git a/idea/testData/script/definition/highlighting/emptyAsyncResolver/template/template.kt b/idea/testData/script/definition/highlighting/emptyAsyncResolver/template/template.kt
new file mode 100644
index 00000000000..c3f2948e6f9
--- /dev/null
+++ b/idea/testData/script/definition/highlighting/emptyAsyncResolver/template/template.kt
@@ -0,0 +1,23 @@
+package custom.scriptDefinition
+
+import java.io.File
+import kotlin.script.dependencies.*
+import kotlin.script.dependencies.experimental.*
+import kotlin.script.templates.ScriptTemplateDefinition
+
+class TestDependenciesResolver : AsyncDependenciesResolver {
+    override suspend fun resolveAsync(
+            scriptContents: ScriptContents,
+            environment: Environment
+    ): DependenciesResolver.ResolveResult {
+        return ScriptDependencies.Empty.asSuccess()
+    }
+}
+
+@ScriptTemplateDefinition(TestDependenciesResolver::class, scriptFilePattern = "script.kts")
+open class Template: Base()
+
+open class Base {
+    val i = 3
+    val str = ""
+}
\ No newline at end of file
diff --git a/idea/testData/script/definition/highlighting/noResolver/script.kts b/idea/testData/script/definition/highlighting/noResolver/script.kts
new file mode 100644
index 00000000000..328e15eb5f1
--- /dev/null
+++ b/idea/testData/script/definition/highlighting/noResolver/script.kts
@@ -0,0 +1,3 @@
+val j: Int = i
+
+val s: String = str
\ No newline at end of file
diff --git a/idea/testData/script/definition/highlighting/noResolver/template/template.kt b/idea/testData/script/definition/highlighting/noResolver/template/template.kt
new file mode 100644
index 00000000000..db1336fe7fd
--- /dev/null
+++ b/idea/testData/script/definition/highlighting/noResolver/template/template.kt
@@ -0,0 +1,13 @@
+package custom.scriptDefinition
+
+import java.io.File
+import kotlin.script.dependencies.*
+import kotlin.script.templates.ScriptTemplateDefinition
+
+@ScriptTemplateDefinition(scriptFilePattern = "script.kts")
+class Template: Base()
+
+open class Base {
+    val i = 3
+    val str = ""
+}
\ No newline at end of file
diff --git a/idea/tests/org/jetbrains/kotlin/idea/script/ScriptConfigurationHighlightingTestGenerated.java b/idea/tests/org/jetbrains/kotlin/idea/script/ScriptConfigurationHighlightingTestGenerated.java
index 2d9470ffd29..3c8b406425e 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/script/ScriptConfigurationHighlightingTestGenerated.java
+++ b/idea/tests/org/jetbrains/kotlin/idea/script/ScriptConfigurationHighlightingTestGenerated.java
@@ -78,6 +78,18 @@ public class ScriptConfigurationHighlightingTestGenerated extends AbstractScript
         doTest(fileName);
     }
 
+    @TestMetadata("emptyAsyncResolver")
+    public void testEmptyAsyncResolver() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/script/definition/highlighting/emptyAsyncResolver/");
+        doTest(fileName);
+    }
+
+    @TestMetadata("noResolver")
+    public void testNoResolver() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/script/definition/highlighting/noResolver/");
+        doTest(fileName);
+    }
+
     @TestMetadata("simple")
     public void testSimple() throws Exception {
         String fileName = KotlinTestUtils.navigationMetadata("idea/testData/script/definition/highlighting/simple/");