Add lazy discovery test

libraries/scripting/common/src/kotlin/script/experimental/definitions/definitionFromAnnotation.kt

@@ -22,7 +22,7 @@ open class ScriptDefinitionFromAnnotatedBaseClass(val environment: ScriptingEnvi
             ?: throw IllegalArgumentException("${ERROR_MSG_PREFIX}Expecting baseClass parameter in the scripting environment")
 
     private val mainAnnotation = baseClass.findAnnotation<KotlinScript>()
-            ?: throw IllegalArgumentException("${ERROR_MSG_PREFIX}Expecting KotlinScript on the $baseClass")
+            ?: throw IllegalArgumentException("${ERROR_MSG_PREFIX}Expecting KotlinScript annotation on the $baseClass")
 
     private val explicitDefinition: ScriptDefinition? =
         baseClass.findAnnotation<KotlinScriptDefinition>()?.definition.takeIf { it != this::class }?.let { it.instantiateScriptHandler() }

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

@@ -29,12 +29,13 @@ 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 {
+        // should use this cl to allow smooth interop with classes explicitly mentioned here, see e.g. scriptDefinition body
+        val parentClassloader = LazyScriptDefinitionFromDiscoveredClass::class.java.classLoader
         if (classpath.isEmpty()) parentClassloader
         else URLClassLoader(classpath.map { it.toURI().toURL() }.toTypedArray(), parentClassloader)
     }

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

@@ -7,8 +7,6 @@ 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
@@ -26,26 +24,23 @@ import kotlin.script.experimental.definitions.ScriptDefinitionFromAnnotatedBaseC
 internal const val SCRIPT_DEFINITION_MARKERS_PATH = "META-INF/kotlin/script/templates/"
 
 class ScriptDefinitionsFromClasspathDiscoverySource(
-    private val configuration: CompilerConfiguration,
+    private val classpath: List<File>,
     private val defaultScriptDefinitionClasspath: List<File>,
-    private val scriptDefinitionParentClassloader: ClassLoader,
     private val messageCollector: MessageCollector
 ) : ScriptDefinitionsSource {
 
     override val definitions: Sequence<KotlinScriptDefinition> = run {
         discoverScriptTemplatesInClasspath(
-            configuration.jvmClasspathRoots,
+            classpath,
             defaultScriptDefinitionClasspath,
-            scriptDefinitionParentClassloader,
             messageCollector
         )
     }
 }
 
-private fun discoverScriptTemplatesInClasspath(
+internal fun discoverScriptTemplatesInClasspath(
     classpath: Iterable<File>,
     defaultScriptDefinitionClasspath: List<File>,
-    scriptDefinitionParentClassloader: ClassLoader,
     messageCollector: MessageCollector
 ): Sequence<LazyScriptDefinitionFromDiscoveredClass> = buildSequence {
     for (dep in classpath) {
@@ -73,7 +68,7 @@ private fun discoverScriptTemplatesInClasspath(
                                         LazyScriptDefinitionFromDiscoveredClass(
                                             jar.getInputStream(templateClass).readBytes(),
                                             templateClassName, listOf(dep) + jar.extractClasspath(defaultScriptDefinitionClasspath),
-                                            scriptDefinitionParentClassloader, messageCollector
+                                            messageCollector
                                         )
                                     )
                                 }
@@ -100,7 +95,7 @@ private fun discoverScriptTemplatesInClasspath(
                                     LazyScriptDefinitionFromDiscoveredClass(
                                         templateClass.readBytes(),
                                         it.name, listOf(dep) + defaultScriptDefinitionClasspath,
-                                        scriptDefinitionParentClassloader, messageCollector
+                                        messageCollector
                                     )
                                 )
                             }

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

@@ -55,9 +55,8 @@ class ScriptingCompilerConfigurationExtension(val project: MockProject) : Compil
             configuration.add(
                 JVMConfigurationKeys.SCRIPT_DEFINITIONS_SOURCES,
                 ScriptDefinitionsFromClasspathDiscoverySource(
-                    configuration,
+                    configuration.jvmClasspathRoots,
                     emptyList(),
-                    Thread.currentThread().contextClassLoader, // TODO: consider isolation here
                     messageCollector
                 )
             )

plugins/scripting/scripting-cli/testData/lazyDefinitions/definitions/TestScriptWithReceivers.kt

@@ -0,0 +1,17 @@
+
+import kotlin.script.experimental.annotations.*
+import kotlin.script.experimental.api.*
+import kotlin.script.experimental.util.*
+import kotlin.script.experimental.misc.*
+import kotlin.reflect.full.starProjectedType
+
+object TestScriptWithReceiversConfiguration : ArrayList<Pair<TypedKey<*>, Any?>>(
+    listOf(
+        ScriptCompileConfigurationProperties.scriptImplicitReceivers(String::class.starProjectedType)
+    )
+)
+
+@KotlinScript
+@KotlinScriptFileExtension("1.kts")
+@KotlinScriptDefaultCompilationConfiguration(TestScriptWithReceiversConfiguration::class)
+abstract class TestScriptWithReceivers

plugins/scripting/scripting-cli/testData/lazyDefinitions/definitions/TestScriptWithSimpleEnvVars.kt

@@ -0,0 +1,18 @@
+
+import kotlin.script.experimental.annotations.*
+import kotlin.script.experimental.api.*
+import kotlin.script.experimental.util.*
+import kotlin.script.experimental.misc.*
+import kotlin.reflect.full.starProjectedType
+
+object TestScriptWithSimpleEnvVarsConfiguration : ArrayList<Pair<TypedKey<*>, Any?>>(
+    listOf(
+        ScriptCompileConfigurationProperties.contextVariables("stringVar1" to String::class.starProjectedType)
+    )
+)
+
+@KotlinScript
+@KotlinScriptFileExtension("2.kts")
+@KotlinScriptDefaultCompilationConfiguration(TestScriptWithSimpleEnvVarsConfiguration::class)
+abstract class TestScriptWithSimpleEnvVars
+

plugins/scripting/scripting-cli/testData/lazyDefinitions/scripts/simpleEnvVars.2.kts

@@ -0,0 +1,2 @@
+
+val res = stringVar1.drop(4)

plugins/scripting/scripting-cli/testData/lazyDefinitions/scripts/stringReceiver.1.kts

@@ -0,0 +1,2 @@
+
+val res = drop(4)

plugins/scripting/scripting-cli/tests/org/jetbrains/kotlin/scripting/compiler/plugin/ScriptingCompilerPluginTest.kt

@@ -5,11 +5,60 @@
 
 package org.jetbrains.kotlin.scripting.compiler.plugin
 
+import junit.framework.TestCase
+import org.jetbrains.kotlin.cli.common.CLIConfigurationKeys
+import org.jetbrains.kotlin.cli.common.ExitCode
+import org.jetbrains.kotlin.cli.common.messages.CompilerMessageLocation
+import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSeverity
+import org.jetbrains.kotlin.cli.common.messages.MessageCollector
+import org.jetbrains.kotlin.cli.jvm.K2JVMCompiler
+import org.jetbrains.kotlin.cli.jvm.compiler.EnvironmentConfigFiles
+import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
+import org.jetbrains.kotlin.cli.jvm.compiler.KotlinToJVMBytecodeCompiler
+import org.jetbrains.kotlin.cli.jvm.config.addJvmClasspathRoot
+import org.jetbrains.kotlin.cli.jvm.config.addJvmClasspathRoots
 import org.jetbrains.kotlin.config.CompilerConfiguration
+import org.jetbrains.kotlin.config.JVMConfigurationKeys
+import org.jetbrains.kotlin.config.addKotlinSourceRoots
+import org.jetbrains.kotlin.script.KotlinScriptDefinition
+import org.jetbrains.kotlin.test.ConfigurationKind
+import org.jetbrains.kotlin.test.KotlinTestUtils
 import org.jetbrains.kotlin.test.TestCaseWithTmpdir
+import org.jetbrains.kotlin.test.TestJdkKind
+import org.jetbrains.kotlin.utils.KotlinPaths
+import org.jetbrains.kotlin.utils.PathUtil
 import org.junit.Assert
+import java.io.File
 
 class ScriptingCompilerPluginTest : TestCaseWithTmpdir() {
+
+    companion object {
+        const val TEST_DATA_DIR = "plugins/scripting/scripting-cli/testData"
+    }
+
+    private val kotlinPaths: KotlinPaths by lazy {
+        val paths = PathUtil.kotlinPathsForDistDirectory
+        TestCase.assertTrue("Lib directory doesn't exist. Run 'ant dist'", paths.libPath.absoluteFile.isDirectory)
+        paths
+    }
+
+    val compilerClasspath = listOf(kotlinPaths.compilerPath)
+    val runtimeClasspath = listOf( kotlinPaths.stdlibPath, kotlinPaths.scriptRuntimePath, kotlinPaths.reflectPath)
+    val scriptingClasspath = listOf("kotlin-scripting-common.jar", "kotlin-scripting-misc.jar").map { File(kotlinPaths.libPath, it) }
+
+    private fun createEnvironment(
+        sources: List<String>, destDir: File, messageCollector: MessageCollector, confBody: CompilerConfiguration.() -> Unit
+    ): KotlinCoreEnvironment {
+        val configuration = KotlinTestUtils.newConfiguration(ConfigurationKind.NO_KOTLIN_REFLECT, TestJdkKind.FULL_JDK).apply {
+            put<MessageCollector>(CLIConfigurationKeys.MESSAGE_COLLECTOR_KEY, messageCollector)
+            addKotlinSourceRoots(sources)
+            put(JVMConfigurationKeys.OUTPUT_DIRECTORY, destDir)
+            confBody()
+        }
+
+        return KotlinCoreEnvironment.createForTests(testRootDisposable, configuration, EnvironmentConfigFiles.JVM_CONFIG_FILES)
+    }
+
     fun testScriptResolverEnvironmentArgsParsing() {
 
         val longStr = (1..100).joinToString { """\" $it aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa \\""" }
@@ -30,4 +79,114 @@ class ScriptingCompilerPluginTest : TestCaseWithTmpdir() {
             res
         )
     }
+
+    fun testLazyScriptDefinition() {
+
+        // Two tests in one function: the discovery code separately, and as a part of regular compilation
+        // tests are combined to avoid double compilation of script definition modules
+
+        val defsOut = File(tmpdir, "testLazyScriptDefinition/out/defs")
+        val defsSrc = File(TEST_DATA_DIR, "lazyDefinitions/definitions")
+        val scriptsOut = File(tmpdir, "testLazyScriptDefinition/out/scripts")
+        val scriptsSrc = File(TEST_DATA_DIR, "lazyDefinitions/scripts")
+        val scriptsOut2 = File(tmpdir, "testLazyScriptDefinition/out/scripts2")
+        val defClasses = listOf("TestScriptWithReceivers", "TestScriptWithSimpleEnvVars")
+
+        val messageCollector = TestMessageCollector()
+
+        val definitionsCompileResult = KotlinToJVMBytecodeCompiler.compileBunchOfSources(
+            createEnvironment(defClasses.map { File(defsSrc, "$it.kt").canonicalPath }, defsOut, messageCollector) {
+                addJvmClasspathRoots(runtimeClasspath)
+                addJvmClasspathRoots(scriptingClasspath)
+            }
+        )
+
+        assertTrue(definitionsCompileResult) {
+            "Compilation of script definitions failed: $messageCollector"
+        }
+
+        val templatesDir = File(defsOut, SCRIPT_DEFINITION_MARKERS_PATH).also { it.mkdirs() }
+        for (def in defClasses) {
+            File(templatesDir, def).createNewFile()
+        }
+
+        messageCollector.clear()
+
+        val lazyDefsSeq =
+            discoverScriptTemplatesInClasspath(listOf(defsOut), emptyList(), messageCollector)
+
+        assertTrue(messageCollector.messages.isEmpty()) {
+            "Unexpected messages from discovery sequence (should be empty):\n$messageCollector"
+        }
+
+        val lazyDefs = lazyDefsSeq.toList()
+
+        for (def in defClasses) {
+            assertTrue(
+                messageCollector.messages.any { it.message.contains("Configure scripting: Added template $def") }
+            ) {
+                "Missing messages from discovery sequence (should contain \"Added template $def\"):\n$messageCollector"
+            }
+            assertTrue(
+                messageCollector.messages.none { it.message.contains("Configure scripting: loading script definition class $def") }
+            ) {
+                "Unexpected messages from discovery sequence (should not contain \"loading script definition class $def\"):\n$messageCollector"
+            }
+        }
+
+        messageCollector.clear()
+
+        val scriptFiles = scriptsSrc.listFiles { file: File -> file.extension == "kts" }.map { it.canonicalPath}
+
+        val scriptsCompileEnv = createEnvironment(scriptFiles, scriptsOut, messageCollector) {
+            addJvmClasspathRoots(runtimeClasspath)
+            addJvmClasspathRoots(scriptingClasspath)
+            addJvmClasspathRoot(defsOut)
+            addAll(JVMConfigurationKeys.SCRIPT_DEFINITIONS, lazyDefs)
+        }
+
+        val res = KotlinToJVMBytecodeCompiler.compileBunchOfSources(scriptsCompileEnv)
+
+        assertTrue(res) {
+            "Failed to compile scripts:\n$messageCollector"
+        }
+
+        val cp = (runtimeClasspath + scriptingClasspath + defsOut).joinToString(File.pathSeparator)
+        val exitCode = K2JVMCompiler().exec(System.err, "-cp", cp, *(scriptFiles.toTypedArray()), "-d", scriptsOut2.canonicalPath)
+
+        Assert.assertEquals(ExitCode.OK, exitCode)
+    }
 }
+
+class TestMessageCollector : MessageCollector {
+    data class Message(val severity: CompilerMessageSeverity, val message: String, val location: CompilerMessageLocation?)
+
+    val messages = arrayListOf<Message>()
+
+    override fun clear() {
+        messages.clear()
+    }
+
+    override fun report(severity: CompilerMessageSeverity, message: String, location: CompilerMessageLocation?) {
+        messages.add(Message(severity, message, location))
+    }
+
+    override fun hasErrors(): Boolean = messages.any { it.severity == CompilerMessageSeverity.EXCEPTION || it.severity == CompilerMessageSeverity.ERROR }
+
+    override fun toString(): String {
+        return messages.joinToString("\n") { "${it.severity}: ${it.message}${it.location?.let{" at $it"} ?: ""}" }
+    }
+}
+
+fun TestMessageCollector.assertHasMessage(msg: String, desiredSeverity: CompilerMessageSeverity? = null) {
+    assert(messages.any { it.message.contains(msg) && (desiredSeverity == null || it.severity == desiredSeverity) }) {
+        "Expecting message \"$msg\" with severity ${desiredSeverity?.toString() ?: "Any"}, actual:\n" +
+                messages.joinToString("\n") { it.severity.toString() + ": " + it.message }
+    }
+}
+
+fun assertTrue(exp: Boolean, msg: () -> String) {
+    if (!exp) {
+        Assert.fail(msg())
+    }
+}