Add Serializable Vars State with nested fields descriptors

build.gradle.kts

@@ -114,7 +114,8 @@ tasks {
 
             "junit.jupiter.execution.parallel.enabled" to doParallelTesting.toString() as Any,
             "junit.jupiter.execution.parallel.mode.default" to "concurrent",
-            "junit.jupiter.execution.parallel.mode.classes.default" to "concurrent"
+            "junit.jupiter.execution.parallel.mode.classes.default" to "concurrent",
+            "jupyter.serialization.enabled" to "true"
         )
     }
 

jupyter-lib/api/src/main/kotlin/org/jetbrains/kotlinx/jupyter/api/VariableState.kt

@@ -1,33 +1,49 @@
 package org.jetbrains.kotlinx.jupyter.api
 
+import java.lang.reflect.Field
 import kotlin.reflect.KProperty
-import kotlin.reflect.KProperty1
 import kotlin.reflect.jvm.isAccessible
 
 interface VariableState {
-    val property: KProperty<*>
+    val property: Field
     val scriptInstance: Any?
     val stringValue: String?
     val value: Result<Any?>
+    val isRecursive: Boolean
+}
+
+class DependentLazyDelegate<T>(val initializer: () -> T?) {
+    private var cachedPropertyValue: T? = null
+    var isChanged: Boolean = true
+
+    operator fun getValue(thisRef: Any?, property: KProperty<*>): T? {
+        if (isChanged) {
+            cachedPropertyValue = initializer()
+        }
+        return cachedPropertyValue
+    }
 }
 
 data class VariableStateImpl(
-    override val property: KProperty1<Any, *>,
+    override val property: Field,
     override val scriptInstance: Any,
 ) : VariableState {
     private val stringCache = VariableStateCache<String?> {
         value.getOrNull()?.let { value ->
             try {
                 value.toString()
             } catch (e: Throwable) {
+                if (e is StackOverflowError) {
+                    isRecursive = true
+                }
                 "${value::class.simpleName}: [exception thrown: $e]"
             }
         }
     }
+    override var isRecursive: Boolean = false
 
-    private val valCache = VariableStateCache<Result<Any?>> (
-        {
-            oldValue, newValue ->
+    private val valCache = VariableStateCache<Result<Any?>>(
+        { oldValue, newValue ->
             oldValue.getOrNull() !== newValue.getOrNull()
         },
         {
@@ -47,49 +63,50 @@ data class VariableStateImpl(
         }
     }
 
-    override val stringValue: String? get() = stringCache.get()
+    override val stringValue: String? get() = stringCache.getOrNull()
 
     override val value: Result<Any?> get() = valCache.get()
 
     companion object {
-        private fun <T : KProperty<*>, R> T.asAccessible(action: (T) -> R): R {
+        @SuppressWarnings("DEPRECATED")
+        private fun <R> Field.asAccessible(action: (Field) -> R): R {
             val wasAccessible = isAccessible
             isAccessible = true
             val res = action(this)
             isAccessible = wasAccessible
             return res
         }
     }
-}
 
-private class VariableStateCache<T>(
-    val equalityChecker: (T, T) -> Boolean = { x, y -> x == y },
-    val calculate: (T?) -> T
-) {
-    private var cachedVal: T? = null
-    private var shouldRenew: Boolean = true
+    private class VariableStateCache<T>(
+        val equalityChecker: (T, T) -> Boolean = { x, y -> x == y },
+        val calculate: (T?) -> T
+    ) {
+        private var cachedVal: T? = null
+        private var shouldRenew: Boolean = true
 
-    fun getOrNull(): T? {
-        return if (shouldRenew) {
-            calculate(cachedVal).also {
-                cachedVal = it
-                shouldRenew = false
+        fun getOrNull(): T? {
+            return if (shouldRenew) {
+                calculate(cachedVal).also {
+                    cachedVal = it
+                    shouldRenew = false
+                }
+            } else {
+                cachedVal
             }
-        } else {
-            cachedVal
         }
-    }
 
-    fun get(): T = getOrNull()!!
+        fun get(): T = getOrNull()!!
 
-    fun update() {
-        shouldRenew = true
-    }
+        fun update() {
+            shouldRenew = true
+        }
 
-    fun forceUpdate(): Boolean {
-        val oldVal = getOrNull()
-        update()
-        val newVal = get()
-        return oldVal != null && equalityChecker(oldVal, newVal)
+        fun forceUpdate(): Boolean {
+            val oldVal = getOrNull()
+            update()
+            val newVal = get()
+            return oldVal != null && equalityChecker(oldVal, newVal)
+        }
     }
 }

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

@@ -20,12 +20,88 @@ 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: SerializableTypeInfo = SerializableTypeInfo(),
+    val value: String? = null,
+    val isContainer: Boolean = false,
+    val stateId: String = ""
+) {
+    // todo: not null
+    val fieldDescriptor: MutableMap<String, SerializedVariablesState?> = mutableMapOf()
+    override fun equals(other: Any?): Boolean {
+        if (this === other) return true
+        if (javaClass != other?.javaClass) return false
+
+        other as SerializedVariablesState
+
+        if (type != other.type) return false
+        if (value != other.value) return false
+        if (isContainer != other.isContainer) return false
+
+        return true
+    }
+
+    override fun hashCode(): Int {
+        var result = type.hashCode()
+        result = 31 * result + (value?.hashCode() ?: 0)
+        result = 31 * result + isContainer.hashCode()
+        return result
+    }
+}
+
+@Serializable
+class SerializationReply(
+    val cell_id: Int = 1,
+    val descriptorsState: Map<String, SerializedVariablesState> = emptyMap(),
+    val comm_id: String = ""
+)
+
 @Serializable
 class EvaluatedSnippetMetadata(
     val newClasspath: Classpath = emptyList(),
     val compiledData: SerializedCompiledScriptsData = SerializedCompiledScriptsData.EMPTY,
     val newImports: List<String> = emptyList(),
-    val evaluatedVariablesState: Map<String, String?> = mutableMapOf()
+    val evaluatedVariablesState: Map<String, SerializedVariablesState> = emptyMap()
 ) {
     companion object {
         val EMPTY = EvaluatedSnippetMetadata()

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

@@ -12,6 +12,7 @@ import org.jetbrains.kotlinx.jupyter.api.RenderersProcessor
 import org.jetbrains.kotlinx.jupyter.api.ResultsAccessor
 import org.jetbrains.kotlinx.jupyter.api.VariableState
 import org.jetbrains.kotlinx.jupyter.api.libraries.LibraryResolutionRequest
+import org.jetbrains.kotlinx.jupyter.repl.InternalEvaluator
 import org.jetbrains.kotlinx.jupyter.repl.impl.SharedReplContext
 
 class DisplayResultWrapper private constructor(
@@ -133,7 +134,9 @@ class NotebookImpl(
 
     private val history = arrayListOf<CodeCellImpl>()
     private var mainCellCreated = false
+    private val _unchangedVariables: MutableSet<String> = mutableSetOf()
 
+    val unchangedVariables: Set<String> get() = _unchangedVariables
     val displays = DisplayContainerImpl()
 
     override fun getAllDisplays(): List<DisplayResultWithCell> {
@@ -149,6 +152,11 @@ class NotebookImpl(
     override val jreInfo: JREInfoProvider
         get() = JavaRuntime
 
+    fun updateVariablesState(evaluator: InternalEvaluator) {
+        _unchangedVariables.clear()
+        _unchangedVariables.addAll(evaluator.getUnchangedVariables())
+    }
+
     fun variablesReportAsHTML(): String {
         return generateHTMLVarsReport(variablesState)
     }

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

@@ -23,6 +23,7 @@ import kotlinx.serialization.json.decodeFromJsonElement
 import kotlinx.serialization.json.encodeToJsonElement
 import kotlinx.serialization.json.jsonObject
 import kotlinx.serialization.serializer
+import org.jetbrains.kotlinx.jupyter.compiler.util.SerializedVariablesState
 import org.jetbrains.kotlinx.jupyter.config.LanguageInfo
 import org.jetbrains.kotlinx.jupyter.exceptions.ReplException
 import kotlin.reflect.KClass
@@ -87,7 +88,11 @@ enum class MessageType(val contentClass: KClass<out MessageContent>) {
     COMM_CLOSE(CommClose::class),
 
     LIST_ERRORS_REQUEST(ListErrorsRequest::class),
-    LIST_ERRORS_REPLY(ListErrorsReply::class);
+    LIST_ERRORS_REPLY(ListErrorsReply::class),
+
+    // from Serialization_Request
+    VARIABLES_VIEW_REQUEST(SerializationRequest::class),
+    VARIABLES_VIEW_REPLY(SerializationReply::class);
 
     val type: String
         get() = name.lowercase()
@@ -552,6 +557,22 @@ class ListErrorsReply(
     val errors: List<ScriptDiagnostic>
 ) : MessageContent()
 
+@Serializable
+class SerializationRequest(
+    val cellId: Int,
+    val descriptorsState: Map<String, SerializedVariablesState>,
+    val topLevelDescriptorName: String = "",
+    val pathToDescriptor: List<String> = emptyList(),
+    val commId: String = ""
+) : MessageContent()
+
+@Serializable
+class SerializationReply(
+    val cell_id: Int = 1,
+    val descriptorsState: Map<String, SerializedVariablesState> = emptyMap(),
+    val comm_id: String = ""
+) : MessageContent()
+
 @Serializable(MessageDataSerializer::class)
 data class MessageData(
     val header: MessageHeader? = null,

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,
@@ -307,6 +309,27 @@ fun JupyterConnection.Socket.shellMessagesHandler(msg: Message, repl: ReplForJup
         is CommInfoRequest -> {
             sendWrapped(msg, makeReplyMessage(msg, MessageType.COMM_INFO_REPLY, content = CommInfoReply(mapOf())))
         }
+        is CommOpen -> {
+            if (!content.targetName.equals("kotlin_serialization", ignoreCase = true)) {
+                send(makeReplyMessage(msg, MessageType.NONE))
+                return
+            }
+            log.debug("Message type in CommOpen: $msg, ${msg.type}")
+            val data = content.data ?: return sendWrapped(msg, makeReplyMessage(msg, MessageType.VARIABLES_VIEW_REPLY))
+            if (data.isEmpty()) return sendWrapped(msg, makeReplyMessage(msg, MessageType.VARIABLES_VIEW_REPLY))
+            log.debug("Message data: $data")
+            val messageContent = getVariablesDescriptorsFromJson(data)
+            connection.launchJob {
+                repl.serializeVariables(
+                    messageContent.topLevelDescriptorName,
+                    messageContent.descriptorsState,
+                    content.commId,
+                    messageContent.pathToDescriptor
+                ) { result ->
+                    sendWrapped(msg, makeReplyMessage(msg, MessageType.COMM_MSG, content = result))
+                }
+            }
+        }
         is CompleteRequest -> {
             connection.launchJob {
                 repl.complete(content.code, content.cursorPos) { result ->
@@ -321,6 +344,17 @@ fun JupyterConnection.Socket.shellMessagesHandler(msg: Message, repl: ReplForJup
                 }
             }
         }
+        is SerializationRequest -> {
+            connection.launchJob {
+                if (content.topLevelDescriptorName.isNotEmpty()) {
+                    repl.serializeVariables(content.topLevelDescriptorName, content.descriptorsState, commID = content.commId, content.pathToDescriptor) { result ->
+                        sendWrapped(msg, makeReplyMessage(msg, MessageType.VARIABLES_VIEW_REPLY, content = result))
+                    }
+                } else {
+                    sendWrapped(msg, makeReplyMessage(msg, MessageType.VARIABLES_VIEW_REPLY, content = null))
+                }
+            }
+        }
         is IsCompleteRequest -> {
             // We are in console mode, so switch off all the loggers
             if (mainLoggerLevel() != Level.OFF) disableLogging()
@@ -513,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
+}