Implement lazy script definition and light discovery in cli plugin, ...

update appropriate parts of the scripting infrastructure

plugins/scripting/scripting-cli/src/org/jetbrains/kotlin/scripting/compiler/plugin/KotlinScriptDefinitionAdapterFromNewAPI.kt

@@ -10,7 +10,7 @@ import org.jetbrains.kotlin.idea.KotlinFileType
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.name.NameUtils
 import org.jetbrains.kotlin.psi.KtScript
-import org.jetbrains.kotlin.script.KotlinScriptDefinition
+import org.jetbrains.kotlin.script.*
 import kotlin.reflect.KClass
 import kotlin.reflect.KType
 import kotlin.script.experimental.api.ScriptCompileConfigurationProperties
@@ -19,21 +19,25 @@ import kotlin.script.experimental.api.ScriptDefinitionProperties
 import kotlin.script.experimental.api.ScriptingEnvironmentProperties
 import kotlin.script.experimental.dependencies.DependenciesResolver
 import kotlin.script.experimental.jvm.impl.BridgeDependenciesResolver
+import kotlin.script.experimental.location.ScriptExpectedLocation
 
-class KotlinScriptDefinitionAdapterFromNewAPI(val scriptDefinition: ScriptDefinition) :
-    KotlinScriptDefinition(scriptDefinition.compilationConfigurator.defaultConfiguration[ScriptingEnvironmentProperties.baseClass]) {
+// temporary trick with passing Any as a template and overwriting it below, TODO: fix after introducing new script definitions hierarchy
+abstract class KotlinScriptDefinitionAdapterFromNewAPIBase : KotlinScriptDefinition(Any::class) {
 
-    override val name: String get() = scriptDefinition.properties.getOrNull(ScriptDefinitionProperties.name) ?: super.name
+    protected abstract val scriptDefinition: ScriptDefinition
+
+    protected abstract val scriptFileExtensionWithDot: String
+
+    open val baseClass: KClass<*>
+        get() = scriptDefinition.compilationConfigurator.defaultConfiguration[ScriptingEnvironmentProperties.baseClass]
+
+    override val template: KClass<*> get() = baseClass
+
+    override val name: String
+        get() = scriptDefinition.properties.getOrNull(ScriptDefinitionProperties.name) ?: "Kotlin Script"
 
-    // TODO: consider creating separate type (subtype? for kotlin scripts)
     override val fileType: LanguageFileType = KotlinFileType.INSTANCE
 
-    override val annotationsForSamWithReceivers: List<String>
-        get() = emptyList()
-
-    private val scriptFileExtensionWithDot =
-        "." + (scriptDefinition.properties.getOrNull(ScriptDefinitionProperties.fileExtension) ?: "kts")
-
     override fun isScript(fileName: String): Boolean =
         fileName.endsWith(scriptFileExtensionWithDot)
 
@@ -42,12 +46,15 @@ class KotlinScriptDefinitionAdapterFromNewAPI(val scriptDefinition: ScriptDefini
         return Name.identifier(fileBasedName.identifier.removeSuffix(scriptFileExtensionWithDot))
     }
 
+    override val annotationsForSamWithReceivers: List<String>
+        get() = emptyList()
+
     override val dependencyResolver: DependenciesResolver by lazy {
         BridgeDependenciesResolver(scriptDefinition.compilationConfigurator)
     }
 
     override val acceptedAnnotations: List<KClass<out Annotation>> by lazy {
-        scriptDefinition.compilationConfigurator.defaultConfiguration.getOrNull(ScriptCompileConfigurationProperties.refineConfigurationOnAnnotations)?.toList()
+        scriptDefinition.compilationConfigurator.defaultConfiguration.getOrNull(ScriptCompileConfigurationProperties.refineConfigurationOnAnnotations)
                 ?: emptyList()
     }
 
@@ -60,6 +67,25 @@ class KotlinScriptDefinitionAdapterFromNewAPI(val scriptDefinition: ScriptDefini
         scriptDefinition.compilationConfigurator.defaultConfiguration.getOrNull(ScriptCompileConfigurationProperties.contextVariables)?.map { (k, v) -> k to v }
                 ?: emptyList()
     }
+
+    override val additionalCompilerArguments: List<String>
+        get() = scriptDefinition.compilationConfigurator.defaultConfiguration.getOrNull(ScriptCompileConfigurationProperties.compilerOptions)
+                ?: emptyList()
+
+    override val scriptExpectedLocations: List<ScriptExpectedLocation> =
+        listOf(
+            ScriptExpectedLocation.SourcesOnly,
+            ScriptExpectedLocation.TestsOnly
+        )
 }
 
 
+class KotlinScriptDefinitionAdapterFromNewAPI(
+    override val scriptDefinition: ScriptDefinition
+) : KotlinScriptDefinitionAdapterFromNewAPIBase() {
+
+    override val name: String get() = scriptDefinition.properties.getOrNull(ScriptDefinitionProperties.name) ?: super.name
+
+    override val scriptFileExtensionWithDot =
+        "." + (scriptDefinition.properties.getOrNull(ScriptDefinitionProperties.fileExtension) ?: "kts")
+}

plugins/scripting/scripting-cli/src/org/jetbrains/kotlin/scripting/compiler/plugin/LazyScriptDefinitionFromDiscoveredClass.kt

@@ -0,0 +1,125 @@
+/*
+ * Copyright 2010-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+/*
+ * Copyright 2010-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.scripting.compiler.plugin
+
+import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSeverity
+import org.jetbrains.kotlin.cli.common.messages.MessageCollector
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.org.objectweb.asm.AnnotationVisitor
+import org.jetbrains.org.objectweb.asm.ClassReader
+import org.jetbrains.org.objectweb.asm.ClassVisitor
+import org.jetbrains.org.objectweb.asm.Opcodes
+import java.io.File
+import java.net.URLClassLoader
+import kotlin.script.experimental.annotations.KotlinScript
+import kotlin.script.experimental.annotations.KotlinScriptFileExtension
+import kotlin.script.experimental.api.*
+import kotlin.script.experimental.definitions.ScriptDefinitionFromAnnotatedBaseClass
+
+class LazyScriptDefinitionFromDiscoveredClass(
+    classBytes: ByteArray,
+    private val className: String,
+    private val classpath: List<File>,
+    private val parentClassloader: ClassLoader,
+    private val messageCollector: MessageCollector
+) : KotlinScriptDefinitionAdapterFromNewAPIBase() {
+    private val annotationsFromAsm = loadAnnotationsFromClass(classBytes)
+
+    private val classloader by lazy {
+        if (classpath.isEmpty()) parentClassloader
+        else URLClassLoader(classpath.map { it.toURI().toURL() }.toTypedArray(), parentClassloader)
+    }
+
+    override val scriptDefinition: ScriptDefinition by lazy {
+        messageCollector.report(
+            CompilerMessageSeverity.LOGGING,
+            "Configure scripting: loading script definition class $className using classpath $classpath\n.  ${Thread.currentThread().stackTrace}"
+        )
+        try {
+            val cls = classloader.loadClass(className).kotlin
+            ScriptDefinitionFromAnnotatedBaseClass(
+                ScriptingEnvironment(
+                    ScriptingEnvironmentProperties.baseClass to cls
+                )
+            )
+        } catch (ex: ClassNotFoundException) {
+            messageCollector.report(CompilerMessageSeverity.ERROR, "Cannot find script definition class $className")
+            InvalidScriptDefinition
+        } catch (ex: Exception) {
+            messageCollector.report(
+                CompilerMessageSeverity.ERROR,
+                "Error processing script definition class $className: ${ex.message}"
+            )
+            InvalidScriptDefinition
+        }
+    }
+
+    override val scriptFileExtensionWithDot: String by lazy {
+        val ext = annotationsFromAsm.find { it.name == KotlinScriptFileExtension::class.simpleName!! }?.args?.first()
+                ?: scriptDefinition.properties.let {
+                    it.getOrNull(ScriptDefinitionProperties.fileExtension) ?: "kts"
+                }
+        ".$ext"
+    }
+
+    override val name: String by lazy {
+        annotationsFromAsm.find { it.name == KotlinScript::class.simpleName!! }?.args?.first()
+                ?: super.name
+    }
+}
+
+object InvalidScriptDefinition : ScriptDefinition {
+    override val properties: ScriptDefinitionPropertiesBag = ScriptDefinitionPropertiesBag()
+    override val compilationConfigurator: ScriptCompilationConfigurator = object : ScriptCompilationConfigurator {
+        override val defaultConfiguration: ScriptCompileConfiguration = ScriptDefinitionPropertiesBag()
+    }
+    override val evaluator: ScriptEvaluator<*>? = null
+}
+
+private class BinAnnData(
+    val name: String,
+    val args: ArrayList<String> = arrayListOf()
+)
+
+private class TemplateAnnotationVisitor(val anns: ArrayList<BinAnnData> = arrayListOf()) : AnnotationVisitor(Opcodes.ASM5) {
+    override fun visit(name: String?, value: Any?) {
+        anns.last().args.add(value.toString())
+    }
+}
+
+private class TemplateClassVisitor(val annVisitor: TemplateAnnotationVisitor) : ClassVisitor(Opcodes.ASM5) {
+    override fun visitAnnotation(desc: String, visible: Boolean): AnnotationVisitor {
+        val shortName = jvmDescToClassId(desc).shortClassName.asString()
+        if (shortName.startsWith("KotlinScript")) {
+            annVisitor.anns.add(BinAnnData(shortName))
+        }
+        return annVisitor
+    }
+}
+
+private fun jvmDescToClassId(desc: String): ClassId {
+    assert(desc.startsWith("L") && desc.endsWith(";")) { "Not a JVM descriptor: $desc" }
+    val name = desc.substring(1, desc.length - 1)
+    val cid = ClassId.topLevel(FqName(name.replace('/', '.')))
+    return cid
+}
+
+private fun loadAnnotationsFromClass(fileContents: ByteArray): ArrayList<BinAnnData> {
+
+    val visitor =
+        TemplateClassVisitor(TemplateAnnotationVisitor())
+
+    ClassReader(fileContents).accept(visitor, ClassReader.SKIP_CODE or ClassReader.SKIP_DEBUG or ClassReader.SKIP_FRAMES)
+
+    return visitor.annVisitor.anns
+}
+

plugins/scripting/scripting-cli/src/org/jetbrains/kotlin/scripting/compiler/plugin/ScriptiDefinitionsFromClasspathDiscoverySource.kt

@@ -0,0 +1,165 @@
+/*
+ * Copyright 2010-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.scripting.compiler.plugin
+
+import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSeverity
+import org.jetbrains.kotlin.cli.common.messages.MessageCollector
+import org.jetbrains.kotlin.cli.jvm.config.jvmClasspathRoots
+import org.jetbrains.kotlin.config.CompilerConfiguration
+import org.jetbrains.kotlin.script.KotlinScriptDefinition
+import org.jetbrains.kotlin.script.KotlinScriptDefinitionFromAnnotatedTemplate
+import org.jetbrains.kotlin.script.ScriptDefinitionsSource
+import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
+import java.io.File
+import java.io.IOException
+import java.net.URLClassLoader
+import java.util.jar.JarFile
+import kotlin.coroutines.experimental.buildSequence
+import kotlin.script.experimental.annotations.KotlinScript
+import kotlin.script.experimental.api.ScriptingEnvironment
+import kotlin.script.experimental.api.ScriptingEnvironmentProperties
+import kotlin.script.experimental.definitions.ScriptDefinitionFromAnnotatedBaseClass
+
+internal const val SCRIPT_DEFINITION_MARKERS_PATH = "META-INF/kotlin/script/templates/"
+
+class ScriptDefinitionsFromClasspathDiscoverySource(
+    private val configuration: CompilerConfiguration,
+    private val defaultScriptDefinitionClasspath: List<File>,
+    private val scriptDefinitionParentClassloader: ClassLoader,
+    private val messageCollector: MessageCollector
+) : ScriptDefinitionsSource {
+
+    override val definitions: Sequence<KotlinScriptDefinition> = run {
+        discoverScriptTemplatesInClasspath(
+            configuration.jvmClasspathRoots,
+            defaultScriptDefinitionClasspath,
+            scriptDefinitionParentClassloader,
+            messageCollector
+        )
+    }
+}
+
+private fun discoverScriptTemplatesInClasspath(
+    classpath: Iterable<File>,
+    defaultScriptDefinitionClasspath: List<File>,
+    scriptDefinitionParentClassloader: ClassLoader,
+    messageCollector: MessageCollector
+): Sequence<LazyScriptDefinitionFromDiscoveredClass> = buildSequence {
+    for (dep in classpath) {
+        try {
+            when {
+                // checking for extension is the compiler current behaviour, so the same logic is implemented here
+                dep.isFile && dep.extension == "jar" -> {
+                    val jar = JarFile(dep)
+                    if (jar.getJarEntry(SCRIPT_DEFINITION_MARKERS_PATH) != null) {
+                        for (template in jar.entries()) {
+                            if (!template.isDirectory && template.name.startsWith(SCRIPT_DEFINITION_MARKERS_PATH)) {
+                                val templateClassName = template.name.removePrefix(SCRIPT_DEFINITION_MARKERS_PATH)
+                                val templateClass = jar.getJarEntry("${templateClassName.replace('.', '/')}.class")
+                                if (templateClass == null) {
+                                    messageCollector.report(
+                                        CompilerMessageSeverity.WARNING,
+                                        "Configure scripting: class not found $templateClassName"
+                                    )
+                                } else {
+                                    messageCollector.report(
+                                        CompilerMessageSeverity.LOGGING,
+                                        "Configure scripting: Added template $templateClassName from $dep"
+                                    )
+                                    yield(
+                                        LazyScriptDefinitionFromDiscoveredClass(
+                                            jar.getInputStream(templateClass).readBytes(),
+                                            templateClassName, listOf(dep) + jar.extractClasspath(defaultScriptDefinitionClasspath),
+                                            scriptDefinitionParentClassloader, messageCollector
+                                        )
+                                    )
+                                }
+                            }
+                        }
+                    }
+                }
+                dep.isDirectory -> {
+                    val dir = File(dep, SCRIPT_DEFINITION_MARKERS_PATH)
+                    if (dir.isDirectory) {
+                        dir.listFiles().forEach {
+                            val templateClass = File(dep, "${it.name.replace('.', '/')}.class")
+                            if (!templateClass.exists() || !templateClass.isFile) {
+                                messageCollector.report(
+                                    CompilerMessageSeverity.WARNING,
+                                    "Configure scripting: class not found ${it.name}"
+                                )
+                            } else {
+                                messageCollector.report(
+                                    CompilerMessageSeverity.LOGGING,
+                                    "Configure scripting: Added template ${it.name} from $dep"
+                                )
+                                yield(
+                                    LazyScriptDefinitionFromDiscoveredClass(
+                                        templateClass.readBytes(),
+                                        it.name, listOf(dep) + defaultScriptDefinitionClasspath,
+                                        scriptDefinitionParentClassloader, messageCollector
+                                    )
+                                )
+                            }
+                        }
+                    }
+                }
+                else -> {
+                    // assuming that invalid classpath entries will be reported elsewhere anyway, so do not spam user with additional warnings here
+                    messageCollector.report(
+                        CompilerMessageSeverity.LOGGING,
+                        "Configure scripting: Unknown classpath entry $dep"
+                    )
+                }
+            }
+        } catch (e: IOException) {
+            messageCollector.report(
+                CompilerMessageSeverity.WARNING,
+                "Configure scripting: unable to process classpath entry $dep: $e"
+            )
+        }
+    }
+}
+
+private fun JarFile.extractClasspath(defaultClasspath: List<File>): List<File> =
+    manifest.mainAttributes.getValue("Class-Path")?.split(" ")?.map(::File) ?: defaultClasspath
+
+internal fun loadScriptDefinition(
+    classloader: URLClassLoader,
+    template: String,
+    scriptResolverEnv: Map<String, Any?>,
+    messageCollector: MessageCollector
+): KotlinScriptDefinition? {
+    try {
+        val cls = classloader.loadClass(template)
+        val def =
+            if (cls.annotations.firstIsInstanceOrNull<KotlinScript>() != null) {
+                KotlinScriptDefinitionAdapterFromNewAPI(
+                    ScriptDefinitionFromAnnotatedBaseClass(
+                        ScriptingEnvironment(
+                            ScriptingEnvironmentProperties.baseClass to cls.kotlin
+                        )
+                    )
+                )
+            } else {
+                KotlinScriptDefinitionFromAnnotatedTemplate(cls.kotlin, scriptResolverEnv)
+            }
+        messageCollector.report(
+            CompilerMessageSeverity.INFO,
+            "Added script definition $template to configuration: name = ${def.name}, " +
+                    "resolver = ${def.dependencyResolver.javaClass.name}"
+        )
+        return def
+    } catch (ex: ClassNotFoundException) {
+        messageCollector.report(CompilerMessageSeverity.ERROR, "Cannot find script definition template class $template")
+    } catch (ex: Exception) {
+        messageCollector.report(
+            CompilerMessageSeverity.ERROR,
+            "Error processing script definition template $template: ${ex.message}"
+        )
+    }
+    return null
+}

plugins/scripting/scripting-cli/src/org/jetbrains/kotlin/scripting/compiler/plugin/ScriptingCompilerConfigurationExtension.kt

@@ -14,20 +14,9 @@ import org.jetbrains.kotlin.compiler.plugin.ComponentRegistrar
 import org.jetbrains.kotlin.config.CompilerConfiguration
 import org.jetbrains.kotlin.config.JVMConfigurationKeys
 import org.jetbrains.kotlin.extensions.CompilerConfigurationExtension
-import org.jetbrains.kotlin.script.KotlinScriptDefinition
-import org.jetbrains.kotlin.script.KotlinScriptDefinitionFromAnnotatedTemplate
-import org.jetbrains.kotlin.script.ScriptDefinitionsSource
 import org.jetbrains.kotlin.script.StandardScriptDefinition
-import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 import java.io.File
-import java.io.IOException
 import java.net.URLClassLoader
-import java.util.jar.JarFile
-import kotlin.coroutines.experimental.buildSequence
-import kotlin.script.experimental.annotations.KotlinScript
-import kotlin.script.experimental.api.ScriptingEnvironment
-import kotlin.script.experimental.api.ScriptingEnvironmentProperties
-import kotlin.script.experimental.definitions.ScriptDefinitionFromAnnotatedBaseClass
 
 class ScriptingCompilerConfigurationExtension(val project: MockProject) : CompilerConfigurationExtension {
 
@@ -67,7 +56,8 @@ class ScriptingCompilerConfigurationExtension(val project: MockProject) : Compil
                 JVMConfigurationKeys.SCRIPT_DEFINITIONS_SOURCES,
                 ScriptDefinitionsFromClasspathDiscoverySource(
                     configuration,
-                    scriptResolverEnv,
+                    emptyList(),
+                    Thread.currentThread().contextClassLoader, // TODO: consider isolation here
                     messageCollector
                 )
             )
@@ -81,72 +71,6 @@ class ScriptingCompilerConfigurationComponentRegistrar : ComponentRegistrar {
     }
 }
 
-class ScriptDefinitionsFromClasspathDiscoverySource(
-    private val configuration: CompilerConfiguration,
-    private val scriptResolverEnv: Map<String, Any?>,
-    private val messageCollector: MessageCollector
-) : ScriptDefinitionsSource {
-
-    override val definitions: Sequence<KotlinScriptDefinition> = run {
-        val classpath = configuration.jvmClasspathRoots
-        // TODO: consider using escaping to allow kotlin escaped names in class names
-        val classloader =
-            URLClassLoader(classpath.map { it.toURI().toURL() }.toTypedArray(), Thread.currentThread().contextClassLoader)
-        discoverScriptTemplatesInClasspath(configuration, messageCollector).mapNotNull {
-            loadScriptDefinition(classloader, it, scriptResolverEnv, messageCollector)
-        }
-    }
-}
-
-private fun discoverScriptTemplatesInClasspath(configuration: CompilerConfiguration, messageCollector: MessageCollector): Sequence<String> {
-    val templatesPath = "META-INF/kotlin/script/templates/"
-    return buildSequence {
-        for (dep in configuration.jvmClasspathRoots) {
-            when {
-                dep.isFile -> {
-                    // this is the compiler behaviour, so the same logic implemented here
-                    if (dep.extension == "jar") {
-                        try {
-                            with(JarFile(dep)) {
-                                for (template in entries()) {
-                                    if (!template.isDirectory && template.name.startsWith(templatesPath)) {
-                                        val templateClassName = template.name.removePrefix(templatesPath)
-                                        yield(templateClassName)
-                                        messageCollector.report(
-                                            CompilerMessageSeverity.LOGGING,
-                                            "Configure scripting: Added template $templateClassName from $dep"
-                                        )
-                                    }
-                                }
-                            }
-                        } catch (e: IOException) {
-                            messageCollector.report(
-                                CompilerMessageSeverity.WARNING,
-                                "Configure scripting: unable to process classpath entry $dep: $e"
-                            )
-                        }
-                    }
-                }
-                dep.isDirectory -> {
-                    val dir = File(dep, templatesPath)
-                    if (dir.isDirectory) {
-                        dir.listFiles().forEach {
-                            yield(it.name)
-                            messageCollector.report(
-                                CompilerMessageSeverity.LOGGING,
-                                "Configure scripting: Added template ${it.name} from $dep"
-                            )
-                        }
-                    }
-                }
-                else -> {
-                    // assuming that invalid classpath entries will be reported elsewhere anyway, so do not spam user with additional warnings here
-                    messageCollector.report(CompilerMessageSeverity.LOGGING, "Configure scripting: Unknown classpath entry $dep")
-                }
-            }
-        }
-    }
-}
 
 fun configureScriptDefinitions(
     scriptTemplates: List<String>,
@@ -161,7 +85,12 @@ fun configureScriptDefinitions(
             URLClassLoader(classpath.map { it.toURI().toURL() }.toTypedArray(), Thread.currentThread().contextClassLoader)
         var hasErrors = false
         for (template in scriptTemplates) {
-            val def = loadScriptDefinition(classloader, template, scriptResolverEnv, messageCollector)
+            val def = loadScriptDefinition(
+                classloader,
+                template,
+                scriptResolverEnv,
+                messageCollector
+            )
             if (!hasErrors && def == null) hasErrors = true
             if (def != null) {
                 configuration.add(JVMConfigurationKeys.SCRIPT_DEFINITIONS, def)
@@ -174,35 +103,3 @@ fun configureScriptDefinitions(
     }
 }
 
-private fun loadScriptDefinition(
-    classloader: URLClassLoader,
-    template: String,
-    scriptResolverEnv: Map<String, Any?>,
-    messageCollector: MessageCollector
-): KotlinScriptDefinition? {
-    try {
-        val cls = classloader.loadClass(template)
-        val def =
-            if (cls.annotations.firstIsInstanceOrNull<KotlinScript>() != null) {
-                KotlinScriptDefinitionAdapterFromNewAPI(
-                    ScriptDefinitionFromAnnotatedBaseClass(ScriptingEnvironment(ScriptingEnvironmentProperties.baseClass to cls.kotlin))
-                )
-            } else {
-                KotlinScriptDefinitionFromAnnotatedTemplate(cls.kotlin, scriptResolverEnv)
-            }
-        messageCollector.report(
-            CompilerMessageSeverity.INFO,
-            "Added script definition $template to configuration: name = ${def.name}, " +
-                    "resolver = ${def.dependencyResolver.javaClass.name}"
-        )
-        return def
-    } catch (ex: ClassNotFoundException) {
-        messageCollector.report(CompilerMessageSeverity.ERROR, "Cannot find script definition template class $template")
-    } catch (ex: Exception) {
-        messageCollector.report(
-            CompilerMessageSeverity.ERROR,
-            "Error processing script definition template $template: ${ex.message}"
-        )
-    }
-    return null
-}