Refactor script-util context tools to simplify its usage in the new scripting

idea/idea-repl/src/org/jetbrains/kotlin/jsr223/KotlinJsr223StandardScriptEngineFactory4Idea.kt

@@ -18,7 +18,7 @@ package org.jetbrains.kotlin.jsr223
 
 import org.jetbrains.kotlin.cli.common.repl.KotlinJsr223JvmScriptEngineFactoryBase
 import org.jetbrains.kotlin.cli.common.repl.ScriptArgsWithTypes
-import org.jetbrains.kotlin.script.util.scriptCompilationClasspathFromContext
+import org.jetbrains.kotlin.script.util.scriptCompilationClasspathFromContextOrStlib
 import org.jetbrains.kotlin.utils.PathUtil
 import javax.script.ScriptContext
 import javax.script.ScriptEngine
@@ -28,11 +28,11 @@ class KotlinJsr223StandardScriptEngineFactory4Idea : KotlinJsr223JvmScriptEngine
 
     override fun getScriptEngine(): ScriptEngine =
             KotlinJsr223JvmScriptEngine4Idea(
-                    this,
+                this,
-                    scriptCompilationClasspathFromContext(PathUtil.KOTLIN_JAVA_SCRIPT_RUNTIME_JAR),
+                scriptCompilationClasspathFromContextOrStlib(PathUtil.KOTLIN_JAVA_SCRIPT_RUNTIME_JAR, wholeClasspath = true),
-                    "kotlin.script.templates.standard.ScriptTemplateWithBindings",
+                "kotlin.script.templates.standard.ScriptTemplateWithBindings",
-                    { ctx, argTypes -> ScriptArgsWithTypes(arrayOf(ctx.getBindings(ScriptContext.ENGINE_SCOPE)), argTypes ?: emptyArray()) },
+                { ctx, argTypes -> ScriptArgsWithTypes(arrayOf(ctx.getBindings(ScriptContext.ENGINE_SCOPE)), argTypes ?: emptyArray()) },
-                    arrayOf(Map::class)
+                arrayOf(Map::class)
             )
 }
 

libraries/tools/kotlin-script-util/src/main/kotlin/org/jetbrains/kotlin/script/jsr223/KotlinJsr223ScriptEngineFactoryExamples.kt

@@ -29,11 +29,11 @@ class KotlinJsr223JvmLocalScriptEngineFactory : KotlinJsr223JvmScriptEngineFacto
 
     override fun getScriptEngine(): ScriptEngine =
             KotlinJsr223JvmLocalScriptEngine(
-                    this,
+                this,
-                    scriptCompilationClasspathFromContext("kotlin-script-util.jar"),
+                scriptCompilationClasspathFromContextOrStlib("kotlin-script-util.jar", wholeClasspath = true),
-                    KotlinStandardJsr223ScriptTemplate::class.qualifiedName!!,
+                KotlinStandardJsr223ScriptTemplate::class.qualifiedName!!,
-                    { ctx, types -> ScriptArgsWithTypes(arrayOf(ctx.getBindings(ScriptContext.ENGINE_SCOPE)), types ?: emptyArray()) },
+                { ctx, types -> ScriptArgsWithTypes(arrayOf(ctx.getBindings(ScriptContext.ENGINE_SCOPE)), types ?: emptyArray()) },
-                    arrayOf(Bindings::class)
+                arrayOf(Bindings::class)
             )
 }
 
@@ -41,12 +41,12 @@ class KotlinJsr223JvmDaemonLocalEvalScriptEngineFactory : KotlinJsr223JvmScriptE
 
     override fun getScriptEngine(): ScriptEngine =
             KotlinJsr223JvmDaemonCompileScriptEngine(
-                    this,
+                this,
-                    KotlinJars.compilerClasspath,
+                KotlinJars.compilerClasspath,
-                    scriptCompilationClasspathFromContext("kotlin-script-util.jar"),
+                scriptCompilationClasspathFromContextOrStlib("kotlin-script-util.jar", wholeClasspath = true),
-                    KotlinStandardJsr223ScriptTemplate::class.qualifiedName!!,
+                KotlinStandardJsr223ScriptTemplate::class.qualifiedName!!,
-                    { ctx, types -> ScriptArgsWithTypes(arrayOf(ctx.getBindings(ScriptContext.ENGINE_SCOPE)), types ?: emptyArray()) },
+                { ctx, types -> ScriptArgsWithTypes(arrayOf(ctx.getBindings(ScriptContext.ENGINE_SCOPE)), types ?: emptyArray()) },
-                    arrayOf(Bindings::class)
+                arrayOf(Bindings::class)
             )
 }
 

libraries/tools/kotlin-script-util/src/main/kotlin/org/jetbrains/kotlin/script/util/context.kt

@@ -65,16 +65,51 @@ internal fun List<File>.takeIfContainsAll(vararg keyNames: String): List<File>?
             keyNames.all { key -> classpath.any { it.matchMaybeVersionedFile(key) } }
         }
 
+internal fun List<File>.filterIfContainsAll(vararg keyNames: String): List<File>? {
+    val res = hashMapOf<String, File>()
+    for (jar in this) {
+        for (prefix in keyNames) {
+            if (jar.matchMaybeVersionedFile(prefix)) {
+                res[prefix] = jar
+                break
+            }
+        }
+    }
+    return if (keyNames.all { res.containsKey(it) }) res.values.toList()
+    else null
+}
+
 internal fun List<File>.takeIfContainsAny(vararg keyNames: String): List<File>? =
         takeIf { classpath ->
             keyNames.any { key -> classpath.any { it.matchMaybeVersionedFile(key) } }
         }
 
-fun scriptCompilationClasspathFromContext(vararg keyNames: String, classLoader: ClassLoader = Thread.currentThread().contextClassLoader): List<File> =
+fun scriptCompilationClasspathFromContextOrNull(
-        System.getProperty(KOTLIN_SCRIPT_CLASSPATH_PROPERTY)?.split(File.pathSeparator)?.map(::File)
+    vararg keyNames: String,
-        ?: classpathFromClassloader(classLoader)?.takeIfContainsAll(*keyNames)
+    classLoader: ClassLoader = Thread.currentThread().contextClassLoader,
-        ?: classpathFromClasspathProperty()?.takeIfContainsAll(*keyNames)
+    wholeClasspath: Boolean = false
-        ?: KotlinJars.kotlinScriptStandardJars
+): List<File>? {
+    fun List<File>.takeAndFilter() = if (wholeClasspath) takeIfContainsAll(*keyNames) else filterIfContainsAll(*keyNames)
+    return System.getProperty(KOTLIN_SCRIPT_CLASSPATH_PROPERTY)?.split(File.pathSeparator)?.map(::File)
+            ?: classpathFromClassloader(classLoader)?.takeAndFilter()
+            ?: classpathFromClasspathProperty()?.takeAndFilter()
+}
+
+fun scriptCompilationClasspathFromContextOrStlib(
+    vararg keyNames: String,
+    classLoader: ClassLoader = Thread.currentThread().contextClassLoader,
+    wholeClasspath: Boolean = false
+): List<File> =
+    scriptCompilationClasspathFromContextOrNull(*keyNames, classLoader = classLoader, wholeClasspath = wholeClasspath)
+            ?: KotlinJars.kotlinScriptStandardJars
+
+fun scriptCompilationClasspathFromContext(
+    vararg keyNames: String,
+    classLoader: ClassLoader = Thread.currentThread().contextClassLoader,
+    wholeClasspath: Boolean = false
+): List<File> =
+    scriptCompilationClasspathFromContextOrNull(*keyNames, classLoader = classLoader, wholeClasspath = wholeClasspath)
+            ?: throw Exception("Unable to get script compilation classpath from context, please specify explicit classpath via \"$KOTLIN_SCRIPT_CLASSPATH_PROPERTY\" property")
 
 object KotlinJars {
 
@@ -101,19 +136,29 @@ object KotlinJars {
         classpath!!
     }
 
-    private fun getLib(propertyName: String, jarName: String, markerClass: KClass<*>): File? =
+    fun getLib(propertyName: String, jarName: String, markerClass: KClass<*>): File? =
             System.getProperty(propertyName)?.let(::File)?.takeIf(File::exists)
             ?: explicitCompilerClasspath?.firstOrNull { it.matchMaybeVersionedFile(jarName) }?.takeIf(File::exists)
             ?: getResourcePathForClass(markerClass.java).takeIf(File::exists)
 
-    val stdlib: File? by lazy {
+    val stdlibOrNull: File? by lazy {
         System.getProperty(KOTLIN_STDLIB_JAR_PROPERTY)?.let(::File)?.takeIf(File::exists)
-        ?: getLib(KOTLIN_RUNTIME_JAR_PROPERTY, KOTLIN_JAVA_STDLIB_JAR, JvmStatic::class)
+                ?: getLib(KOTLIN_RUNTIME_JAR_PROPERTY, KOTLIN_JAVA_STDLIB_JAR, JvmStatic::class)
     }
 
-    val scriptRuntime: File? by lazy {
+    val stdlib: File by lazy {
+        stdlibOrNull
+                ?: throw Exception("Unable to find kotlin stdlib, please specify it explicitly via \"$KOTLIN_STDLIB_JAR_PROPERTY\" property")
+    }
+
+    val scriptRuntimeOrNull: File? by lazy {
         getLib(KOTLIN_SCRIPT_RUNTIME_JAR_PROPERTY, KOTLIN_JAVA_SCRIPT_RUNTIME_JAR, ScriptTemplateWithArgs::class)
     }
 
-    val kotlinScriptStandardJars get() = listOf(stdlib, scriptRuntime).filterNotNull()
+    val scriptRuntime: File by lazy {
+        scriptRuntimeOrNull
+                ?: throw Exception("Unable to find kotlin script runtime, please specify it explicitly via \"$KOTLIN_SCRIPT_RUNTIME_JAR_PROPERTY\" property")
+    }
+
+    val kotlinScriptStandardJars get() = listOf(stdlibOrNull, scriptRuntimeOrNull).filterNotNull()
 }

libraries/tools/kotlin-script-util/src/test/kotlin/org/jetbrains/kotlin/script/util/ScriptUtilContextTests.kt

@@ -24,7 +24,7 @@ class ScriptUtilContextTests {
     @Test
     fun testExplicitScriptClasspath() {
         withSysProp(KOTLIN_SCRIPT_CLASSPATH_PROPERTY, "a:b:c") {
-            val classpath = scriptCompilationClasspathFromContext()
+            val classpath = scriptCompilationClasspathFromContextOrStlib(wholeClasspath = true)
             Assert.assertEquals("a:b:c", classpath.joinToString(":"))
         }
     }