Implement compiled script caching with tests

compiler/backend/src/org/jetbrains/kotlin/codegen/BytesUrlUtils.java

@@ -0,0 +1,47 @@
+/*
+ * 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.codegen;
+
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+import java.io.ByteArrayInputStream;
+import java.io.InputStream;
+import java.net.MalformedURLException;
+import java.net.URL;
+import java.net.URLConnection;
+import java.net.URLStreamHandler;
+import java.util.Base64;
+
+public class BytesUrlUtils {
+
+    private static final URLStreamHandler BYTES_URL_HANDLER = new URLStreamHandler() {
+        @Override
+        protected URLConnection openConnection(URL url) {
+            return new URLConnection(url) {
+                @Override
+                public void connect() {
+                }
+
+                @Override
+                public InputStream getInputStream() {
+                    return new ByteArrayInputStream(Base64.getDecoder().decode(url.getPath()));
+                }
+            };
+        }
+    };
+
+    /**
+     * Encode the entire [bytes] array in the self-contained URL with "bytes" protocol
+     * @param bytes byte array to encode into the URL
+     * @return the URL containing encoded [bytes] contents
+     * @throws MalformedURLException
+     */
+    @Nullable
+    public static URL createBytesUrl(@NotNull byte[] bytes) throws MalformedURLException {
+        return new URL(null, "bytes:" + Base64.getEncoder().encodeToString(bytes), BYTES_URL_HANDLER);
+    }
+}

compiler/backend/src/org/jetbrains/kotlin/codegen/GeneratedClassLoader.java

@@ -27,30 +27,12 @@ import java.io.IOException;
 import java.io.InputStream;
 import java.net.URL;
 import java.net.URLClassLoader;
-import java.net.URLConnection;
-import java.net.URLStreamHandler;
-import java.util.Base64;
 import java.util.Collections;
 import java.util.Enumeration;
 import java.util.List;
 import java.util.jar.Manifest;
 
 public class GeneratedClassLoader extends URLClassLoader {
-    private static final URLStreamHandler FAKE_BASE64_URL_HANDLER = new URLStreamHandler() {
-        @Override
-        protected URLConnection openConnection(URL url) {
-            return new URLConnection(url) {
-                @Override
-                public void connect() {
-                }
-
-                @Override
-                public InputStream getInputStream() {
-                    return new ByteArrayInputStream(Base64.getDecoder().decode(url.getPath()));
-                }
-            };
-        }
-    };
 
     private ClassFileFactory factory;
 
@@ -92,10 +74,9 @@ public class GeneratedClassLoader extends URLClassLoader {
     private URL createFakeURLForResource(@NotNull String name) {
         try {
             OutputFile outputFile = factory.get(name);
-            // Encode the byte array in the URL path to prevent creating unneeded temporary files
             return outputFile == null
                    ? null
-                   : new URL(null, "bytes:" + Base64.getEncoder().encodeToString(outputFile.asByteArray()), FAKE_BASE64_URL_HANDLER);
+                   : BytesUrlUtils.createBytesUrl(outputFile.asByteArray());
         } catch (IOException e) {
             throw ExceptionUtilsKt.rethrow(e);
         }

libraries/scripting/common/src/kotlin/script/experimental/util/propertiesCollection.kt

@@ -31,6 +31,8 @@ open class PropertiesCollection(private val properties: Map<Key<*>, Any> = empty
     fun <T> getNoDefault(key: PropertiesCollection.Key<T>): T? =
         properties[key]?.let { it as T }
 
+    fun entries(): Set<Map.Entry<Key<*>, Any>> = properties.entries
+
     companion object {
         fun <T> key(defaultValue: T? = null) = PropertyKeyDelegate(defaultValue)
         fun <T> keyCopy(source: Key<T>) = PropertyKeyCopyDelegate(source)

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

@@ -0,0 +1,35 @@
+/*
+ * 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 kotlin.script.experimental.jvmhost.impl
+
+import org.jetbrains.kotlin.codegen.BytesUrlUtils
+import java.io.ByteArrayInputStream
+import java.io.InputStream
+import java.net.URL
+import java.util.*
+
+internal class CompiledScriptClassLoader(parent: ClassLoader, private val entries: Map<String, ByteArray>) : ClassLoader(parent) {
+
+    override fun findClass(name: String): Class<*>? {
+        val classPathName = name.replace('.', '/') + ".class"
+        val classBytes = entries[classPathName] ?: return null
+        return defineClass(name, classBytes, 0, classBytes.size)
+    }
+
+    override fun getResourceAsStream(name: String): InputStream? =
+        entries[name]?.let(::ByteArrayInputStream) ?: super.getResourceAsStream(name)
+
+    override fun findResources(name: String?): Enumeration<URL>? {
+        val fromParent = super.findResources(name)
+
+        val url = entries[name]?.let { BytesUrlUtils.createBytesUrl(it) } ?: return fromParent
+
+        return Collections.enumeration(listOf(url) + fromParent.asSequence())
+    }
+
+    override fun findResource(name: String?): URL? =
+        entries[name]?.let { BytesUrlUtils.createBytesUrl(it) } ?: super.findResource(name)
+}

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

@@ -0,0 +1,81 @@
+/*
+ * 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 kotlin.script.experimental.jvmhost.impl
+
+import org.jetbrains.kotlin.codegen.state.GenerationState
+import java.io.ObjectInputStream
+import java.io.ObjectOutputStream
+import java.io.Serializable
+import java.net.URLClassLoader
+import kotlin.reflect.KClass
+import kotlin.script.experimental.api.*
+import kotlin.script.experimental.jvm.JvmDependency
+import kotlin.script.experimental.jvmhost.JvmScriptEvaluationConfiguration
+import kotlin.script.experimental.jvmhost.baseClassLoader
+
+class KJvmCompiledScript<out ScriptBase : Any>(
+    compilationConfiguration: ScriptCompilationConfiguration,
+    generationState: GenerationState,
+    private var scriptClassFQName: String
+) : CompiledScript<ScriptBase>, Serializable {
+
+    private var _compilationConfiguration: ScriptCompilationConfiguration? = compilationConfiguration
+    private var compilerOutputFiles: Map<String, ByteArray> = run {
+        val res = sortedMapOf<String, ByteArray>()
+        for (it in generationState.factory.asList()) {
+            res[it.relativePath] = it.asByteArray()
+        }
+        res
+    }
+
+    override val compilationConfiguration: ScriptCompilationConfiguration
+        get() = _compilationConfiguration!!
+
+    override suspend fun getClass(scriptEvaluationConfiguration: ScriptEvaluationConfiguration?): ResultWithDiagnostics<KClass<*>> = try {
+        val baseClassLoader = scriptEvaluationConfiguration?.get(JvmScriptEvaluationConfiguration.baseClassLoader)
+            ?: Thread.currentThread().contextClassLoader
+        val dependencies = compilationConfiguration[ScriptCompilationConfiguration.dependencies]
+            ?.flatMap { (it as? JvmDependency)?.classpath?.map { it.toURI().toURL() } ?: emptyList() }
+        // TODO: previous dependencies and classloaders should be taken into account here
+        val classLoaderWithDeps =
+            if (dependencies == null) baseClassLoader
+            else URLClassLoader(dependencies.toTypedArray(), baseClassLoader)
+        val classLoader = CompiledScriptClassLoader(classLoaderWithDeps, compilerOutputFiles)
+
+        val clazz = classLoader.loadClass(scriptClassFQName).kotlin
+        clazz.asSuccess()
+    } catch (e: Throwable) {
+        ResultWithDiagnostics.Failure(
+            ScriptDiagnostic(
+                "Unable to instantiate class $scriptClassFQName",
+                exception = e
+            )
+        )
+    }
+
+    // This method is exposed because the compilation configuration is not generally serializable (yet), but since it is supposed to
+    // be deserialized only from the cache, the configuration could be assigned from the cache.load method
+    fun setCompilationConfiguration(configuration: ScriptCompilationConfiguration) {
+        if (_compilationConfiguration != null) throw IllegalStateException("This method is applicable only in deserialization context")
+        _compilationConfiguration = configuration
+    }
+
+    private fun writeObject(outputStream: ObjectOutputStream) {
+        outputStream.writeObject(compilerOutputFiles)
+        outputStream.writeObject(scriptClassFQName)
+    }
+
+    private fun readObject(inputStream: ObjectInputStream) {
+        _compilationConfiguration = null
+        compilerOutputFiles = inputStream.readObject() as Map<String, ByteArray>
+        scriptClassFQName = inputStream.readObject() as String
+    }
+
+    companion object {
+        @JvmStatic
+        private val serialVersionUID = 0L
+    }
+}

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

@@ -25,7 +25,6 @@ import org.jetbrains.kotlin.cli.jvm.config.JvmModulePathRoot
 import org.jetbrains.kotlin.cli.jvm.config.addJvmClasspathRoots
 import org.jetbrains.kotlin.cli.jvm.modules.CoreJrtFileSystem
 import org.jetbrains.kotlin.codegen.ClassBuilderFactories
-import org.jetbrains.kotlin.codegen.GeneratedClassLoader
 import org.jetbrains.kotlin.codegen.KotlinCodegenFacade
 import org.jetbrains.kotlin.codegen.state.GenerationState
 import org.jetbrains.kotlin.config.*
@@ -35,8 +34,6 @@ import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.script.KotlinScriptDefinition
 import org.jetbrains.kotlin.script.util.KotlinJars
 import java.io.File
-import java.net.URLClassLoader
-import kotlin.reflect.KClass
 import kotlin.script.experimental.api.*
 import kotlin.script.experimental.dependencies.DependenciesResolver
 import kotlin.script.experimental.host.ScriptingHostConfiguration
@@ -46,35 +43,9 @@ import kotlin.script.experimental.jvm.JvmDependency
 import kotlin.script.experimental.jvm.impl.BridgeDependenciesResolver
 import kotlin.script.experimental.jvm.javaHome
 import kotlin.script.experimental.jvm.jvm
-import kotlin.script.experimental.jvmhost.JvmScriptEvaluationConfiguration
 import kotlin.script.experimental.jvmhost.KJvmCompilerProxy
-import kotlin.script.experimental.jvmhost.baseClassLoader
 import kotlin.script.experimental.util.getOrError
 
-class KJvmCompiledScript<out ScriptBase : Any>(
-    override val compilationConfiguration: ScriptCompilationConfiguration,
-    private val generationState: GenerationState,
-    private val scriptClassFQName: String
-) : CompiledScript<ScriptBase> {
-
-    override suspend fun getClass(scriptEvaluationConfiguration: ScriptEvaluationConfiguration?): ResultWithDiagnostics<KClass<*>> = try {
-        val baseClassLoader = scriptEvaluationConfiguration?.get(JvmScriptEvaluationConfiguration.baseClassLoader)
-            ?: Thread.currentThread().contextClassLoader
-        val dependencies = compilationConfiguration[ScriptCompilationConfiguration.dependencies]
-            ?.flatMap { (it as? JvmDependency)?.classpath?.map { it.toURI().toURL() } ?: emptyList() }
-        // TODO: previous dependencies and classloaders should be taken into account here
-        val classLoaderWithDeps =
-            if (dependencies == null) baseClassLoader
-            else URLClassLoader(dependencies.toTypedArray(), baseClassLoader)
-        val classLoader = GeneratedClassLoader(generationState.factory, classLoaderWithDeps)
-
-        val clazz = classLoader.loadClass(scriptClassFQName).kotlin
-        clazz.asSuccess()
-    } catch (e: Throwable) {
-        ResultWithDiagnostics.Failure(ScriptDiagnostic("Unable to instantiate class $scriptClassFQName", exception = e))
-    }
-}
-
 class KJvmCompilerImpl(val hostConfiguration: ScriptingHostConfiguration) : KJvmCompilerProxy {
 
     override fun compile(
@@ -204,7 +175,7 @@ class KJvmCompilerImpl(val hostConfiguration: ScriptingHostConfiguration) : KJvm
     }
 }
 
-class ScriptDiagnosticsMessageCollector : MessageCollector {
+internal class ScriptDiagnosticsMessageCollector : MessageCollector {
 
     private val _diagnostics = arrayListOf<ScriptDiagnostic>()
 

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

@@ -44,7 +44,7 @@ open class BasicJvmScriptEvaluator : ScriptEvaluator {
                     scriptEvaluationConfiguration?.get(ScriptEvaluationConfiguration.constructorArgs)?.let {
                         args.addAll(it)
                     }
-                    val ctor = scriptClass.java.constructors.first()
+                    val ctor = scriptClass.java.constructors.single()
                     val instance = ctor.newInstance(*args.toArray())
 
                     // TODO: fix result value

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

@@ -5,14 +5,21 @@
 
 package kotlin.script.experimental.jvmhost.test
 
+import kotlinx.coroutines.experimental.runBlocking
+import org.jetbrains.kotlin.daemon.common.toHexString
 import org.junit.Assert
 import org.junit.Test
-import java.io.ByteArrayOutputStream
+import java.io.*
-import java.io.PrintStream
+import java.nio.file.Files
-import kotlin.script.experimental.api.ResultWithDiagnostics
+import java.security.MessageDigest
+import kotlin.reflect.KClass
+import kotlin.script.experimental.api.*
+import kotlin.script.experimental.host.BasicScriptingHost
 import kotlin.script.experimental.host.toScriptSource
-import kotlin.script.experimental.jvmhost.BasicJvmScriptingHost
+import kotlin.script.experimental.jvm.defaultJvmScriptingHostConfiguration
-import kotlin.script.experimental.jvmhost.createCompilationConfigurationFromTemplate
+import kotlin.script.experimental.jvmhost.*
+import kotlin.script.experimental.jvmhost.impl.CompiledScriptClassLoader
+import kotlin.script.experimental.jvmhost.impl.KJvmCompiledScript
 import kotlin.script.templates.standard.SimpleScriptTemplate
 
 class ScriptingHostTest {
@@ -22,20 +29,188 @@ class ScriptingHostTest {
         val greeting = "Hello from script!"
         val output = captureOut {
             evalScript("println(\"$greeting\")")
-        }.trim()
+        }
         Assert.assertEquals(greeting, output)
     }
-}
 
-internal fun evalScript(script: String) {
+    @Test
-    val compilationConfiguration = createCompilationConfigurationFromTemplate<SimpleScriptTemplate>()
+    fun testMemoryCache() {
-    val res = BasicJvmScriptingHost().eval(script.toScriptSource(), compilationConfiguration, null)
+        val script = "val x = 1\nprintln(\"x = \$x\")"
-    if (res is ResultWithDiagnostics.Failure) {
+        val cache = SimpleMemoryScriptsCache()
-        throw Exception("Compilation/evaluation failed:\n  ${res.reports.joinToString("\n  ") { it.exception?.toString() ?: it.message }}")
+        val compiler = JvmScriptCompiler(defaultJvmScriptingHostConfiguration, cache = cache)
+        val evaluator = BasicJvmScriptEvaluator()
+        val host = BasicJvmScriptingHost(compiler = compiler, evaluator = evaluator)
+        Assert.assertTrue(cache.data.isEmpty())
+
+        val output = captureOut { evalScript(script, host) }
+        Assert.assertEquals("x = 1", output)
+
+        Assert.assertEquals(1, cache.data.size)
+        val compiled = cache.data.values.first()
+
+        val output2 = captureOut { runBlocking { evaluator(compiled, null) } }
+        Assert.assertEquals(output, output2)
+
+        // TODO: check if cached script is actually used
+        val output3 = captureOut { evalScript(script, host) }.trim()
+        Assert.assertEquals(output, output3)
+    }
+
+    @Test
+    fun testFileCache() {
+        val script = "val x = 1\nprintln(\"x = \$x\")"
+        val cacheDir = Files.createTempDirectory("scriptingTestCache").toFile()
+        try {
+            val cache = FileBasedScriptCache(cacheDir)
+            val compiler = JvmScriptCompiler(defaultJvmScriptingHostConfiguration, cache = cache)
+            val evaluator = BasicJvmScriptEvaluator()
+            val host = BasicJvmScriptingHost(compiler = compiler, evaluator = evaluator)
+            Assert.assertTrue(cache.baseDir.listFiles().isEmpty())
+
+            val scriptCompilationConfiguration = createCompilationConfigurationFromTemplate<SimpleScriptTemplate>()
+
+            var compiledScript: CompiledScript<*>? = null
+            val output = captureOut {
+                runBlocking {
+                    compiler(script.toScriptSource(), scriptCompilationConfiguration).onSuccess {
+                        compiledScript = it
+                        evaluator(it, null)
+                    }.throwOnFailure()
+                }
+            }
+            Assert.assertEquals("x = 1", output)
+
+            val cachedCompiledScript = cache.baseDir.listFiles().let { files ->
+                Assert.assertEquals(1, files.size)
+                files.first().readCompiledScript(scriptCompilationConfiguration)
+            }
+
+            runBlocking {
+                Assert.assertEquals(
+                    compiledScript!!.getClass(null).resultOrNull()?.qualifiedName,
+                    cachedCompiledScript.getClass(null).resultOrNull()?.qualifiedName
+                )
+            }
+
+            val output2 = captureOut {
+                runBlocking {
+                    evaluator(cachedCompiledScript, null).throwOnFailure()
+                }
+            }
+            Assert.assertEquals(output, output2)
+
+            // TODO: check if cached script is actually used
+            val output3 = captureOut { evalScript(script, host) }.trim()
+            Assert.assertEquals(output, output3)
+        } finally {
+            cacheDir.deleteRecursively()
+        }
+    }
+
+    @Test
+    fun testCompiledScriptClassLoader() {
+        val script = "val x = 1"
+        val scriptCompilationConfiguration = createCompilationConfigurationFromTemplate<SimpleScriptTemplate>()
+        val compiler = JvmScriptCompiler(defaultJvmScriptingHostConfiguration)
+        val compiledScript = runBlocking {
+            val res = compiler(script.toScriptSource(), scriptCompilationConfiguration).throwOnFailure()
+            (res as ResultWithDiagnostics.Success<CompiledScript<*>>).value
+        }
+        val compiledScriptClass = runBlocking { compiledScript.getClass(null).throwOnFailure().resultOrNull()!! as KClass<*> }
+        val classLoader = compiledScriptClass.java.classLoader
+
+        Assert.assertTrue(classLoader is CompiledScriptClassLoader)
+        val anotherClass = classLoader.loadClass(compiledScriptClass.qualifiedName)
+
+        Assert.assertEquals(compiledScriptClass.java, anotherClass)
+
+        val classResourceName = compiledScriptClass.qualifiedName!!.replace('.', '/') + ".class"
+        val classAsResourceUrl = classLoader.getResource(classResourceName)
+        val classAssResourceStream = classLoader.getResourceAsStream(classResourceName)
+
+        Assert.assertNotNull(classAsResourceUrl)
+        Assert.assertNotNull(classAssResourceStream)
+
+        val classAsResourceData = classAsResourceUrl.openConnection().getInputStream().readBytes()
+        val classAsResourceStreamData = classAssResourceStream.readBytes()
+
+        Assert.assertArrayEquals(classAsResourceData, classAsResourceStreamData)
+
+        // TODO: consider testing getResources as well
     }
 }
 
-internal fun captureOut(body: () -> Unit): String {
+fun ResultWithDiagnostics<*>.throwOnFailure(): ResultWithDiagnostics<*> = apply {
+    if (this is ResultWithDiagnostics.Failure) {
+        throw Exception("Compilation/evaluation failed:\n  ${reports.joinToString("\n  ") { it.exception?.toString() ?: it.message }}")
+    }
+}
+
+private fun evalScript(script: String, host: BasicScriptingHost = BasicJvmScriptingHost()) {
+    val compilationConfiguration = createCompilationConfigurationFromTemplate<SimpleScriptTemplate>()
+    host.eval(script.toScriptSource(), compilationConfiguration, null).throwOnFailure()
+}
+
+
+private class SimpleMemoryScriptsCache : CompiledJvmScriptsCache {
+
+    internal val data = hashMapOf<Pair<SourceCode, ScriptCompilationConfiguration>, CompiledScript<*>>()
+
+    override fun get(script: SourceCode, scriptCompilationConfiguration: ScriptCompilationConfiguration): CompiledScript<*>? =
+        data[script to scriptCompilationConfiguration]
+
+    override fun store(
+        compiledScript: CompiledScript<*>,
+        script: SourceCode,
+        scriptCompilationConfiguration: ScriptCompilationConfiguration
+    ) {
+        data[script to scriptCompilationConfiguration] = compiledScript
+    }
+}
+
+private fun File.readCompiledScript(scriptCompilationConfiguration: ScriptCompilationConfiguration): CompiledScript<*> {
+    return inputStream().use { fs ->
+        ObjectInputStream(fs).use { os ->
+            (os.readObject() as KJvmCompiledScript<*>).apply {
+                setCompilationConfiguration(scriptCompilationConfiguration)
+            }
+        }
+    }
+}
+
+
+private class FileBasedScriptCache(val baseDir: File) : CompiledJvmScriptsCache {
+
+    internal fun uniqueHash(script: SourceCode, scriptCompilationConfiguration: ScriptCompilationConfiguration): String {
+        val digestWrapper = MessageDigest.getInstance("MD5")
+        digestWrapper.update(script.text.toByteArray())
+        scriptCompilationConfiguration.entries().sortedBy { it.key.name }.forEach {
+            digestWrapper.update(it.key.name.toByteArray())
+            digestWrapper.update(it.value.toString().toByteArray())
+        }
+        return digestWrapper.digest().toHexString()
+    }
+
+    override fun get(script: SourceCode, scriptCompilationConfiguration: ScriptCompilationConfiguration): CompiledScript<*>? {
+        val file = File(baseDir, uniqueHash(script, scriptCompilationConfiguration))
+        return if (!file.exists()) null else file.readCompiledScript(scriptCompilationConfiguration)
+    }
+
+    override fun store(
+        compiledScript: CompiledScript<*>,
+        script: SourceCode,
+        scriptCompilationConfiguration: ScriptCompilationConfiguration
+    ) {
+        val file = File(baseDir, uniqueHash(script, scriptCompilationConfiguration))
+        file.outputStream().use { fs ->
+            ObjectOutputStream(fs).use { os ->
+                os.writeObject(compiledScript)
+            }
+        }
+    }
+}
+
+private fun captureOut(body: () -> Unit): String {
     val outStream = ByteArrayOutputStream()
     val prevOut = System.out
     System.setOut(PrintStream(outStream))
@@ -45,5 +220,5 @@ internal fun captureOut(body: () -> Unit): String {
         System.out.flush()
         System.setOut(prevOut)
     }
-    return outStream.toString()
+    return outStream.toString().trim()
 }