Immplement default cache in main-kts

libraries/tools/kotlin-main-kts-test/test/org/jetbrains/kotlin/mainKts/test/mainKtsIT.kt

@@ -6,6 +6,8 @@
 package org.jetbrains.kotlin.mainKts.test
 
 import junit.framework.Assert
+import org.jetbrains.kotlin.mainKts.COMPILED_SCRIPTS_CACHE_DIR_ENV_VAR
+import org.jetbrains.kotlin.mainKts.COMPILED_SCRIPTS_CACHE_DIR_PROPERTY
 import org.jetbrains.kotlin.scripting.compiler.plugin.runWithK2JVMCompiler
 import org.jetbrains.kotlin.scripting.compiler.plugin.runWithKotlinc
 import org.junit.Ignore
@@ -49,26 +51,35 @@ class MainKtsIT {
 fun runWithKotlincAndMainKts(
     scriptPath: String,
     expectedOutPatterns: List<String> = emptyList(),
-    expectedExitCode: Int = 0
+    expectedExitCode: Int = 0,
-) = runWithKotlinc(
+    cacheDir: File? = null
-    scriptPath, expectedOutPatterns, expectedExitCode,
+) {
-    classpath = listOf(
+    runWithKotlinc(
-        File("dist/kotlinc/lib/kotlin-main-kts.jar").also {
+        scriptPath, expectedOutPatterns, expectedExitCode,
-            Assert.assertTrue("kotlin-main-kts.jar not found, run dist task: ${it.absolutePath}", it.exists())
+        classpath = listOf(
-        }
+            File("dist/kotlinc/lib/kotlin-main-kts.jar").also {
+                Assert.assertTrue("kotlin-main-kts.jar not found, run dist task: ${it.absolutePath}", it.exists())
+            }
+        ),
+        additionalEnvVars = listOf(COMPILED_SCRIPTS_CACHE_DIR_ENV_VAR to (cacheDir?.absolutePath ?: ""))
     )
-)
+}
 
 fun runWithK2JVMCompilerAndMainKts(
     scriptPath: String,
     expectedOutPatterns: List<String> = emptyList(),
-    expectedExitCode: Int = 0
+    expectedExitCode: Int = 0,
-) = runWithK2JVMCompiler(
+    cacheDir: File? = null
-    scriptPath, expectedOutPatterns, expectedExitCode,
+) {
-    classpath = listOf(
+    withProperty(COMPILED_SCRIPTS_CACHE_DIR_PROPERTY, cacheDir?.absolutePath ?: "") {
-        File("dist/kotlinc/lib/kotlin-main-kts.jar").also {
+        runWithK2JVMCompiler(
-            Assert.assertTrue("kotlin-main-kts.jar not found, run dist task: ${it.absolutePath}", it.exists())
+            scriptPath, expectedOutPatterns, expectedExitCode,
-        }
+            classpath = listOf(
-    )
+                File("dist/kotlinc/lib/kotlin-main-kts.jar").also {
-)
+                    Assert.assertTrue("kotlin-main-kts.jar not found, run dist task: ${it.absolutePath}", it.exists())
+                }
+            )
+        )
+    }
+}
 

libraries/tools/kotlin-main-kts-test/test/org/jetbrains/kotlin/mainKts/test/mainKtsTest.kt

@@ -4,10 +4,12 @@
  */
 package org.jetbrains.kotlin.mainKts.test
 
+import org.jetbrains.kotlin.mainKts.COMPILED_SCRIPTS_CACHE_DIR_PROPERTY
 import org.jetbrains.kotlin.mainKts.MainKtsScript
 import org.junit.Assert
 import org.junit.Test
 import java.io.*
+import java.net.URLClassLoader
 import kotlin.script.experimental.api.*
 import kotlin.script.experimental.host.toScriptSource
 import kotlin.script.experimental.jvm.baseClassLoader
@@ -15,20 +17,21 @@ import kotlin.script.experimental.jvm.jvm
 import kotlin.script.experimental.jvmhost.BasicJvmScriptingHost
 import kotlin.script.experimental.jvmhost.createJvmCompilationConfigurationFromTemplate
 
-fun evalFile(scriptFile: File): ResultWithDiagnostics<EvaluationResult> {
+fun evalFile(scriptFile: File, cacheDir: File? = null): ResultWithDiagnostics<EvaluationResult> =
+    withProperty(COMPILED_SCRIPTS_CACHE_DIR_PROPERTY, cacheDir?.absolutePath ?: "") {
+        val scriptDefinition = createJvmCompilationConfigurationFromTemplate<MainKtsScript>()
 
-    val scriptDefinition = createJvmCompilationConfigurationFromTemplate<MainKtsScript>()
+        val evaluationEnv = ScriptEvaluationConfiguration {
-
+            jvm {
-    val evaluationEnv = ScriptEvaluationConfiguration {
+                baseClassLoader(null)
-        jvm {
+            }
-            baseClassLoader(null)
+            constructorArgs(emptyArray<String>())
+            enableScriptsInstancesSharing()
         }
-        constructorArgs(emptyArray<String>())
+
-        enableScriptsInstancesSharing()
+        BasicJvmScriptingHost().eval(scriptFile.toScriptSource(), scriptDefinition, evaluationEnv)
     }
 
-    return BasicJvmScriptingHost().eval(scriptFile.toScriptSource(), scriptDefinition, evaluationEnv)
-}
 
 const val TEST_DATA_ROOT = "libraries/tools/kotlin-main-kts-test/testData"
 
@@ -77,6 +80,8 @@ class MainKtsTest {
         assertSucceeded(res)
     }
 
+    private val outFromImportTest = listOf("Hi from common", "Hi from middle", "sharedVar == 5")
+
     @Test
     fun testImport() {
 
@@ -85,7 +90,7 @@ class MainKtsTest {
             assertSucceeded(res)
         }.lines()
 
-        Assert.assertEquals(listOf("Hi from common", "Hi from middle", "sharedVar == 5"), out)
+        Assert.assertEquals(outFromImportTest, out)
     }
 
     @Test
@@ -100,6 +105,34 @@ class MainKtsTest {
         Assert.assertEquals(listOf("Hi from sub", "Hi from super", "Hi from random"), out)
     }
 
+    @Test
+    fun testCache() {
+        val script = File("$TEST_DATA_ROOT/import-test.main.kts")
+        val cache = createTempDir("main.kts.test")
+
+        try {
+            Assert.assertTrue(cache.exists() && cache.listFiles { f: File -> f.extension == "jar" }?.isEmpty() == true)
+            val out1 = evalSuccessWithOut(script)
+            Assert.assertEquals(outFromImportTest, out1)
+            Assert.assertTrue(cache.listFiles { f: File -> f.extension.equals("jar", ignoreCase = true) }?.isEmpty() == true)
+
+            val out2 = evalSuccessWithOut(script, cache)
+            Assert.assertEquals(outFromImportTest, out2)
+            val casheFile = cache.listFiles { f: File -> f.extension.equals("jar", ignoreCase = true) }?.firstOrNull()
+            Assert.assertTrue(casheFile != null && casheFile.exists())
+
+            val out3 = captureOut {
+                val classLoader = URLClassLoader(arrayOf(casheFile!!.toURI().toURL()), null)
+                val clazz = classLoader.loadClass("Import_test_main")
+                val mainFn = clazz.getDeclaredMethod("main", Array<String>::class.java)
+                mainFn.invoke(null, arrayOf<String>())
+            }.lines()
+            Assert.assertEquals(outFromImportTest, out3)
+        } finally {
+            cache.deleteRecursively()
+        }
+    }
+
     private fun assertIsJava6Bytecode(res: ResultWithDiagnostics<EvaluationResult>) {
         val scriptClassResource = res.valueOrThrow().returnValue.scriptClass!!.java.run {
             getResource("$simpleName.class")
@@ -129,6 +162,12 @@ class MainKtsTest {
             res is ResultWithDiagnostics.Failure && res.reports.any { it.message.contains("$expectedError") }
         )
     }
+
+    private fun evalSuccessWithOut(scriptFile: File, cacheDir: File? = null): List<String> =
+        captureOut {
+            val res = evalFile(scriptFile, cacheDir)
+            assertSucceeded(res)
+        }.lines()
 }
 
 internal fun captureOut(body: () -> Unit): String {
@@ -143,3 +182,15 @@ internal fun captureOut(body: () -> Unit): String {
     }
     return outStream.toString().trim()
 }
+
+internal fun <T> withProperty(name: String, value: String?, body: () -> T): T {
+    val prevCacheDir = System.getProperty(name)
+    if (value == null) System.clearProperty(name)
+    else System.setProperty(name, value)
+    try {
+        return body()
+    } finally {
+        if (prevCacheDir == null) System.clearProperty(name)
+        else System.setProperty(name, prevCacheDir)
+    }
+}

libraries/tools/kotlin-main-kts/src/org/jetbrains/kotlin/mainKts/scriptDef.kt

@@ -13,6 +13,7 @@ import org.jetbrains.kotlin.script.util.DependsOn
 import org.jetbrains.kotlin.script.util.Import
 import org.jetbrains.kotlin.script.util.Repository
 import java.io.File
+import java.security.MessageDigest
 import kotlin.script.dependencies.ScriptContents
 import kotlin.script.dependencies.ScriptDependenciesResolver
 import kotlin.script.experimental.annotations.KotlinScript
@@ -21,11 +22,11 @@ import kotlin.script.experimental.dependencies.CompoundDependenciesResolver
 import kotlin.script.experimental.dependencies.FileSystemDependenciesResolver
 import kotlin.script.experimental.host.FileBasedScriptSource
 import kotlin.script.experimental.host.FileScriptSource
+import kotlin.script.experimental.host.ScriptingHostConfiguration
+import kotlin.script.experimental.jvm.*
 import kotlin.script.experimental.jvm.compat.mapLegacyDiagnosticSeverity
 import kotlin.script.experimental.jvm.compat.mapLegacyScriptPosition
-import kotlin.script.experimental.jvm.dependenciesFromClassContext
+import kotlin.script.experimental.jvmhost.CompiledScriptJarsCache
-import kotlin.script.experimental.jvm.jvm
-import kotlin.script.experimental.jvm.updateClasspath
 import kotlin.script.experimental.jvmhost.jsr223.configureProvidedPropertiesFromJsr223Context
 import kotlin.script.experimental.jvmhost.jsr223.importAllBindings
 import kotlin.script.experimental.jvmhost.jsr223.jsr223
@@ -38,7 +39,10 @@ import kotlin.script.experimental.jvmhost.jsr223.jsr223
 )
 abstract class MainKtsScript(val args: Array<String>)
 
-object MainKtsScriptDefinition : ScriptCompilationConfiguration(
+const val COMPILED_SCRIPTS_CACHE_DIR_ENV_VAR = "KOTLIN_MAIN_KTS_COMPILED_SCRIPTS_CACHE_DIR"
+const val COMPILED_SCRIPTS_CACHE_DIR_PROPERTY = "kotlin.main.kts.compiled.scripts.cache.dir"
+
+class MainKtsScriptDefinition : ScriptCompilationConfiguration(
     {
         defaultImports(DependsOn::class, Repository::class, Import::class, CompilerOptions::class)
         jvm {
@@ -54,15 +58,45 @@ object MainKtsScriptDefinition : ScriptCompilationConfiguration(
         jsr223 {
             importAllBindings(true)
         }
+        hostConfiguration(ScriptingHostConfiguration {
+            jvm {
+                val cacheExtSetting = System.getProperty(COMPILED_SCRIPTS_CACHE_DIR_PROPERTY)
+                    ?: System.getenv(COMPILED_SCRIPTS_CACHE_DIR_ENV_VAR)
+                val cacheBaseDir = when {
+                    cacheExtSetting == null -> System.getProperty("java.io.tmpdir")
+                        ?.let(::File)?.takeIf { it.exists() && it.isDirectory }
+                        ?.let { File(it, "main.kts.compiled.cache").apply { mkdir() } }
+                    cacheExtSetting.isBlank() -> null
+                    else -> File(cacheExtSetting)
+                }?.takeIf { it.exists() && it.isDirectory }
+                if (cacheBaseDir != null)
+                    compilationCache(
+                        CompiledScriptJarsCache { script, scriptCompilationConfiguration ->
+                            File(cacheBaseDir, compiledScriptUniqueName(script, scriptCompilationConfiguration) + ".jar")
+                        }
+                    )
+            }
+        })
     })
 
 object MainKtsEvaluationConfiguration : ScriptEvaluationConfiguration(
     {
         scriptsInstancesSharing(true)
         refineConfigurationBeforeEvaluate(::configureProvidedPropertiesFromJsr223Context)
+        refineConfigurationBeforeEvaluate(::configureConstructorArgsFromMainArgs)
     }
 )
 
+fun configureConstructorArgsFromMainArgs(context: ScriptEvaluationConfigurationRefinementContext): ResultWithDiagnostics<ScriptEvaluationConfiguration> {
+    val mainArgs = context.evaluationConfiguration[ScriptEvaluationConfiguration.jvm.mainArguments]
+    val res = if (context.evaluationConfiguration[ScriptEvaluationConfiguration.constructorArgs] == null && mainArgs != null) {
+        context.evaluationConfiguration.with {
+            constructorArgs(mainArgs)
+        }
+    } else context.evaluationConfiguration
+    return res.asSuccess()
+}
+
 class MainKtsConfigurator : RefineScriptCompilationConfigurationHandler {
     private val resolver = CompoundDependenciesResolver(FileSystemDependenciesResolver(), IvyResolver())
 
@@ -114,3 +148,16 @@ class MainKtsConfigurator : RefineScriptCompilationConfigurationHandler {
     }
 }
 
+private fun compiledScriptUniqueName(script: SourceCode, scriptCompilationConfiguration: ScriptCompilationConfiguration): String {
+    val digestWrapper = MessageDigest.getInstance("MD5")
+    digestWrapper.update(script.text.toByteArray())
+    scriptCompilationConfiguration.notTransientData.entries
+        .sortedBy { it.key.name }
+        .forEach {
+            digestWrapper.update(it.key.name.toByteArray())
+            digestWrapper.update(it.value.toString().toByteArray())
+        }
+    return digestWrapper.digest().toHexString()
+}
+
+private fun ByteArray.toHexString(): String = joinToString("", transform = { "%02x".format(it) })

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

@@ -24,7 +24,8 @@ fun runWithKotlinc(
     expectedOutPatterns: List<String> = emptyList(),
     expectedExitCode: Int = 0,
     workDirectory: File? = null,
-    classpath: List<File> = emptyList()
+    classpath: List<File> = emptyList(),
+    additionalEnvVars: Iterable<Pair<String, String>>? = null
 ) {
     val executableName = "kotlinc"
     // TODO:
@@ -45,6 +46,9 @@ fun runWithKotlinc(
     if (workDirectory != null) {
         processBuilder.directory(workDirectory)
     }
+    if (additionalEnvVars != null) {
+        processBuilder.environment().putAll(additionalEnvVars)
+    }
     val process = processBuilder.start()
 
     data class ExceptionContainer(