Implement advanced gradle script template selection in the provider, improve diagnostics

compiler/frontend.script/src/org/jetbrains/kotlin/script/scriptTemplateProviderExtensionPoint.kt

@@ -46,6 +46,9 @@ interface ScriptTemplatesProvider {
 
     val environment: Map<String, Any?>?
 
+    // for caching already loaded definitions, when needed
+    val scriptDefinitions: List<KotlinScriptDefinition>? get() = null
+
     companion object {
         val EP_NAME: ExtensionPointName<ScriptTemplatesProvider> =
                 ExtensionPointName.create<ScriptTemplatesProvider>("org.jetbrains.kotlin.scriptTemplatesProvider")
@@ -54,26 +57,27 @@ interface ScriptTemplatesProvider {
 
 fun makeScriptDefsFromTemplatesProviderExtensions(project: Project,
                                                   errorsHandler: ((ScriptTemplatesProvider, Exception) -> Unit) = { _, ex -> throw ex }
-): List<KotlinScriptDefinitionFromAnnotatedTemplate> =
+): List<KotlinScriptDefinition> =
         makeScriptDefsFromTemplatesProviders(Extensions.getArea(project).getExtensionPoint(ScriptTemplatesProvider.EP_NAME).extensions.asIterable(),
                                              errorsHandler)
 
 fun makeScriptDefsFromTemplatesProviders(providers: Iterable<ScriptTemplatesProvider>,
                                          errorsHandler: ((ScriptTemplatesProvider, Exception) -> Unit) = { _, ex -> throw ex }
-): List<KotlinScriptDefinitionFromAnnotatedTemplate> {
+): List<KotlinScriptDefinition> = providers.flatMap { provider ->
-    return providers.filter { it.isValid }.flatMap { provider ->
+    try {
-        try {
+        LOG.info("[kts] loading script definitions ${provider.templateClassNames} using cp: ${provider.dependenciesClasspath.joinToString(File.pathSeparator)}")
-            Logger.getInstance("makeScriptDefsFromTemplatesProviders")
+        provider.scriptDefinitions ?: run {
-                    .info("[kts] loading script definitions ${provider.templateClassNames} using cp: ${provider.dependenciesClasspath.joinToString(File.pathSeparator)}")
             val loader = URLClassLoader(provider.dependenciesClasspath.map { File(it).toURI().toURL() }.toTypedArray(), ScriptTemplatesProvider::class.java.classLoader)
             provider.templateClassNames.map {
-                val cl = loader.loadClass(it)
+                KotlinScriptDefinitionFromAnnotatedTemplate(loader.loadClass(it).kotlin, provider.resolver, provider.filePattern, provider.environment)
-                KotlinScriptDefinitionFromAnnotatedTemplate(cl.kotlin, provider.resolver, provider.filePattern, provider.environment)
             }
         }
-        catch (ex: Exception) {
-            errorsHandler(provider, ex)
-            emptyList<KotlinScriptDefinitionFromAnnotatedTemplate>()
-        }
     }
-}
+    catch (ex: Exception) {
+        LOG.info("Templates provider ${provider.id} is invalid: ${ex.message}")
+        errorsHandler(provider, ex)
+        emptyList<KotlinScriptDefinition>()
+    }
+}
+
+private val LOG = Logger.getInstance("ScriptTemplatesProviders")

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

@@ -23,6 +23,8 @@ import com.intellij.openapi.project.Project
 import org.gradle.tooling.ProjectConnection
 import org.jetbrains.kotlin.lexer.KotlinLexer
 import org.jetbrains.kotlin.lexer.KtTokens
+import org.jetbrains.kotlin.script.KotlinScriptDefinition
+import org.jetbrains.kotlin.script.KotlinScriptDefinitionFromAnnotatedTemplate
 import org.jetbrains.kotlin.script.ScriptTemplatesProvider
 import org.jetbrains.plugins.gradle.service.execution.GradleExecutionHelper
 import org.jetbrains.plugins.gradle.settings.GradleExecutionSettings
@@ -56,31 +58,51 @@ class GradleScriptTemplatesProvider(project: Project): ScriptTemplatesProvider {
     }
 
     override val id: String = "Gradle"
-    override val isValid: Boolean get() = gradleExeSettings?.gradleHome != null
+    override val isValid: Boolean get() = templatesData is TemplateDataOrError.Data
 
-    override val templateClassNames: Iterable<String> by lazy {
+    private sealed class TemplateDataOrError {
-        val loader = URLClassLoader(dependenciesClasspath.map { File(it).toURI().toURL() }.toTypedArray(), ScriptTemplatesProvider::class.java.classLoader)
+        class Data(val templateClassNames: Iterable<String>,
-        val tpl = listOf("org.gradle.kotlin.dsl.KotlinBuildScript", "org.gradle.script.lang.kotlin.KotlinBuildScript").firstOrNull {
+                   val dependenciesClasspath: Iterable<File>,
-            try {
+                   val scriptDefinitions: List<KotlinScriptDefinition>) : TemplateDataOrError()
-                // TODO: consider introducing caching mechanism (e.g. for ScriptDefinition) in ScriptTemplatesProvider to avoid double loading
+        class Error(val message: String) : TemplateDataOrError()
-                loader.loadClass(it)
+    }
-                true
+
-            }
+    private val templatesData: TemplateDataOrError by lazy {
-            catch (e: ClassNotFoundException) { false }
+
-            catch (e: NoClassDefFoundError) { false }
+        if (gradleExeSettings?.gradleHome == null) return@lazy TemplateDataOrError.Error("Unable to get Gradle home directory")
+
+        val gradleLibDir = File(gradleExeSettings!!.gradleHome, "lib").let {
+            it.takeIf { it.exists() && it.isDirectory }
+            ?: return@lazy TemplateDataOrError.Error("Invalid Gradle libraries directory $it")
         }
-        if (tpl != null) listOf(tpl) else emptyList()
+
+        for ((template, selector) in templates2DependenciesSelectors) {
+            val cp = gradleLibDir.listFiles { it /* an inference problem without explicit 'it', TODO: remove when fixed */ -> selector.matches(it.name) }.takeIf { it.isNotEmpty() } ?: continue
+
+            val loader = URLClassLoader(cp.map { it.toURI().toURL() }.toTypedArray(), ScriptTemplatesProvider::class.java.classLoader)
+            try {
+                val cl = loader.loadClass(template)
+                val def = KotlinScriptDefinitionFromAnnotatedTemplate(cl.kotlin, resolver, filePattern, environment)
+                return@lazy TemplateDataOrError.Data(listOf(template), cp.asIterable(), listOf(def))
+            }
+            catch (e: ClassNotFoundException) {}
+            catch (e: NoClassDefFoundError) {}
+        }
+
+        return@lazy TemplateDataOrError.Error("Unable to find a suitable template in the Gradle libraries directory $gradleLibDir")
     }
 
-    override val dependenciesClasspath: Iterable<String> by lazy {
+    override val templateClassNames: Iterable<String> get() = when(templatesData) {
-        gradleExeSettings?.gradleHome?.let { File(it, "lib") }
+        is GradleScriptTemplatesProvider.TemplateDataOrError.Data -> (templatesData as TemplateDataOrError.Data).templateClassNames
-                ?.let { if (it.exists()) it else null }
+        is GradleScriptTemplatesProvider.TemplateDataOrError.Error -> throw IllegalStateException((templatesData as TemplateDataOrError.Error).message)
-                ?.listFiles { file -> file.extension == "jar" && depLibsPrefixes.any { file.name.startsWith(it) } }
-                    ?.map { it.canonicalPath }
-                ?: emptyList()
     }
+
+    override val dependenciesClasspath: Iterable<String> get() = when(templatesData) {
+        is GradleScriptTemplatesProvider.TemplateDataOrError.Data -> (templatesData as TemplateDataOrError.Data).dependenciesClasspath.map { it.canonicalPath }
+        is GradleScriptTemplatesProvider.TemplateDataOrError.Error -> throw IllegalStateException((templatesData as TemplateDataOrError.Error).message)
+    }
+
     override val environment: Map<String, Any?>? by lazy {
-
         mapOf(
                 "gradleHome" to gradleExeSettings?.gradleHome?.let(::File),
                 "projectRoot" to (project.basePath ?: project.baseDir.canonicalPath)?.let(::File),
@@ -91,8 +113,15 @@ class GradleScriptTemplatesProvider(project: Project): ScriptTemplatesProvider {
                 "getScriptSectionTokens" to ::topLevelSectionCodeTextTokens)
     }
 
+    override val scriptDefinitions: List<KotlinScriptDefinition>? get() = when(templatesData) {
+        is GradleScriptTemplatesProvider.TemplateDataOrError.Data -> (templatesData as TemplateDataOrError.Data).scriptDefinitions
+        is GradleScriptTemplatesProvider.TemplateDataOrError.Error -> throw IllegalStateException((templatesData as TemplateDataOrError.Error).message)
+    }
+
     companion object {
-        private val depLibsPrefixes = listOf("gradle-kotlin-dsl", "gradle-script-kotlin", "gradle-core")
+        private val templates2DependenciesSelectors = listOf(
+                "org.gradle.kotlin.dsl.KotlinBuildScript" to Regex("^gradle-(?:kotlin-dsl|core).*\\.jar\$"),
+                "org.gradle.script.lang.kotlin.KotlinBuildScript" to Regex("^gradle-(?:script-kotlin|core).*\\.jar\$"))
     }
 }