Simplify and fix classloaders handling on script evaluation

#KT-37558 fixed

libraries/scripting/jvm-host-test/test/kotlin/script/experimental/jvmhost/test/CachingTest.kt

@@ -19,6 +19,7 @@ import kotlin.script.experimental.host.with
 import kotlin.script.experimental.jvm.*
 import kotlin.script.experimental.jvm.impl.KJvmCompiledScript
 import kotlin.script.experimental.jvm.util.KotlinJars
+import kotlin.script.experimental.jvm.util.classpathFromClass
 import kotlin.script.experimental.jvmhost.BasicJvmScriptingHost
 import kotlin.script.experimental.jvmhost.CompiledScriptJarsCache
 import kotlin.script.experimental.jvmhost.JvmScriptCompiler
@@ -40,7 +41,10 @@ class CachingTest : TestCase() {
     @Test
     fun testSimpleImportWithMemoryCache() {
         val cache = SimpleMemoryScriptsCache()
-        checkWithCache(cache, scriptWithImport, scriptWithImportExpectedOutput) { makeSimpleConfigurationWithTestImport() }
+        checkWithCache(
+            cache, scriptWithImport, scriptWithImportExpectedOutput,
+            compilationConfiguration = { makeSimpleConfigurationWithTestImport() }
+        )
     }
 
 
@@ -60,7 +64,10 @@ class CachingTest : TestCase() {
             val cache = FileBasedScriptCache(cacheDir)
             Assert.assertTrue(cache.baseDir.listFiles().isEmpty())
 
-            checkWithCache(cache, scriptWithImport, scriptWithImportExpectedOutput) { makeSimpleConfigurationWithTestImport() }
+            checkWithCache(
+                cache, scriptWithImport, scriptWithImportExpectedOutput,
+                compilationConfiguration = { makeSimpleConfigurationWithTestImport() }
+            )
         }
     }
 
@@ -84,7 +91,10 @@ class CachingTest : TestCase() {
             val cache = TestCompiledScriptJarsCache(cacheDir)
             Assert.assertTrue(cache.baseDir.listFiles()!!.isEmpty())
 
-            checkWithCache(cache, scriptWithImport, scriptWithImportExpectedOutput) { makeSimpleConfigurationWithTestImport() }
+            checkWithCache(
+                cache, scriptWithImport, scriptWithImportExpectedOutput,
+                compilationConfiguration = { makeSimpleConfigurationWithTestImport() }
+            )
 
             // cannot make it work in this form - it requires a dependency on the current test classes, but classes directory seems
             // not work when specified in the manifest
@@ -95,31 +105,54 @@ class CachingTest : TestCase() {
         }
     }
 
+    @Test
+    fun testImplicitReceiversWithJarCache() {
+        withTempDir("scriptingTestJarCache") { cacheDir ->
+            val cache = TestCompiledScriptJarsCache(cacheDir)
+            Assert.assertTrue(cache.baseDir.listFiles()!!.isEmpty())
+
+            checkWithCache(
+                cache, simpleScript, simpleScriptExpectedOutput, checkDirectEval = false,
+                compilationConfiguration = {
+                    updateClasspath(classpathFromClass<ScriptingHostTest>()) // the class defined here should be in the classpath
+                    implicitReceivers(Implicit::class)
+                },
+                evaluationConfiguration = {
+                    implicitReceivers(Implicit)
+                }
+            )
+        }
+    }
+
     private fun checkWithCache(
-        cache: ScriptingCacheWithCounters, script: String, expectedOutput: List<String>,
+        cache: ScriptingCacheWithCounters, script: String, expectedOutput: List<String>, checkDirectEval: Boolean = true,
-        configurationBuilder: ScriptCompilationConfiguration.Builder.() -> Unit = {}
+        compilationConfiguration: ScriptCompilationConfiguration.Builder.() -> Unit = {},
+        evaluationConfiguration: ScriptEvaluationConfiguration.Builder.() -> Unit = {}
     ) {
-        val hostConfiguration = defaultJvmScriptingHostConfiguration.with {
+        val myHostConfiguration = defaultJvmScriptingHostConfiguration.with {
             jvm {
                 baseClassLoader.replaceOnlyDefault(null)
                 compilationCache(cache)
             }
         }
-        val compiler = JvmScriptCompiler(hostConfiguration)
+        val compiler = JvmScriptCompiler(myHostConfiguration)
         val evaluator = BasicJvmScriptEvaluator()
         val host = BasicJvmScriptingHost(compiler = compiler, evaluator = evaluator)
 
-        val scriptCompilationConfiguration = ScriptCompilationConfiguration(body = configurationBuilder).with {
+        val scriptCompilationConfiguration = ScriptCompilationConfiguration(body = compilationConfiguration).with {
             updateClasspath(KotlinJars.kotlinScriptStandardJarsWithReflect)
+            hostConfiguration.update { myHostConfiguration }
         }
 
+        val scriptEvaluationConfiguration = ScriptEvaluationConfiguration(body = evaluationConfiguration)
+
         Assert.assertEquals(0, cache.storedScripts)
-        var compiledScript: CompiledScript<*>? = null
+        var compiledScript: CompiledScript? = null
         val output = captureOut {
             runBlocking {
                 compiler(script.toScriptSource(), scriptCompilationConfiguration).onSuccess {
                     compiledScript = it
-                    evaluator(it)
+                    evaluator(it, scriptEvaluationConfiguration)
                 }.throwOnFailure()
             }
         }.lines()
@@ -127,34 +160,40 @@ class CachingTest : TestCase() {
         Assert.assertEquals(1, cache.storedScripts)
         Assert.assertEquals(0, cache.retrievedScripts)
 
-        val cachedScript = cache.get(script.toScriptSource(), scriptCompilationConfiguration)
+        if (checkDirectEval) {
-        Assert.assertNotNull(cachedScript)
+            val cachedScript = cache.get(script.toScriptSource(), scriptCompilationConfiguration)
-        Assert.assertEquals(1, cache.retrievedScripts)
+            Assert.assertNotNull(cachedScript)
+            Assert.assertEquals(1, cache.retrievedScripts)
 
-        val compiledScriptClassRes = runBlocking { compiledScript!!.getClass(null) }
+            val compiledScriptClassRes = runBlocking { compiledScript!!.getClass(null) }
-        val cachedScriptClassRes = runBlocking { cachedScript!!.getClass(null) }
+            val cachedScriptClassRes = runBlocking { cachedScript!!.getClass(null) }
 
-        val compiledScriptClass = compiledScriptClassRes.valueOrThrow()
+            val compiledScriptClass = compiledScriptClassRes.valueOrThrow()
-        val cachedScriptClass = cachedScriptClassRes.valueOrThrow()
+            val cachedScriptClass = cachedScriptClassRes.valueOrThrow()
 
-        Assert.assertEquals(compiledScriptClass.qualifiedName, cachedScriptClass.qualifiedName)
+            Assert.assertEquals(compiledScriptClass.qualifiedName, cachedScriptClass.qualifiedName)
-        Assert.assertEquals(compiledScriptClass.supertypes, cachedScriptClass.supertypes)
+            Assert.assertEquals(compiledScriptClass.supertypes, cachedScriptClass.supertypes)
 
-        val output2 = captureOut {
+            val output2 = captureOut {
-            runBlocking {
+                runBlocking {
-                evaluator(cachedScript!!).throwOnFailure()
+                    evaluator(cachedScript!!, scriptEvaluationConfiguration).throwOnFailure()
-            }
+                }
-        }.lines()
+            }.lines()
-        Assert.assertEquals(output, output2)
+            Assert.assertEquals(output, output2)
+        }
 
         val output3 = captureOut {
-            host.eval(script.toScriptSource(), scriptCompilationConfiguration, null).throwOnFailure()
+            host.eval(script.toScriptSource(), scriptCompilationConfiguration, scriptEvaluationConfiguration).throwOnFailure()
         }.lines()
-        Assert.assertEquals(2, cache.retrievedScripts)
+        Assert.assertEquals(if (checkDirectEval) 2 else 1, cache.retrievedScripts)
         Assert.assertEquals(output, output3)
     }
 }
 
+object Implicit {
+    val ix = 0
+}
+
 private interface ScriptingCacheWithCounters : CompiledJvmScriptsCache {
 
     val storedScripts: Int
@@ -163,7 +202,7 @@ private interface ScriptingCacheWithCounters : CompiledJvmScriptsCache {
 
 private class SimpleMemoryScriptsCache : ScriptingCacheWithCounters {
 
-    internal val data = hashMapOf<Pair<SourceCode, Map<*, *>>, CompiledScript<*>>()
+    internal val data = hashMapOf<Pair<SourceCode, Map<*, *>>, CompiledScript>()
 
     private var _storedScripts = 0
     private var _retrievedScripts = 0
@@ -174,11 +213,11 @@ private class SimpleMemoryScriptsCache : ScriptingCacheWithCounters {
     override val retrievedScripts: Int
         get() = _retrievedScripts
 
-    override fun get(script: SourceCode, scriptCompilationConfiguration: ScriptCompilationConfiguration): CompiledScript<*>? =
+    override fun get(script: SourceCode, scriptCompilationConfiguration: ScriptCompilationConfiguration): CompiledScript? =
         data[script to scriptCompilationConfiguration.notTransientData]?.also { _retrievedScripts++ }
 
     override fun store(
-        compiledScript: CompiledScript<*>,
+        compiledScript: CompiledScript,
         script: SourceCode,
         scriptCompilationConfiguration: ScriptCompilationConfiguration
     ) {
@@ -189,13 +228,13 @@ private class SimpleMemoryScriptsCache : ScriptingCacheWithCounters {
 
 private class FileBasedScriptCache(val baseDir: File) : ScriptingCacheWithCounters {
 
-    override fun get(script: SourceCode, scriptCompilationConfiguration: ScriptCompilationConfiguration): CompiledScript<*>? {
+    override fun get(script: SourceCode, scriptCompilationConfiguration: ScriptCompilationConfiguration): CompiledScript? {
         val file = File(baseDir, uniqueScriptHash(script, scriptCompilationConfiguration))
         return if (!file.exists()) null else file.readCompiledScript(scriptCompilationConfiguration)?.also { retrievedScripts++ }
     }
 
     override fun store(
-        compiledScript: CompiledScript<*>,
+        compiledScript: CompiledScript,
         script: SourceCode,
         scriptCompilationConfiguration: ScriptCompilationConfiguration
     ) {
@@ -220,11 +259,11 @@ class TestCompiledScriptJarsCache(val baseDir: File) : CompiledScriptJarsCache(
         File(baseDir, uniqueScriptHash(script, scriptCompilationConfiguration) + ".jar")
     }), ScriptingCacheWithCounters
 {
-    override fun get(script: SourceCode, scriptCompilationConfiguration: ScriptCompilationConfiguration): CompiledScript<*>? =
+    override fun get(script: SourceCode, scriptCompilationConfiguration: ScriptCompilationConfiguration): CompiledScript? =
         super.get(script, scriptCompilationConfiguration)?.also { retrievedScripts++ }
 
     override fun store(
-        compiledScript: CompiledScript<*>,
+        compiledScript: CompiledScript,
         script: SourceCode,
         scriptCompilationConfiguration: ScriptCompilationConfiguration
     ) {
@@ -250,10 +289,10 @@ internal fun uniqueScriptHash(script: SourceCode, scriptCompilationConfiguration
     return digestWrapper.digest().toHexString()
 }
 
-private fun File.readCompiledScript(scriptCompilationConfiguration: ScriptCompilationConfiguration): CompiledScript<*> {
+private fun File.readCompiledScript(scriptCompilationConfiguration: ScriptCompilationConfiguration): CompiledScript {
     return inputStream().use { fs ->
         ObjectInputStream(fs).use {
-            it.readObject() as KJvmCompiledScript<*>
+            it.readObject() as KJvmCompiledScript
         }
     }
 }

libraries/scripting/jvm-host/src/kotlin/script/experimental/jvmhost/CompiledScriptJarsCache.kt

@@ -8,11 +8,8 @@ package kotlin.script.experimental.jvmhost
 import java.io.File
 import java.net.URLClassLoader
 import java.util.jar.JarInputStream
-import kotlin.script.experimental.api.CompiledScript
+import kotlin.reflect.KClass
-import kotlin.script.experimental.api.ScriptCompilationConfiguration
+import kotlin.script.experimental.api.*
-import kotlin.script.experimental.api.SourceCode
-import kotlin.script.experimental.api.hostConfiguration
-import kotlin.script.experimental.host.ScriptingHostConfiguration
 import kotlin.script.experimental.jvm.CompiledJvmScriptsCache
 import kotlin.script.experimental.jvm.baseClassLoader
 import kotlin.script.experimental.jvm.impl.KJvmCompiledScript
@@ -33,12 +30,7 @@ open class CompiledScriptJarsCache(val scriptToFile: (SourceCode, ScriptCompilat
                 it.manifest.mainAttributes.getValue("Main-Class")
             }
         }
-        val classLoader = URLClassLoader(
+        return KJvmCompiledScriptLazilyLoadedFromClasspath(className, listOf(file))
-            arrayOf(file.toURI().toURL()),
-            scriptCompilationConfiguration[ScriptCompilationConfiguration.hostConfiguration]?.get(ScriptingHostConfiguration.jvm.baseClassLoader)
-        )
-
-        return createScriptFromClassLoader(className, classLoader)
     }
 
     override fun store(
@@ -54,4 +46,40 @@ open class CompiledScriptJarsCache(val scriptToFile: (SourceCode, ScriptCompilat
 
         jvmScript.saveToJar(file)
     }
-}
+}
+
+private class KJvmCompiledScriptLazilyLoadedFromClasspath(
+    private val scriptClassFQName: String,
+    private val classPath: List<File>
+) : CompiledScript {
+
+    private var loadedScript: KJvmCompiledScript? = null
+
+    fun getScriptOrError(): KJvmCompiledScript = loadedScript ?: throw RuntimeException("Compiled script is not loaded yet")
+
+    override suspend fun getClass(scriptEvaluationConfiguration: ScriptEvaluationConfiguration?): ResultWithDiagnostics<KClass<*>> {
+        if (loadedScript == null) {
+            val actualEvaluationConfiguration = scriptEvaluationConfiguration ?: ScriptEvaluationConfiguration()
+            val baseClassLoader = actualEvaluationConfiguration[ScriptEvaluationConfiguration.jvm.baseClassLoader]
+            val classLoader = URLClassLoader(
+                classPath.map { it.toURI().toURL() }.toTypedArray(),
+                baseClassLoader
+            )
+            loadedScript = createScriptFromClassLoader(scriptClassFQName, classLoader)
+        }
+        return getScriptOrError().getClass(scriptEvaluationConfiguration)
+    }
+
+    override val compilationConfiguration: ScriptCompilationConfiguration
+        get() = getScriptOrError().compilationConfiguration
+
+    override val sourceLocationId: String?
+        get() = getScriptOrError().sourceLocationId
+
+    override val otherScripts: List<CompiledScript>
+        get() = getScriptOrError().otherScripts
+
+    override val resultField: Pair<String, KotlinType>?
+        get() = getScriptOrError().resultField
+}
+

libraries/scripting/jvm-host/src/kotlin/script/experimental/jvmhost/jvmScriptSaving.kt

@@ -14,7 +14,6 @@ import java.util.jar.Manifest
 import kotlin.script.experimental.api.*
 import kotlin.script.experimental.jvm.JvmDependency
 import kotlin.script.experimental.jvm.impl.*
-import kotlin.script.experimental.jvm.util.scriptCompilationClasspathFromContext
 import kotlin.script.experimental.jvm.util.scriptCompilationClasspathFromContextOrNull
 
 // TODO: generate execution code (main)

libraries/scripting/jvm/src/kotlin/script/experimental/jvm/BasicJvmScriptEvaluator.kt

@@ -8,7 +8,6 @@ package kotlin.script.experimental.jvm
 import java.lang.reflect.InvocationTargetException
 import kotlin.reflect.KClass
 import kotlin.script.experimental.api.*
-import kotlin.script.experimental.jvm.impl.getConfigurationWithClassloader
 
 open class BasicJvmScriptEvaluator : ScriptEvaluator {
 
@@ -16,26 +15,28 @@ open class BasicJvmScriptEvaluator : ScriptEvaluator {
         compiledScript: CompiledScript,
         scriptEvaluationConfiguration: ScriptEvaluationConfiguration
     ): ResultWithDiagnostics<EvaluationResult> = try {
-        val configuration = getConfigurationWithClassloader(compiledScript, scriptEvaluationConfiguration)
+        compiledScript.getClass(scriptEvaluationConfiguration).onSuccess { scriptClass ->
 
-        compiledScript.getClass(configuration).onSuccess { scriptClass ->
+            // configuration shared between all module scripts
-            // in the future, when (if) we'll stop to compile everything into constructor
+            val sharedConfiguration = scriptEvaluationConfiguration.getOrPrepareShared(scriptClass.java.classLoader)
-            // run as SAM
-            // return res
 
-            val sharedScripts = configuration[ScriptEvaluationConfiguration.jvm.scriptsInstancesSharingMap]
+            val sharedScripts = sharedConfiguration[ScriptEvaluationConfiguration.jvm.scriptsInstancesSharingMap]
 
             sharedScripts?.get(scriptClass)?.asSuccess()
                 ?: compiledScript.otherScripts.mapSuccess {
-                    invoke(it, configuration)
+                    invoke(it, sharedConfiguration)
                 }.onSuccess { importedScriptsEvalResults ->
 
                     val refinedEvalConfiguration =
-                        configuration.refineBeforeEvaluation(compiledScript).valueOr {
+                        sharedConfiguration.refineBeforeEvaluation(compiledScript).valueOr {
                             return@invoke ResultWithDiagnostics.Failure(it.reports)
                         }
 
                     val resultValue = try {
+                        // in the future, when (if) we'll stop to compile everything into constructor
+                        // run as SAM
+                        // return res
+
                         val instance =
                             scriptClass.evalWithConfigAndOtherScriptsResults(refinedEvalConfiguration, importedScriptsEvalResults)
 
@@ -99,3 +100,13 @@ open class BasicJvmScriptEvaluator : ScriptEvaluator {
     }
 }
 
+private fun ScriptEvaluationConfiguration.getOrPrepareShared(classLoader: ClassLoader): ScriptEvaluationConfiguration =
+    if (this[ScriptEvaluationConfiguration.jvm.actualClassLoader] != null)
+        this
+    else
+        with {
+            ScriptEvaluationConfiguration.jvm.actualClassLoader(classLoader)
+            if (this[ScriptEvaluationConfiguration.scriptsInstancesSharing] == true) {
+                ScriptEvaluationConfiguration.jvm.scriptsInstancesSharingMap(mutableMapOf())
+            }
+        }

libraries/scripting/jvm/src/kotlin/script/experimental/jvm/impl/KJvmCompiledScript.kt

@@ -11,7 +11,6 @@ import java.net.URLClassLoader
 import kotlin.reflect.KClass
 import kotlin.script.experimental.api.*
 import kotlin.script.experimental.jvm.*
-import kotlin.script.experimental.jvm.actualClassLoader
 
 internal class KJvmCompiledScriptData(
     var sourceLocationId: String?,
@@ -122,25 +121,6 @@ fun KJvmCompiledScript.getOrCreateActualClassloader(evaluationConfiguration: Scr
         return module.createClassLoader(classLoaderWithDeps)
     }
 
-fun getConfigurationWithClassloader(
-    script: CompiledScript, baseConfiguration: ScriptEvaluationConfiguration
-): ScriptEvaluationConfiguration =
-    if (baseConfiguration.containsKey(ScriptEvaluationConfiguration.jvm.actualClassLoader))
-        baseConfiguration
-    else {
-        val jvmScript = (script as? KJvmCompiledScript)
-            ?: throw IllegalArgumentException("Unexpected compiled script type: $script")
-
-        val classloader = jvmScript.getOrCreateActualClassloader(baseConfiguration)
-
-        ScriptEvaluationConfiguration(baseConfiguration) {
-            ScriptEvaluationConfiguration.jvm.actualClassLoader(classloader)
-            if (baseConfiguration[ScriptEvaluationConfiguration.scriptsInstancesSharing] == true) {
-                ScriptEvaluationConfiguration.jvm.scriptsInstancesSharingMap(mutableMapOf())
-            }
-        }
-    }
-
 private fun CompiledScript.makeClassLoaderFromDependencies(baseClassLoader: ClassLoader?): ClassLoader? {
     val processedScripts = mutableSetOf<CompiledScript>()
     fun recursiveScriptsSeq(res: Sequence<CompiledScript>, script: CompiledScript): Sequence<CompiledScript> =