Add special type classes in shared-compiler

Author: nikolay-egorov

jupyter-lib/shared-compiler/src/main/kotlin/org/jetbrains/kotlinx/jupyter/compiler/util/serializedCompiledScript.kt

@@ -20,9 +20,48 @@ data class SerializedCompiledScriptsData(
     }
 }
 
+@Serializable
+data class SerializableTypeInfo(val type: Type = Type.Custom, val isPrimitive: Boolean = false, val fullType: String = "") {
+    companion object {
+        val ignoreSet = setOf("int", "double", "boolean", "char", "float", "byte", "string", "entry")
+
+        val propertyNamesForNullFilter = setOf("data", "size")
+
+        fun makeFromSerializedVariablesState(type: String?, isContainer: Boolean?): SerializableTypeInfo {
+            val fullType = type.orEmpty()
+            val enumType = fullType.toTypeEnum()
+            val isPrimitive = !(
+                if (fullType != "Entry") (isContainer ?: false)
+                else true
+                )
+
+            return SerializableTypeInfo(enumType, isPrimitive, fullType)
+        }
+    }
+}
+
+@Serializable
+enum class Type {
+    Map,
+    Entry,
+    Array,
+    List,
+    Custom
+}
+
+fun String.toTypeEnum(): Type {
+    return when (this) {
+        "Map" -> Type.Map
+        "Entry" -> Type.Entry
+        "Array" -> Type.Array
+        "List" -> Type.List
+        else -> Type.Custom
+    }
+}
+
 @Serializable
 data class SerializedVariablesState(
-    val type: String = "",
+    val type: SerializableTypeInfo = SerializableTypeInfo(),
     val value: String? = null,
     val isContainer: Boolean = false,
     val stateId: String = ""

src/main/kotlin/org/jetbrains/kotlinx/jupyter/apiImpl.kt

@@ -133,8 +133,9 @@ class NotebookImpl(
 
     private val history = arrayListOf<CodeCellImpl>()
     private var mainCellCreated = false
-    private val unchangedVariables: MutableSet<String> = mutableSetOf()
+    private val _unchangedVariables: MutableSet<String> = mutableSetOf()
 
+    val unchangedVariables: Set<String> get() = _unchangedVariables
     val displays = DisplayContainerImpl()
 
     override fun getAllDisplays(): List<DisplayResultWithCell> {
@@ -153,16 +154,14 @@ class NotebookImpl(
     fun updateVariablesState(evaluator: InternalEvaluator) {
         variablesState += evaluator.variablesHolder
         currentCellVariables = evaluator.cellVariables
-        unchangedVariables.clear()
-        unchangedVariables.addAll(evaluator.getUnchangedVariables())
+        _unchangedVariables.clear()
+        _unchangedVariables.addAll(evaluator.getUnchangedVariables())
     }
 
     fun updateVariablesState(varsStateUpdate: Map<String, VariableState>) {
         variablesState += varsStateUpdate
     }
 
-    fun unchangedVariables(): Set<String> = unchangedVariables
-
     fun variablesReportAsHTML(): String {
         return generateHTMLVarsReport(variablesState)
     }

src/main/kotlin/org/jetbrains/kotlinx/jupyter/protocol.kt

@@ -2,8 +2,11 @@ package org.jetbrains.kotlinx.jupyter
 
 import ch.qos.logback.classic.Level
 import kotlinx.serialization.json.Json
+import kotlinx.serialization.json.JsonElement
+import kotlinx.serialization.json.JsonNull
 import kotlinx.serialization.json.JsonObject
 import kotlinx.serialization.json.encodeToJsonElement
+import kotlinx.serialization.json.jsonObject
 import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlinx.jupyter.LoggingManagement.disableLogging
 import org.jetbrains.kotlinx.jupyter.LoggingManagement.mainLoggerLevel
@@ -82,7 +85,6 @@ class OkResponseWithMessage(
                 )
             )
         }
-
         socket.send(
             makeReplyMessage(
                 requestMsg,
@@ -92,7 +94,7 @@ class OkResponseWithMessage(
                     "engine" to Json.encodeToJsonElement(requestMsg.data.header?.session),
                     "status" to Json.encodeToJsonElement("ok"),
                     "started" to Json.encodeToJsonElement(startedTime),
-                    "eval_metadata" to Json.encodeToJsonElement(metadata),
+                    "eval_metadata" to Json.encodeToJsonElement(metadata.convertToNullIfEmpty()),
                 ),
                 content = ExecuteReply(
                     MessageStatus.OK,
@@ -317,7 +319,7 @@ fun JupyterConnection.Socket.shellMessagesHandler(msg: Message, repl: ReplForJup
             if (data.isEmpty()) return sendWrapped(msg, makeReplyMessage(msg, MessageType.SERIALIZATION_REPLY))
             log.debug("Message data: $data")
             val messageContent = getVariablesDescriptorsFromJson(data)
-            GlobalScope.launch(Dispatchers.Default) {
+            connection.launchJob {
                 repl.serializeVariables(
                     messageContent.topLevelDescriptorName,
                     messageContent.descriptorsState,
@@ -343,7 +345,7 @@ fun JupyterConnection.Socket.shellMessagesHandler(msg: Message, repl: ReplForJup
             }
         }
         is SerializationRequest -> {
-            GlobalScope.launch(Dispatchers.Default) {
+            connection.launchJob {
                 if (content.topLevelDescriptorName.isNotEmpty()) {
                     repl.serializeVariables(content.topLevelDescriptorName, content.descriptorsState, commID = content.commId, content.pathToDescriptor) { result ->
                         sendWrapped(msg, makeReplyMessage(msg, MessageType.SERIALIZATION_REPLY, content = result))
@@ -545,3 +547,8 @@ fun JupyterConnection.evalWithIO(repl: ReplForJupyter, srcMessage: Message, body
         KernelStreams.setStreams(false, out, err)
     }
 }
+
+fun EvaluatedSnippetMetadata?.convertToNullIfEmpty(): JsonElement? {
+    val jsonNode = Json.encodeToJsonElement(this)
+    return if (jsonNode is JsonNull || jsonNode?.jsonObject.isEmpty()) null else jsonNode
+}

src/main/kotlin/org/jetbrains/kotlinx/jupyter/repl.kt

@@ -433,9 +433,10 @@ class ReplForJupyterImpl(
             val newImports: List<String>
             val oldDeclarations: MutableMap<String, Int> = mutableMapOf()
             oldDeclarations.putAll(internalEvaluator.getVariablesDeclarationInfo())
+            val jupyterId = evalData.jupyterId
             val result = try {
-                log.debug("Current cell id: ${evalData.jupyterId}")
-                executor.execute(evalData.code, evalData.displayHandler, currentCellId = evalData.jupyterId - 1) { internalId, codeToExecute ->
+                log.debug("Current cell id: $jupyterId")
+                executor.execute(evalData.code, evalData.displayHandler, currentCellId = jupyterId - 1) { internalId, codeToExecute ->
                     if (evalData.storeHistory) {
                         cell = notebook.addCell(internalId, codeToExecute, EvalData(evalData))
                     }
@@ -464,7 +465,7 @@ class ReplForJupyterImpl(
             // printVars()
             // printUsagesInfo(jupyterId, cellVariables[jupyterId - 1])
             val variablesCells: Map<String, Int> = notebook.variablesState.mapValues { internalEvaluator.findVariableCell(it.key) }
-            val serializedData = variablesSerializer.serializeVariables(jupyterId - 1, notebook.variablesState, oldDeclarations, variablesCells, notebook.unchangedVariables())
+            val serializedData = variablesSerializer.serializeVariables(jupyterId - 1, notebook.variablesState, oldDeclarations, variablesCells, notebook.unchangedVariables)
 
             GlobalScope.launch(Dispatchers.Default) {
                 variablesSerializer.tryValidateCache(jupyterId - 1, notebook.cellVariables)

src/main/kotlin/org/jetbrains/kotlinx/jupyter/serializationUtils.kt

@@ -5,6 +5,7 @@ import kotlinx.serialization.json.Json
 import kotlinx.serialization.json.JsonObject
 import kotlinx.serialization.json.decodeFromJsonElement
 import org.jetbrains.kotlinx.jupyter.api.VariableState
+import org.jetbrains.kotlinx.jupyter.compiler.util.SerializableTypeInfo
 import org.jetbrains.kotlinx.jupyter.compiler.util.SerializedVariablesState
 import java.lang.reflect.Field
 import kotlin.contracts.ExperimentalContracts
@@ -32,14 +33,14 @@ enum class PropertiesType {
 }
 
 @Serializable
-data class SerializedCommMessageContent(
+data class VariablesStateCommMessageContent(
     val topLevelDescriptorName: String,
     val descriptorsState: Map<String, SerializedVariablesState>,
     val pathToDescriptor: List<String> = emptyList()
 )
 
-fun getVariablesDescriptorsFromJson(json: JsonObject): SerializedCommMessageContent {
-    return Json.decodeFromJsonElement<SerializedCommMessageContent>(json)
+fun getVariablesDescriptorsFromJson(json: JsonObject): VariablesStateCommMessageContent {
+    return Json.decodeFromJsonElement<VariablesStateCommMessageContent>(json)
 }
 
 class ProcessedSerializedVarsState(
@@ -216,7 +217,12 @@ class VariablesSerializer(
             } else {
                 ""
             }
-            val serializedVersion = SerializedVariablesState(simpleTypeName, stringedValue, true, varID)
+            val serializedVersion = SerializedVariablesState(
+                SerializableTypeInfo.makeFromSerializedVariablesState(simpleTypeName, true),
+                stringedValue,
+                true,
+                varID
+            )
             val descriptors = serializedVersion.fieldDescriptor
 
             // only for set case
@@ -700,7 +706,12 @@ class VariablesSerializer(
                 ""
             }
 
-        val serializedVariablesState = SerializedVariablesState(type, getProperString(value), isContainer, finalID)
+        val serializedVariablesState = SerializedVariablesState(
+            SerializableTypeInfo.makeFromSerializedVariablesState(simpleTypeName, isContainer),
+            getProperString(value),
+            isContainer,
+            finalID
+        )
 
         return ProcessedSerializedVarsState(serializedVariablesState, membersProperties?.toTypedArray())
     }

src/test/kotlin/org/jetbrains/kotlinx/jupyter/test/repl/ReplTests.kt

@@ -810,14 +810,13 @@ class ReplVarsTest : AbstractSingleReplTest() {
             """.trimIndent(),
             jupyterId = 1
         )
-        var state = repl.notebook.unchangedVariables()
-        val res = eval(
+        eval(
             """
             l += 11111
             """.trimIndent(),
             jupyterId = 2
         ).metadata.evaluatedVariablesState
-        state = repl.notebook.unchangedVariables()
+        val state: Set<String> = repl.notebook.unchangedVariables
         assertEquals(1, state.size)
         assertTrue(state.contains("m"))
     }
@@ -887,7 +886,7 @@ class ReplVarsTest : AbstractSingleReplTest() {
             """.trimIndent(),
             jupyterId = 1
         )
-        var state = repl.notebook.unchangedVariables()
+        var state = repl.notebook.unchangedVariables
         assertEquals(3, state.size)
 
         eval(
@@ -898,7 +897,7 @@ class ReplVarsTest : AbstractSingleReplTest() {
             """.trimIndent(),
             jupyterId = 2
         )
-        state = repl.notebook.unchangedVariables()
+        state = repl.notebook.unchangedVariables
         assertEquals(0, state.size)
 
         eval(
@@ -907,7 +906,7 @@ class ReplVarsTest : AbstractSingleReplTest() {
             """.trimIndent(),
             jupyterId = 3
         )
-        state = repl.notebook.unchangedVariables()
+        state = repl.notebook.unchangedVariables
         assertEquals(1, state.size)
     }
 }
@@ -943,7 +942,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
         assertEquals(listOf(1, 2, 3, 4).toString().substring(1, actualContainer.value!!.length + 1), actualContainer.value)
 
         val serializer = repl.variablesSerializer
-        val newData = serializer.doIncrementalSerialization(0, "x", "data", actualContainer)
+        serializer.doIncrementalSerialization(0, "x", "data", actualContainer)
     }
 
     @Test
@@ -959,7 +958,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
         assertEquals(2, varsData.size)
         assertTrue(varsData.containsKey("x"))
         assertTrue(varsData.containsKey("f"))
-        var unchangedVariables = repl.notebook.unchangedVariables()
+        var unchangedVariables = repl.notebook.unchangedVariables
         assertTrue(unchangedVariables.isNotEmpty())
 
         eval(
@@ -968,7 +967,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
             """.trimIndent(),
             jupyterId = 1
         )
-        unchangedVariables = repl.notebook.unchangedVariables()
+        unchangedVariables = repl.notebook.unchangedVariables
         assertTrue(unchangedVariables.contains("x"))
         assertTrue(unchangedVariables.contains("f"))
     }
@@ -1032,7 +1031,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
 
         val serializer = repl.variablesSerializer
 
-        val newData = serializer.doIncrementalSerialization(0, "c", "i", descriptor["i"]!!)
+        serializer.doIncrementalSerialization(0, "c", "i", descriptor["i"]!!)
     }
 
     @Test
@@ -1321,7 +1320,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
             """.trimIndent(),
             jupyterId = 1
         )
-        val state = repl.notebook.unchangedVariables()
+        val state = repl.notebook.unchangedVariables
         val setOfCell = setOf("x", "f", "z")
         assertTrue(state.isNotEmpty())
         assertEquals(setOfCell, state)
@@ -1348,7 +1347,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
             """.trimIndent(),
             jupyterId = 1
         )
-        var state = repl.notebook.unchangedVariables()
+        var state = repl.notebook.unchangedVariables
         val setOfCell = setOf("x", "f", "z")
         assertTrue(state.isNotEmpty())
         assertEquals(setOfCell, state)
@@ -1372,9 +1371,9 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
             """.trimIndent(),
             jupyterId = 3
         )
-        state = repl.notebook.unchangedVariables()
-//        assertTrue(state.isNotEmpty())
-//        assertEquals(state, setOfPrevCell)
+        state = repl.notebook.unchangedVariables
+        assertTrue(state.isEmpty())
+        // assertEquals(state, setOfPrevCell)
 
         eval(
             """
@@ -1384,20 +1383,20 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
             """.trimIndent(),
             jupyterId = 4
         )
-        state = repl.notebook.unchangedVariables()
+        state = repl.notebook.unchangedVariables
         assertTrue(state.isEmpty())
     }
 
     @Test
     fun testSerializationClearInfo() {
-        var res = eval(
+        eval(
             """
             val x = listOf(1, 2, 3, 4)
             """.trimIndent(),
             jupyterId = 1
         ).metadata.evaluatedVariablesState
-        var state = repl.notebook.unchangedVariables()
-        res = eval(
+        repl.notebook.unchangedVariables
+        eval(
             """
             val x = listOf(1, 2, 3, 4)
             """.trimIndent(),