Add ReplCompilerMode enum, allowing users of the Kernel to start it in either K1 or K2 mode

Merge-request: KDS-MR-59
Merged-by: Christian Melchior <[email protected]>

Author: cmelchior

Diff for: jupyter-lib/shared-compiler/src/main/kotlin/org/jetbrains/kotlinx/jupyter/config/Compilation.kt

@@ -4,9 +4,12 @@ import jupyter.kotlin.CompilerArgs
 import jupyter.kotlin.DependsOn
 import jupyter.kotlin.Repository
 import org.jetbrains.kotlin.scripting.resolve.skipExtensionsResolutionForImplicitsExceptInnermost
+import org.jetbrains.kotlinx.jupyter.api.KernelLoggerFactory
 import org.jetbrains.kotlinx.jupyter.compiler.CompilerArgsConfigurator
 import org.jetbrains.kotlinx.jupyter.compiler.ScriptDataCollector
 import org.jetbrains.kotlinx.jupyter.messaging.ExecutionCount
+import org.jetbrains.kotlinx.jupyter.startup.DEFAULT
+import org.jetbrains.kotlinx.jupyter.startup.ReplCompilerMode
 import java.io.File
 import kotlin.script.experimental.api.KotlinType
 import kotlin.script.experimental.api.ScriptAcceptedLocation
@@ -61,14 +64,20 @@ val ScriptCompilationConfigurationKeys.jupyterOptions by PropertiesCollection.ke
     defaultValue = JupyterCompilingOptions.DEFAULT,
 )
 
+// Also called from IDEA
 fun getCompilationConfiguration(
     scriptClasspath: List<File> = emptyList(),
     scriptReceivers: List<Any> = emptyList(),
     compilerArgsConfigurator: CompilerArgsConfigurator,
     scriptingClassGetter: GetScriptingClassByClassLoader = JvmGetScriptingClass(),
     scriptDataCollectors: List<ScriptDataCollector> = emptyList(),
+    replCompilerMode: ReplCompilerMode = ReplCompilerMode.DEFAULT,
+    loggerFactory: KernelLoggerFactory,
     body: ScriptCompilationConfiguration.Builder.() -> Unit = {},
 ): ScriptCompilationConfiguration {
+    if (replCompilerMode == ReplCompilerMode.K2) {
+        loggerFactory.getLogger("getCompilationConfiguration").warn("K2 Repl Mode is ignored for now. Falling back to K1")
+    }
     return ScriptCompilationConfiguration {
         hostConfiguration.update {
             it.with {

Diff for: jupyter-lib/shared-compiler/src/main/kotlin/org/jetbrains/kotlinx/jupyter/repl/creating/DefaultReplComponentsProvider.kt

@@ -18,6 +18,7 @@ import org.jetbrains.kotlinx.jupyter.repl.MavenRepositoryCoordinates
 import org.jetbrains.kotlinx.jupyter.repl.ReplRuntimeProperties
 import org.jetbrains.kotlinx.jupyter.repl.config.DefaultReplSettings
 import org.jetbrains.kotlinx.jupyter.repl.embedded.InMemoryReplResultsHolder
+import org.jetbrains.kotlinx.jupyter.startup.ReplCompilerMode
 import org.jetbrains.kotlinx.jupyter.util.toUpperCaseAsciiOnly
 import java.io.File
 
@@ -115,4 +116,8 @@ open class DefaultReplComponentsProvider(
     override fun provideInMemoryReplResultsHolder(): InMemoryReplResultsHolder {
         return _inMemoryResultsHolder
     }
+
+    override fun provideReplCompilerMode(): ReplCompilerMode {
+        return _settings.kernelConfig.replCompilerMode
+    }
 }

Diff for: jupyter-lib/shared-compiler/src/main/kotlin/org/jetbrains/kotlinx/jupyter/repl/creating/LazilyConstructibleReplComponentsProvider.kt

@@ -24,6 +24,7 @@ import org.jetbrains.kotlinx.jupyter.repl.ReplOptions
 import org.jetbrains.kotlinx.jupyter.repl.ReplRuntimeProperties
 import org.jetbrains.kotlinx.jupyter.repl.embedded.InMemoryReplResultsHolder
 import org.jetbrains.kotlinx.jupyter.repl.notebook.MutableNotebook
+import org.jetbrains.kotlinx.jupyter.startup.ReplCompilerMode
 import java.io.File
 
 interface LazilyConstructibleReplComponentsProvider : ReplComponentsProvider {
@@ -80,4 +81,6 @@ interface LazilyConstructibleReplComponentsProvider : ReplComponentsProvider {
     fun provideLibraryReferenceParser(): LibraryReferenceParser
 
     fun provideInMemoryReplResultsHolder(): InMemoryReplResultsHolder
+
+    fun provideReplCompilerMode(): ReplCompilerMode
 }

Diff for: jupyter-lib/shared-compiler/src/main/kotlin/org/jetbrains/kotlinx/jupyter/repl/creating/LazilyConstructibleReplComponentsProviderImpl.kt

@@ -21,6 +21,7 @@ import org.jetbrains.kotlinx.jupyter.repl.MavenRepositoryCoordinates
 import org.jetbrains.kotlinx.jupyter.repl.ReplOptions
 import org.jetbrains.kotlinx.jupyter.repl.ReplRuntimeProperties
 import org.jetbrains.kotlinx.jupyter.repl.embedded.InMemoryReplResultsHolder
+import org.jetbrains.kotlinx.jupyter.startup.ReplCompilerMode
 import java.io.File
 
 abstract class LazilyConstructibleReplComponentsProviderImpl : LazilyConstructibleReplComponentsProvider {
@@ -51,6 +52,7 @@ abstract class LazilyConstructibleReplComponentsProviderImpl : LazilyConstructib
     override val magicsHandler: LibrariesAwareMagicsHandler? by lazy { provideMagicsHandler() }
     override val libraryReferenceParser: LibraryReferenceParser by lazy { provideLibraryReferenceParser() }
     override val inMemoryReplResultsHolder: InMemoryReplResultsHolder by lazy { provideInMemoryReplResultsHolder() }
+    override val replCompilerMode: ReplCompilerMode by lazy { provideReplCompilerMode() }
 
     // TODO: add other methods incl. display handler and socket messages listener
     // Inheritors should be constructed of connection (JupyterConnection)

Diff for: jupyter-lib/shared-compiler/src/main/kotlin/org/jetbrains/kotlinx/jupyter/repl/creating/ReplComponentsProvider.kt

@@ -24,6 +24,7 @@ import org.jetbrains.kotlinx.jupyter.repl.ReplOptions
 import org.jetbrains.kotlinx.jupyter.repl.ReplRuntimeProperties
 import org.jetbrains.kotlinx.jupyter.repl.embedded.InMemoryReplResultsHolder
 import org.jetbrains.kotlinx.jupyter.repl.notebook.MutableNotebook
+import org.jetbrains.kotlinx.jupyter.startup.ReplCompilerMode
 import java.io.File
 
 interface ReplComponentsProvider {
@@ -54,4 +55,5 @@ interface ReplComponentsProvider {
     val magicsHandler: LibrariesAwareMagicsHandler?
     val libraryReferenceParser: LibraryReferenceParser
     val inMemoryReplResultsHolder: InMemoryReplResultsHolder
+    val replCompilerMode: ReplCompilerMode
 }

Diff for: jupyter-lib/shared-compiler/src/main/kotlin/org/jetbrains/kotlinx/jupyter/repl/creating/ReplComponentsProviderBase.kt

@@ -41,6 +41,8 @@ import org.jetbrains.kotlinx.jupyter.repl.embedded.InMemoryReplResultsHolder
 import org.jetbrains.kotlinx.jupyter.repl.embedded.NoOpInMemoryReplResultsHolder
 import org.jetbrains.kotlinx.jupyter.repl.notebook.MutableNotebook
 import org.jetbrains.kotlinx.jupyter.repl.notebook.impl.NotebookImpl
+import org.jetbrains.kotlinx.jupyter.startup.DEFAULT
+import org.jetbrains.kotlinx.jupyter.startup.ReplCompilerMode
 import org.jetbrains.kotlinx.jupyter.util.asCommonFactory
 import java.io.File
 
@@ -130,4 +132,6 @@ abstract class ReplComponentsProviderBase : LazilyConstructibleReplComponentsPro
     override fun provideLibraryReferenceParser(): LibraryReferenceParser = LibraryReferenceParserImpl(libraryInfoCache)
 
     override fun provideInMemoryReplResultsHolder(): InMemoryReplResultsHolder = NoOpInMemoryReplResultsHolder
+
+    override fun provideReplCompilerMode(): ReplCompilerMode = ReplCompilerMode.DEFAULT
 }

Diff for: jupyter-lib/shared-compiler/src/main/kotlin/org/jetbrains/kotlinx/jupyter/startup/Arguments.kt

@@ -18,6 +18,18 @@ import java.util.EnumMap
 
 typealias KernelPorts = Map<JupyterSocketType, Int>
 
+/**
+ * This is used to represent whether the underlying REPL is running in either K1 or K2 mode.
+ */
+enum class ReplCompilerMode {
+    K1,
+    K2,
+    ;
+
+    companion object
+}
+
+val ReplCompilerMode.Companion.DEFAULT get() = ReplCompilerMode.K1
 const val KERNEL_TRANSPORT_SCHEME = "tcp"
 const val KERNEL_SIGNATURE_SCHEME = "HmacSHA256"
 
@@ -37,6 +49,7 @@ data class KernelArgs(
     val debugPort: Int?,
     val clientType: String?,
     val jvmTargetForSnippets: String?,
+    val replCompilerMode: ReplCompilerMode,
 ) {
     fun parseParams(): KernelJupyterParams {
         return KernelJupyterParams.fromFile(cfgFile)
@@ -53,6 +66,7 @@ data class KernelArgs(
             debugPort?.let { add("-debugPort=$it") }
             clientType?.let { add("-client=$it") }
             jvmTargetForSnippets?.let { add("-jvmTarget=$it") }
+            add("-replCompilerMode=${replCompilerMode.name}")
         }
     }
 }
@@ -120,6 +134,7 @@ data class KernelConfig(
     val debugPort: Int? = null,
     val clientType: String? = null,
     val jvmTargetForSnippets: String? = null,
+    val replCompilerMode: ReplCompilerMode = ReplCompilerMode.DEFAULT,
 ) {
     val hmac by lazy {
         HMAC(signatureScheme.replace("-", ""), signatureKey)
@@ -133,7 +148,7 @@ data class KernelConfig(
         val format = Json { prettyPrint = true }
         cfgFile.writeText(format.encodeToString(params))
 
-        return KernelArgs(cfgFile, scriptClasspath, homeDir, debugPort, clientType, jvmTargetForSnippets)
+        return KernelArgs(cfgFile, scriptClasspath, homeDir, debugPort, clientType, jvmTargetForSnippets, replCompilerMode)
     }
 }
 
@@ -149,13 +164,15 @@ fun createClientKotlinKernelConfig(
     host: String,
     ports: KernelPorts,
     signatureKey: String,
+    replCompilerMode: ReplCompilerMode,
 ) = KernelConfig(
     host = host,
     ports = ports,
     transport = KERNEL_TRANSPORT_SCHEME,
     signatureScheme = KERNEL_SIGNATURE_SCHEME,
     signatureKey = signatureKey,
     homeDir = null,
+    replCompilerMode = replCompilerMode,
 )
 
 /**
@@ -179,6 +196,7 @@ fun createKotlinKernelConfig(
     homeDir: File? = null,
     debugPort: Int? = null,
     clientType: String? = null,
+    replCompilerMode: ReplCompilerMode = ReplCompilerMode.DEFAULT,
 ) = KernelConfig(
     ports = ports,
     transport = KERNEL_TRANSPORT_SCHEME,
@@ -188,6 +206,7 @@ fun createKotlinKernelConfig(
     homeDir = homeDir,
     debugPort = debugPort,
     clientType = clientType,
+    replCompilerMode = replCompilerMode,
 )
 
 const val MAIN_CLASS_NAME = "org.jetbrains.kotlinx.jupyter.IkotlinKt"
@@ -230,5 +249,6 @@ fun KernelArgs.getConfig(): KernelConfig {
         debugPort = debugPort,
         clientType = clientType,
         jvmTargetForSnippets = jvmTargetForSnippets,
+        replCompilerMode = replCompilerMode,
     )
 }

Diff for: src/main/kotlin/org/jetbrains/kotlinx/jupyter/Ikotlin.kt

@@ -36,8 +36,10 @@ import org.jetbrains.kotlinx.jupyter.repl.config.DefaultReplSettings
 import org.jetbrains.kotlinx.jupyter.repl.creating.DefaultReplComponentsProvider
 import org.jetbrains.kotlinx.jupyter.repl.creating.createRepl
 import org.jetbrains.kotlinx.jupyter.repl.embedded.NoOpInMemoryReplResultsHolder
+import org.jetbrains.kotlinx.jupyter.startup.DEFAULT
 import org.jetbrains.kotlinx.jupyter.startup.KernelArgs
 import org.jetbrains.kotlinx.jupyter.startup.KernelConfig
+import org.jetbrains.kotlinx.jupyter.startup.ReplCompilerMode
 import org.jetbrains.kotlinx.jupyter.startup.getConfig
 import org.slf4j.Logger
 import java.io.File
@@ -53,6 +55,7 @@ private fun parseCommandLine(vararg args: String): KernelArgs {
     var debugPort: Int? = null
     var clientType: String? = null
     var jvmTargetForSnippets: String? = null
+    var replCompilerMode: ReplCompilerMode = ReplCompilerMode.DEFAULT
     args.forEach { arg ->
         when {
             arg.startsWith("-cp=") || arg.startsWith("-classpath=") -> {
@@ -73,6 +76,12 @@ private fun parseCommandLine(vararg args: String): KernelArgs {
             arg.startsWith("-jvmTarget") -> {
                 jvmTargetForSnippets = arg.substringAfter('=')
             }
+            arg.startsWith("-replCompilerMode=") -> {
+                val userMode = arg.substringAfter('=')
+                replCompilerMode = ReplCompilerMode.entries.find {
+                    it.name == userMode
+                } ?: throw IllegalArgumentException("Invalid replCompilerMode: $userMode")
+            }
             else -> {
                 cfgFile?.let { throw IllegalArgumentException("config file already set to $it") }
                 cfgFile = File(arg)
@@ -82,7 +91,7 @@ private fun parseCommandLine(vararg args: String): KernelArgs {
     val cfgFileValue = cfgFile ?: throw IllegalArgumentException("config file is not provided")
     if (!cfgFileValue.exists() || !cfgFileValue.isFile) throw IllegalArgumentException("invalid config file $cfgFileValue")
 
-    return KernelArgs(cfgFileValue, classpath ?: emptyList(), homeDir, debugPort, clientType, jvmTargetForSnippets)
+    return KernelArgs(cfgFileValue, classpath ?: emptyList(), homeDir, debugPort, clientType, jvmTargetForSnippets, replCompilerMode)
 }
 
 fun printClassPath(logger: Logger) {
@@ -129,7 +138,16 @@ fun embedKernel(
     scriptReceivers: List<Any>? = null,
 ) {
     val scriptClasspath = System.getProperty("java.class.path").split(File.pathSeparator).toTypedArray().map { File(it) }
-    val kernelConfig = KernelArgs(cfgFile, scriptClasspath, null, null, null, null).getConfig()
+    val kernelConfig =
+        KernelArgs(
+            cfgFile = cfgFile,
+            scriptClasspath = scriptClasspath,
+            homeDir = null,
+            debugPort = null,
+            clientType = null,
+            jvmTargetForSnippets = null,
+            replCompilerMode = ReplCompilerMode.DEFAULT,
+        ).getConfig()
     val replSettings =
         createReplSettings(
             DefaultKernelLoggerFactory,

Diff for: src/main/kotlin/org/jetbrains/kotlinx/jupyter/repl/creating/CreatingUtil.kt

@@ -16,6 +16,8 @@ import org.jetbrains.kotlinx.jupyter.repl.ReplForJupyter
 import org.jetbrains.kotlinx.jupyter.repl.ReplRuntimeProperties
 import org.jetbrains.kotlinx.jupyter.repl.embedded.InMemoryReplResultsHolder
 import org.jetbrains.kotlinx.jupyter.repl.embedded.NoOpInMemoryReplResultsHolder
+import org.jetbrains.kotlinx.jupyter.startup.DEFAULT
+import org.jetbrains.kotlinx.jupyter.startup.ReplCompilerMode
 import java.io.File
 
 fun createRepl(
@@ -32,6 +34,7 @@ fun createRepl(
     communicationFacility: JupyterCommunicationFacility = CommunicationFacilityMock,
     debugPort: Int? = null,
     inMemoryReplResultsHolder: InMemoryReplResultsHolder = NoOpInMemoryReplResultsHolder,
+    replCompilerMode: ReplCompilerMode = ReplCompilerMode.DEFAULT,
 ): ReplForJupyter {
     val componentsProvider =
         object : ReplComponentsProviderBase() {
@@ -66,6 +69,8 @@ fun createRepl(
             override fun provideLibraryReferenceParser() = httpUtil.libraryReferenceParser
 
             override fun provideInMemoryReplResultsHolder() = inMemoryReplResultsHolder
+
+            override fun provideReplCompilerMode(): ReplCompilerMode = replCompilerMode
         }
     return componentsProvider.createRepl()
 }

Diff for: src/main/kotlin/org/jetbrains/kotlinx/jupyter/repl/creating/ReplFactoryBase.kt

@@ -32,6 +32,7 @@ class ReplFactoryBase(
                 sessionOptions,
                 magicsHandler,
                 inMemoryReplResultsHolder,
+                replCompilerMode,
             )
         }
     }

Diff for: src/main/kotlin/org/jetbrains/kotlinx/jupyter/repl/impl/ReplForJupyterImpl.kt

@@ -53,6 +53,7 @@ import org.jetbrains.kotlinx.jupyter.compiler.DefaultCompilerArgsConfigurator
 import org.jetbrains.kotlinx.jupyter.compiler.ScriptDeclarationsCollectorInternal
 import org.jetbrains.kotlinx.jupyter.compiler.ScriptImportsCollector
 import org.jetbrains.kotlinx.jupyter.config.CellId
+import org.jetbrains.kotlinx.jupyter.config.DefaultKernelLoggerFactory
 import org.jetbrains.kotlinx.jupyter.config.addBaseClass
 import org.jetbrains.kotlinx.jupyter.config.catchAll
 import org.jetbrains.kotlinx.jupyter.config.defaultRuntimeProperties
@@ -118,6 +119,7 @@ import org.jetbrains.kotlinx.jupyter.repl.result.InternalMetadataImpl
 import org.jetbrains.kotlinx.jupyter.repl.result.InternalReplResult
 import org.jetbrains.kotlinx.jupyter.repl.result.SerializedCompiledScriptsData
 import org.jetbrains.kotlinx.jupyter.repl.result.buildScriptsData
+import org.jetbrains.kotlinx.jupyter.startup.ReplCompilerMode
 import java.io.Closeable
 import java.io.File
 import java.net.URLClassLoader
@@ -162,6 +164,7 @@ class ReplForJupyterImpl(
     override val sessionOptions: SessionOptions,
     customMagicsHandler: LibrariesAwareMagicsHandler?,
     private val inMemoryReplResultsHolder: InMemoryReplResultsHolder,
+    private val replCompilerMode: ReplCompilerMode,
 ) : ReplForJupyter, BaseKernelHost, UserHandlesProvider, Closeable {
     private val logger = loggerFactory.getLogger(this::class)
 
@@ -244,6 +247,7 @@ class ReplForJupyterImpl(
             scriptReceivers,
             compilerArgsConfigurator,
             scriptDataCollectors = listOf(importsCollector, declarationsCollector),
+            loggerFactory = DefaultKernelLoggerFactory,
             body = { addBaseClass<ScriptTemplateWithDisplayHelpers>() },
         ).with {
             refineConfiguration {
@@ -305,12 +309,25 @@ class ReplForJupyterImpl(
             constructorArgs(this@ReplForJupyterImpl)
         }
 
-    private val jupyterCompiler by lazy {
-        JupyterCompilerWithCompletion.create(compilerConfiguration, evaluatorConfiguration)
+    private val jupyterCompiler: JupyterCompilerWithCompletion by lazy {
+        when (replCompilerMode) {
+            ReplCompilerMode.K1 -> {
+                JupyterCompilerWithCompletion.create(
+                    compilerConfiguration,
+                    evaluatorConfiguration,
+                )
+            }
+            ReplCompilerMode.K2 -> {
+                k2Unsupported()
+            }
+        }
     }
 
     private val evaluator: BasicJvmReplEvaluator by lazy {
-        BasicJvmReplEvaluator()
+        when (replCompilerMode) {
+            ReplCompilerMode.K1 -> BasicJvmReplEvaluator()
+            ReplCompilerMode.K2 -> k2Unsupported()
+        }
     }
 
     private val hostProvider =
@@ -780,4 +797,8 @@ class ReplForJupyterImpl(
     override fun close() {
         notebook.closeIfPossible()
     }
+
+    private fun k2Unsupported(): Nothing {
+        throw IllegalStateException("K2 REPL is not supported yet")
+    }
 }