Reuse script args substitution for replacing bindings in JSR 223 sample engines, fixes KT-15450

compiler/cli/cli-common/src/org/jetbrains/kotlin/cli/common/repl/KotlinJsr223JvmScriptEngineBase.kt

@@ -22,18 +22,20 @@ import java.util.concurrent.atomic.AtomicInteger
 import javax.script.*
 
 val KOTLIN_SCRIPT_HISTORY_BINDINGS_KEY = "kotlin.script.history"
-
+val KOTLIN_SCRIPT_LINE_NUMBER_BINDINGS_KEY = "kotlin.script.line.number"
 
 // TODO consider additional error handling
-val Bindings.kotlinScriptHistory: MutableList<ReplCodeLine>
+var Bindings.kotlinScriptHistory: MutableList<ReplCodeLine>
     @Suppress("UNCHECKED_CAST")
     get() = getOrPut(KOTLIN_SCRIPT_HISTORY_BINDINGS_KEY, { arrayListOf<ReplCodeLine>() }) as MutableList<ReplCodeLine>
+    set(v) { put(KOTLIN_SCRIPT_HISTORY_BINDINGS_KEY, v) }
 
+val Bindings.kotlinScriptLineNumber: AtomicInteger
+    @Suppress("UNCHECKED_CAST")
+    get() = getOrPut(KOTLIN_SCRIPT_LINE_NUMBER_BINDINGS_KEY, { AtomicInteger(0) }) as AtomicInteger
 
 abstract class KotlinJsr223JvmScriptEngineBase(protected val myFactory: ScriptEngineFactory) : AbstractScriptEngine(), ScriptEngine, Compilable {
 
-    protected var codeLineNumber = AtomicInteger(0)
-
     protected abstract val replCompiler: ReplCompileAction
     protected abstract val replScriptEvaluator: ReplFullEvaluator
 
@@ -49,13 +51,20 @@ abstract class KotlinJsr223JvmScriptEngineBase(protected val myFactory: ScriptEn
 
     override fun getFactory(): ScriptEngineFactory = myFactory
 
-    fun nextCodeLine(code: String) = ReplCodeLine(codeLineNumber.incrementAndGet(), code)
+    // the parameter could be used in the future when we decide to keep state completely in the context and solve appropriate problems (now e.g. replCompiler keeps separate state)
+    fun nextCodeLine(@Suppress("UNUSED_PARAMETER") context: ScriptContext, code: String) = ReplCodeLine(this.context.getBindings(ScriptContext.ENGINE_SCOPE).kotlinScriptLineNumber.incrementAndGet(), code)
+
+    private fun getCurrentHistory(@Suppress("UNUSED_PARAMETER") context: ScriptContext) = this.context.getBindings(ScriptContext.ENGINE_SCOPE).kotlinScriptHistory
+
+    private fun setContextHistory(@Suppress("UNUSED_PARAMETER") context: ScriptContext, history: ArrayList<ReplCodeLine>) {
+        this.context.getBindings(ScriptContext.ENGINE_SCOPE).kotlinScriptHistory = history
+    }
 
     open fun overrideScriptArgs(context: ScriptContext): ScriptArgsWithTypes? = null
 
     open fun compileAndEval(script: String, context: ScriptContext): Any? {
-        val codeLine = nextCodeLine(script)
+        val codeLine = nextCodeLine(context, script)
-        val history = context.getBindings(ScriptContext.ENGINE_SCOPE).kotlinScriptHistory
+        val history = getCurrentHistory(context)
         val result = replScriptEvaluator.compileAndEval(codeLine, scriptArgs = overrideScriptArgs(context), verifyHistory = history)
         val ret = when (result) {
             is ReplEvalResult.ValueResult -> result.value
@@ -64,13 +73,13 @@ abstract class KotlinJsr223JvmScriptEngineBase(protected val myFactory: ScriptEn
             is ReplEvalResult.Incomplete -> throw ScriptException("error: incomplete code")
             is ReplEvalResult.HistoryMismatch -> throw ScriptException("Repl history mismatch at line: ${result.lineNo}")
         }
-        context.getBindings(ScriptContext.ENGINE_SCOPE).put(KOTLIN_SCRIPT_HISTORY_BINDINGS_KEY, ArrayList(result.completedEvalHistory))
+        setContextHistory(context, ArrayList(result.completedEvalHistory))
         return ret
     }
 
     open fun compile(script: String, context: ScriptContext): CompiledScript {
-        val codeLine = nextCodeLine(script)
+        val codeLine = nextCodeLine(context, script)
-        val history = context.getBindings(ScriptContext.ENGINE_SCOPE).kotlinScriptHistory
+        val history = getCurrentHistory(context)
 
         val result = replCompiler.compile(codeLine, history)
         val compiled = when (result) {
@@ -79,7 +88,8 @@ abstract class KotlinJsr223JvmScriptEngineBase(protected val myFactory: ScriptEn
             is ReplCompileResult.HistoryMismatch -> throw ScriptException("Repl history mismatch at line: ${result.lineNo}")
             is ReplCompileResult.CompiledClasses -> result
         }
-        context.getBindings(ScriptContext.ENGINE_SCOPE).put(KOTLIN_SCRIPT_HISTORY_BINDINGS_KEY, ArrayList(result.compiledHistory))
+        // TODO: check if it is ok to keep compiled history in the same place as compiledEval one
+        setContextHistory(context, ArrayList(result.compiledHistory))
         return CompiledKotlinScript(this, codeLine, compiled)
     }
 
@@ -98,7 +108,7 @@ abstract class KotlinJsr223JvmScriptEngineBase(protected val myFactory: ScriptEn
             is ReplEvalResult.Incomplete -> throw ScriptException("error: incomplete code")
             is ReplEvalResult.HistoryMismatch -> throw ScriptException("Repl history mismatch at line: ${result.lineNo}")
         }
-        context.getBindings(ScriptContext.ENGINE_SCOPE).put(KOTLIN_SCRIPT_HISTORY_BINDINGS_KEY, ArrayList(result.completedEvalHistory))
+        setContextHistory(context, ArrayList(result.completedEvalHistory))
         return ret
     }
 

compiler/cli/src/org/jetbrains/kotlin/cli/jvm/repl/GenericReplCompiler.kt

@@ -58,7 +58,7 @@ open class GenericReplCompiler(disposable: Disposable,
 
             removedCompiledLines.zip(removedAnalyzedLines).forEach {
                 if (it.first.first != it.second) {
-                    throw IllegalStateException("History mistmatch when resetting lines")
+                    throw IllegalStateException("History mismatch when resetting lines")
                 }
             }
 

libraries/examples/kotlin-jsr223-daemon-local-eval-example/src/test/kotlin/org/jetbrains/kotlin/script/jsr223/KotlinJsr223ScriptEngineIT.kt

@@ -92,6 +92,25 @@ obj
         val res3 = invocator!!.invokeMethod(res1, "fn1", 3)
         Assert.assertEquals(6, res3)
     }
+
+    @Test
+    fun testEvalWithContext() {
+        val engine = ScriptEngineManager().getEngineByExtension("kts")!!
+
+        engine.put("z", 33)
+
+        engine.eval("""val x = 10 + bindings["z"] as Int""")
+
+        val result = engine.eval("""x + 20""")
+        Assert.assertEquals(63, result)
+
+        // in the current implementation the history is shared between contexts, so "x" could also be used in this line,
+        // but this behaviour probably will not be preserved in the future, since contexts may become completely isolated
+        val result2 = engine.eval("""11 + bindings["boundValue"] as Int""", engine.createBindings().apply {
+            put("boundValue", 100)
+        })
+        Assert.assertEquals(111, result2)
+    }
 }
 
 fun assertThrows(exceptionClass: Class<*>, body: () -> Unit) {

libraries/examples/kotlin-jsr223-local-example/src/test/kotlin/org/jetbrains/kotlin/script/jsr223/KotlinJsr223ScriptEngineIT.kt

@@ -92,6 +92,25 @@ obj
         val res3 = invocator!!.invokeMethod(res1, "fn1", 3)
         Assert.assertEquals(6, res3)
     }
+
+    @Test
+    fun testEvalWithContext() {
+        val engine = ScriptEngineManager().getEngineByExtension("kts")!!
+
+        engine.put("z", 33)
+
+        engine.eval("""val x = 10 + bindings["z"] as Int""")
+
+        val result = engine.eval("""x + 20""")
+        Assert.assertEquals(63, result)
+
+        // in the current implementation the history is shared between contexts, so "x" could also be used in this line,
+        // but this behaviour probably will not be preserved in the future, since contexts may become completely isolated
+        val result2 = engine.eval("""11 + bindings["boundValue"] as Int""", engine.createBindings().apply {
+            put("boundValue", 100)
+        })
+        Assert.assertEquals(111, result2)
+    }
 }
 
 fun assertThrows(exceptionClass: Class<*>, body: () -> Unit) {

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

@@ -36,8 +36,8 @@ class KotlinJsr223JvmDaemonCompileScriptEngine(
         compilerJar: File,
         templateClasspath: List<File>,
         templateClassName: String,
-        getScriptArgs: (ScriptContext, Array<out KClass<out Any>>?) -> ScriptArgsWithTypes?,
+        val getScriptArgs: (ScriptContext, Array<out KClass<out Any>>?) -> ScriptArgsWithTypes?,
-        scriptArgsTypes: Array<out KClass<out Any>>?,
+        val scriptArgsTypes: Array<out KClass<out Any>>?,
         compilerOut: OutputStream = System.err
 ) : KotlinJsr223JvmScriptEngineBase(factory), KotlinJsr223JvmInvocableScriptEngine {
 
@@ -61,6 +61,8 @@ class KotlinJsr223JvmDaemonCompileScriptEngine(
 
     override val replScriptEvaluator: ReplFullEvaluator get() = localEvaluator
 
+    override fun overrideScriptArgs(context: ScriptContext): ScriptArgsWithTypes? = getScriptArgs(context, scriptArgsTypes)
+
     private fun connectToCompileService(compilerJar: File): CompileService {
         val compilerId = CompilerId.makeCompilerId(compilerJar)
         val daemonOptions = configureDaemonOptions()

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

@@ -38,8 +38,8 @@ class KotlinJsr223JvmLocalScriptEngine(
         factory: ScriptEngineFactory,
         val templateClasspath: List<File>,
         templateClassName: String,
-        getScriptArgs: (ScriptContext, Array<out KClass<out Any>>?) -> ScriptArgsWithTypes?,
+        val getScriptArgs: (ScriptContext, Array<out KClass<out Any>>?) -> ScriptArgsWithTypes?,
-        scriptArgsTypes: Array<out KClass<out Any>>?
+        val scriptArgsTypes: Array<out KClass<out Any>>?
 ) : KotlinJsr223JvmScriptEngineBase(factory), KotlinJsr223JvmInvocableScriptEngine {
 
     override val replCompiler: ReplCompiler by lazy {
@@ -54,6 +54,8 @@ class KotlinJsr223JvmLocalScriptEngine(
 
     override val replScriptEvaluator: ReplFullEvaluator get() = localEvaluator
 
+    override fun overrideScriptArgs(context: ScriptContext): ScriptArgsWithTypes? = getScriptArgs(context, scriptArgsTypes)
+
     private fun makeScriptDefinition(templateClasspath: List<File>, templateClassName: String): KotlinScriptDefinition {
         val classloader = URLClassLoader(templateClasspath.map { it.toURI().toURL() }.toTypedArray(), this.javaClass.classLoader)
         val cls = classloader.loadClass(templateClassName)