gradle.kts: cache definitions

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

@@ -29,6 +29,7 @@ import com.intellij.util.containers.SLRUMap
 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.configuration.CompositeScriptConfigurationManager
 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
@@ -82,9 +83,15 @@ class ScriptDefinitionsManager(private val project: Project) : LazyScriptDefinit
     private val scriptDefinitionsCacheLock = ReentrantLock()
     private val scriptDefinitionsCache = SLRUMap<String, ScriptDefinition>(10, 10)
 
+    val configurations = (ScriptConfigurationManager.getInstance(project) as CompositeScriptConfigurationManager)
+
     override fun findDefinition(script: SourceCode): ScriptDefinition? {
         val locationId = script.locationId ?: return null
         if (nonScriptId(locationId)) return null
+
+        val fastPath = configurations.tryGetScriptDefinitionFast(locationId)
+        if (fastPath != null) return fastPath
+
         if (!isReady()) return null
 
         val cached = scriptDefinitionsCacheLock.withLock { scriptDefinitionsCache.get(locationId) }

idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/configuration/CompositeScriptConfigurationManager.kt

@@ -19,6 +19,7 @@ import org.jetbrains.kotlin.idea.core.script.configuration.utils.ScriptClassRoot
 import org.jetbrains.kotlin.idea.core.script.configuration.utils.ScriptClassRootsUpdater
 import org.jetbrains.kotlin.idea.core.script.configuration.utils.getKtFile
 import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.scripting.definitions.ScriptDefinition
 import org.jetbrains.kotlin.scripting.resolve.ScriptCompilationConfigurationWrapper
 
 class CompositeScriptConfigurationManager(val project: Project) : ScriptConfigurationManager {
@@ -31,6 +32,10 @@ class CompositeScriptConfigurationManager(val project: Project) : ScriptConfigur
 
     val default = DefaultScriptingSupport(this)
 
+    fun tryGetScriptDefinitionFast(locationId: String): ScriptDefinition? {
+        return classpathRoots.getLightScriptInfo(locationId)?.definition
+    }
+
     private fun getOrLoadConfiguration(
         virtualFile: VirtualFile,
         preloadedKtFile: KtFile? = null

idea/idea-core/src/org/jetbrains/kotlin/idea/core/script/configuration/utils/ScriptClassRootsCache.kt

@@ -21,6 +21,7 @@ import org.jetbrains.kotlin.idea.core.script.LOG
 import org.jetbrains.kotlin.idea.core.script.ScriptConfigurationManager.Companion.classpathEntryToVfs
 import org.jetbrains.kotlin.idea.core.script.ScriptConfigurationManager.Companion.toVfsRoots
 import org.jetbrains.kotlin.idea.util.getProjectJdkTableSafe
+import org.jetbrains.kotlin.scripting.definitions.ScriptDefinition
 import org.jetbrains.kotlin.scripting.resolve.ScriptCompilationConfigurationWrapper
 import java.io.File
 import java.lang.ref.Reference
@@ -34,14 +35,14 @@ class ScriptClassRootsCache(
     private val nonModulesSdks: Collection<Sdk>,
     val customDefinitionsUsed: Boolean
 ) {
-    abstract class LightScriptInfo {
+    abstract class LightScriptInfo(val definition: ScriptDefinition?) {
         @Volatile
         var heavyCache: Reference<HeavyScriptInfo>? = null
 
         abstract fun buildConfiguration(): ScriptCompilationConfigurationWrapper?
     }
 
-    class DirectScriptInfo(val result: ScriptCompilationConfigurationWrapper) : LightScriptInfo() {
+    class DirectScriptInfo(val result: ScriptCompilationConfigurationWrapper) : LightScriptInfo(null) {
         override fun buildConfiguration(): ScriptCompilationConfigurationWrapper = result
     }
 

idea/idea-gradle/src/org/jetbrains/kotlin/idea/scripting/gradle/GradleBuildRoot.kt

@@ -9,10 +9,12 @@ import com.intellij.openapi.project.Project
 import com.intellij.openapi.vfs.LocalFileSystem
 import com.intellij.openapi.vfs.VfsUtil
 import com.intellij.openapi.vfs.VirtualFile
+import org.jetbrains.kotlin.idea.core.script.ScriptDefinitionContributor
 import org.jetbrains.kotlin.idea.core.script.configuration.utils.ScriptClassRootsCache
 import org.jetbrains.kotlin.idea.scripting.gradle.importing.KotlinDslScriptModel
 import org.jetbrains.kotlin.scripting.definitions.ScriptDefinition
 import org.jetbrains.kotlin.scripting.definitions.findScriptDefinition
+import org.jetbrains.kotlin.scripting.resolve.KtFileScriptSource
 import org.jetbrains.kotlin.scripting.resolve.ScriptCompilationConfigurationWrapper
 import org.jetbrains.kotlin.scripting.resolve.VirtualFileScriptSource
 import org.jetbrains.kotlin.scripting.resolve.adjustByDefinition
@@ -74,23 +76,24 @@ sealed class GradleBuildRoot {
             val javaHome = context.javaHome
             javaHome?.let { builder.addSdk(it) }
 
-            val anyScript = data.models.firstOrNull()?.let {
-                LocalFileSystem.getInstance().findFileByPath(it.file)
-            }
+            val definitions = ScriptDefinitionContributor.EP_NAME.getExtensions(project)
+                .filterIsInstance<GradleScriptDefinitionsContributor>()
+                .single().definitions.toList()
 
-            val scriptDefinition: ScriptDefinition? = anyScript?.findScriptDefinition(project)
+            val root = ScriptInfoRoot(this)
 
-            if (scriptDefinition != null) {
-                builder.classes.addAll(data.templateClasspath)
-            }
+            builder.classes.addAll(data.templateClasspath)
+            data.models.forEach { script ->
+                val vFile = LocalFileSystem.getInstance().findFileByPath(script.file)
+                val def = if (vFile != null) {
+                    val src = VirtualFileScriptSource(vFile)
+                    definitions.firstOrNull { it.isScript(src) }
+                } else null
 
-            val root = ScriptInfoRoot(this, scriptDefinition)
+                builder.scripts[script.file] = ScriptInfo(root, def, script)
 
-            data.models.forEach {
-                builder.scripts[it.file] = ScriptInfo(root, it)
-
-                builder.classes.addAll(it.classPath)
-                builder.sources.addAll(it.sourcePath)
+                builder.classes.addAll(script.classPath)
+                builder.sources.addAll(script.sourcePath)
             }
         }
 
@@ -98,37 +101,56 @@ sealed class GradleBuildRoot {
          * Common info between scripts (to save bits inside ScriptInfo)
          */
         data class ScriptInfoRoot(
-            val buildRoot: Imported,
-            val scriptDefinition: ScriptDefinition?
+            val buildRoot: Imported
         )
 
-        data class ScriptInfo(
+        class ScriptInfo(
             val root: ScriptInfoRoot,
+            scriptDefinition: ScriptDefinition?,
             val model: KotlinDslScriptModel
-        ) : ScriptClassRootsCache.LightScriptInfo() {
+        ) : ScriptClassRootsCache.LightScriptInfo(scriptDefinition) {
             val buildRoot get() = root.buildRoot
 
             override fun buildConfiguration(): ScriptCompilationConfigurationWrapper? {
                 val javaHome = buildRoot.context.javaHome
-                val scriptDefinition = root.scriptDefinition
 
                 val scriptFile = File(model.file)
                 val virtualFile = VfsUtil.findFile(scriptFile.toPath(), true)!!
 
-                if (scriptDefinition == null) return null
+                if (definition == null) return null
 
                 return ScriptCompilationConfigurationWrapper.FromCompilationConfiguration(
                     VirtualFileScriptSource(virtualFile),
-                    scriptDefinition.compilationConfiguration.with {
+                    definition.compilationConfiguration.with {
                         if (javaHome != null) {
                             jvm.jdkHome(javaHome)
                         }
                         defaultImports(model.imports)
                         dependencies(JvmDependency(model.classPath.map { File(it) }))
                         ide.dependenciesSources(JvmDependency(model.sourcePath.map { File(it) }))
-                    }.adjustByDefinition(scriptDefinition)
+                    }.adjustByDefinition(definition)
                 )
             }
+
+            override fun equals(other: Any?): Boolean {
+                if (this === other) return true
+                if (javaClass != other?.javaClass) return false
+
+                other as ScriptInfo
+
+                if (root != other.root) return false
+                if (model != other.model) return false
+                if (definition != other.definition) return false
+
+                return true
+            }
+
+            override fun hashCode(): Int {
+                var result = root.hashCode()
+                result = 31 * result + model.hashCode()
+                result = 31 * result + definition.hashCode()
+                return result
+            }
         }
     }
 }