From 1513429613c2025e6528afac528815fabcf277a4 Mon Sep 17 00:00:00 2001
From: Natalia Selezneva <natalia.selezneva@jetbrains.com>
Date: Thu, 2 Apr 2020 16:02:24 +0300
Subject: [PATCH] Load script definitions in IDE on project opening

This will speed up first script opening,
will avoid unexpected indexing because of new definitions
Script definitions are also needed for Kotlin scripting settings
---
 .../core/script/ScriptDefinitionsManager.kt   | 33 ++++++++++++++++---
 idea/resources/META-INF/scripting-support.xml |  1 +
 2 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/ScriptDefinitionsManager.kt b/idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/ScriptDefinitionsManager.kt
index 4bc1937f85a..bc8744e81fa 100644
--- a/idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/ScriptDefinitionsManager.kt
+++ b/idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/ScriptDefinitionsManager.kt
@@ -16,10 +16,12 @@ import com.intellij.openapi.diagnostic.ControlFlowException
 import com.intellij.openapi.extensions.ExtensionPointName
 import com.intellij.openapi.extensions.Extensions
 import com.intellij.openapi.fileTypes.FileTypeManager
+import com.intellij.openapi.progress.runBackgroundableTask
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.projectRoots.JavaSdk
 import com.intellij.openapi.projectRoots.ex.PathUtilEx
 import com.intellij.openapi.roots.ProjectRootManager
+import com.intellij.openapi.startup.StartupActivity
 import com.intellij.openapi.vfs.VfsUtil
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.openapi.vfs.VirtualFileManager
@@ -28,7 +30,9 @@ import org.jetbrains.kotlin.idea.KotlinFileType
 import org.jetbrains.kotlin.idea.caches.project.SdkInfo
 import org.jetbrains.kotlin.idea.caches.project.getScriptRelatedModuleInfo
 import org.jetbrains.kotlin.idea.core.script.settings.KotlinScriptingSettings
+import org.jetbrains.kotlin.idea.util.application.executeOnPooledThread
 import org.jetbrains.kotlin.idea.util.application.getServiceSafe
+import org.jetbrains.kotlin.idea.util.application.isUnitTestMode
 import org.jetbrains.kotlin.idea.util.getProjectJdkTableSafe
 import org.jetbrains.kotlin.script.ScriptTemplatesProvider
 import org.jetbrains.kotlin.scripting.definitions.*
@@ -55,6 +59,20 @@ import kotlin.script.experimental.jvm.defaultJvmScriptingHostConfiguration
 import kotlin.script.experimental.jvm.util.scriptCompilationClasspathFromContextOrStdlib
 import kotlin.script.templates.standard.ScriptTemplateWithArgs
 
+class LoadScriptDefinitionsStartupActivity : StartupActivity {
+    override fun runActivity(project: Project) {
+        if (isUnitTestMode()) {
+            // In tests definitions are loaded synchronously because they are needed to analyze script
+            // In IDE script won't be highlighted before all definitions are loaded, then the highlighting will be restarted
+            ScriptDefinitionsManager.getInstance(project).reloadScriptDefinitionsIfNeeded()
+        } else {
+            executeOnPooledThread {
+                ScriptDefinitionsManager.getInstance(project).reloadScriptDefinitionsIfNeeded()
+            }
+        }
+    }
+}
+
 class ScriptDefinitionsManager(private val project: Project) : LazyScriptDefinitionProvider() {
     private var definitionsBySource = mutableMapOf<ScriptDefinitionsSource, List<ScriptDefinition>>()
     private var definitions: List<ScriptDefinition>? = null
@@ -124,7 +142,17 @@ class ScriptDefinitionsManager(private val project: Project) : LazyScriptDefinit
         return fromNewEp.dropLast(1) + fromDeprecatedEP + fromNewEp.last()
     }
 
+    fun reloadScriptDefinitionsIfNeeded() = lock.write {
+        if (definitions == null) {
+            loadScriptDefinitions()
+        }
+    }
+
     fun reloadScriptDefinitions() = lock.write {
+        loadScriptDefinitions()
+    }
+
+    private fun loadScriptDefinitions() {
         for (source in getSources()) {
             val definitions = source.safeGetDefinitions()
             definitionsBySource[source] = definitions
@@ -147,9 +175,6 @@ class ScriptDefinitionsManager(private val project: Project) : LazyScriptDefinit
     }
 
     fun isReady(): Boolean {
-        if (definitions == null) {
-            reloadScriptDefinitions()
-        }
         return definitions != null && definitionsBySource.keys.all { source ->
             // TODO: implement another API for readiness checking
             (source as? ScriptDefinitionContributor)?.isReady() != false
@@ -319,7 +344,7 @@ fun ScriptDefinitionContributor.asSource(): ScriptDefinitionsSource =
     if (this is ScriptDefinitionsSource) this
     else ScriptDefinitionSourceFromContributor(this)
 
-class StandardScriptDefinitionContributor(project: Project) : ScriptDefinitionContributor {
+class StandardScriptDefinitionContributor(val project: Project) : ScriptDefinitionContributor {
     private val standardIdeScriptDefinition = StandardIdeScriptDefinition(project)
 
     override fun getDefinitions() = listOf(standardIdeScriptDefinition)
diff --git a/idea/resources/META-INF/scripting-support.xml b/idea/resources/META-INF/scripting-support.xml
index e2c1e454ba1..58bc6cdec1b 100644
--- a/idea/resources/META-INF/scripting-support.xml
+++ b/idea/resources/META-INF/scripting-support.xml
@@ -10,5 +10,6 @@
 
     <trafficLightRendererContributor implementation="org.jetbrains.kotlin.idea.core.script.ScriptTrafficLightRendererContributor"/>
     <fileBasedIndex implementation="org.jetbrains.kotlin.idea.script.ScriptTemplatesClassRootsIndex"/>
+    <postStartupActivity implementation="org.jetbrains.kotlin.idea.core.script.LoadScriptDefinitionsStartupActivity"/>
   </extensions>
 </idea-plugin>
\ No newline at end of file