Add possibility to use serialization request with only top-level descriptor name, not cellID

Author: nikolay-egorov

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

@@ -5,15 +5,13 @@ import kotlin.reflect.jvm.isAccessible
 import java.lang.reflect.Field
 
 interface VariableState {
-//    val property: KProperty<*>
     val property: Field
     val scriptInstance: Any?
     val stringValue: String?
     val value: Result<Any?>
 }
 
 data class VariableStateImpl(
-//    override val property: KProperty1<Any, *>,
     override val property: Field,
     override val scriptInstance: Any,
 ) : VariableState {

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

@@ -30,6 +30,12 @@ data class SerializedVariablesState(
     val fieldDescriptor: MutableMap<String, SerializedVariablesState?> = mutableMapOf()
 }
 
+@Serializable
+class SerializationReply(
+    val cellId: Int = 1,
+    val descriptorsState: Map<String, SerializedVariablesState> = emptyMap()
+)
+
 @Serializable
 class EvaluatedSnippetMetadata(
     val newClasspath: Classpath = emptyList(),

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

@@ -133,6 +133,7 @@ class NotebookImpl(
 
     private val history = arrayListOf<CodeCellImpl>()
     private var mainCellCreated = false
+    private val unchangedVariables: MutableSet<String> = mutableSetOf()
 
     val displays = DisplayContainerImpl()
 
@@ -149,6 +150,19 @@ class NotebookImpl(
     override val jreInfo: JREInfoProvider
         get() = JavaRuntime
 
+    fun updateVariablesState(evaluator: InternalEvaluator) {
+        variablesState += evaluator.variablesHolder
+        currentCellVariables = evaluator.cellVariables
+        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/message_types.kt

@@ -93,9 +93,6 @@ enum class MessageType(val contentClass: KClass<out MessageContent>) {
     SERIALIZATION_REQUEST(SerializationRequest::class),
     SERIALIZATION_REPLY(SerializationReply::class);
 
-    // TODO: add custom commands
-    // this custom message should be supported on client-side. either JS or Idea Plugin
-
     val type: String
         get() = name.lowercase()
 }
@@ -562,12 +559,13 @@ class ListErrorsReply(
 @Serializable
 class SerializationRequest(
     val cellId: Int,
-    val descriptorsState: Map<String, SerializedVariablesState>
+    val descriptorsState: Map<String, SerializedVariablesState>,
+    val topLevelDescriptorName: String = ""
 ) : MessageContent()
 
 @Serializable
 class SerializationReply(
-    val cellId: Int,
+    val cellId: Int = 1,
     val descriptorsState: Map<String, SerializedVariablesState> = emptyMap()
 ) : MessageContent()
 

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

@@ -307,6 +307,21 @@ 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.commId.equals(MessageType.SERIALIZATION_REQUEST.name, 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.SERIALIZATION_REPLY))
+
+            val messageContent = getVariablesDescriptorsFromJson(data)
+            GlobalScope.launch(Dispatchers.Default) {
+                repl.serializeVariables(messageContent.topLevelDescriptorName, messageContent.descriptorsState) { result ->
+                    sendWrapped(msg, makeReplyMessage(msg, MessageType.COMM_OPEN, content = result))
+                }
+            }
+        }
         is CompleteRequest -> {
             connection.launchJob {
                 repl.complete(content.code, content.cursorPos) { result ->
@@ -323,8 +338,14 @@ fun JupyterConnection.Socket.shellMessagesHandler(msg: Message, repl: ReplForJup
         }
         is SerializationRequest -> {
             GlobalScope.launch(Dispatchers.Default) {
-                repl.serializeVariables(content.cellId, content.descriptorsState) { result ->
-                    sendWrapped(msg, makeReplyMessage(msg, MessageType.SERIALIZATION_REPLY, content = result))
+                if (content.topLevelDescriptorName.isNotEmpty()) {
+                    repl.serializeVariables(content.topLevelDescriptorName, content.descriptorsState) { result ->
+                        sendWrapped(msg, makeReplyMessage(msg, MessageType.SERIALIZATION_REPLY, content = result))
+                    }
+                } else {
+                    repl.serializeVariables(content.cellId, content.descriptorsState) { result ->
+                        sendWrapped(msg, makeReplyMessage(msg, MessageType.SERIALIZATION_REPLY, content = result))
+                    }
                 }
             }
         }

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

@@ -136,6 +136,8 @@ interface ReplForJupyter {
 
     suspend fun serializeVariables(cellId: Int, descriptorsState: Map<String, SerializedVariablesState>, callback: (SerializationReply) -> Unit)
 
+    suspend fun serializeVariables(topLevelVarName: String, descriptorsState: Map<String, SerializedVariablesState>, callback: (SerializationReply) -> Unit)
+
     val homeDir: File?
 
     val currentClasspath: Collection<String>
@@ -554,15 +556,26 @@ class ReplForJupyterImpl(
 
     private val serializationQueue = LockQueue<SerializationReply, SerializationArgs>()
     override suspend fun serializeVariables(cellId: Int, descriptorsState: Map<String, SerializedVariablesState>, callback: (SerializationReply) -> Unit) {
-        doWithLock(SerializationArgs(cellId, descriptorsState, callback), serializationQueue, SerializationReply(cellId, descriptorsState), ::doSerializeVariables)
+        doWithLock(SerializationArgs(descriptorsState, cellId = cellId, callback = callback), serializationQueue, SerializationReply(cellId, descriptorsState), ::doSerializeVariables)
+    }
+
+    override suspend fun serializeVariables(topLevelVarName: String, descriptorsState: Map<String, SerializedVariablesState>, callback: (SerializationReply) -> Unit) {
+        doWithLock(SerializationArgs(descriptorsState, topLevelVarName = topLevelVarName, callback = callback), serializationQueue, SerializationReply(), ::doSerializeVariables)
     }
 
     private fun doSerializeVariables(args: SerializationArgs): SerializationReply {
         val resultMap = mutableMapOf<String, SerializedVariablesState>()
+        val cellId = if (args.cellId != -1) args.cellId else {
+            val watcherInfo = internalEvaluator.findVariableCell(args.topLevelVarName) + 1
+            val finalAns = if (watcherInfo == - 1) 1 else watcherInfo
+            finalAns
+        }
         args.descriptorsState.forEach { (name, state) ->
-            resultMap[name] = variablesSerializer.doIncrementalSerialization(args.cellId - 1, name, state)
+            resultMap[name] = variablesSerializer.doIncrementalSerialization(cellId - 1, name, state)
         }
-        return SerializationReply(args.cellId, resultMap)
+        log.debug("Serialization cellID: $cellId")
+        log.debug("Serialization answer: ${resultMap.entries.first().value.fieldDescriptor}")
+        return SerializationReply(cellId, resultMap)
     }
 
 
@@ -599,11 +612,13 @@ class ReplForJupyterImpl(
         LockQueueArgs<ListErrorsResult>
 
     private data class SerializationArgs(
-        val cellId: Int,
         val descriptorsState: Map<String, SerializedVariablesState>,
+        var cellId: Int = -1,
+        val topLevelVarName: String = "",
         override val callback: (SerializationReply) -> Unit
     ) : LockQueueArgs<SerializationReply>
 
+
     @JvmInline
     private value class LockQueue<T, Args : LockQueueArgs<T>>(
         private val args: AtomicReference<Args?> = AtomicReference()

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

@@ -30,4 +30,14 @@ interface InternalEvaluator {
      * returns empty data or null
      */
     fun popAddedCompiledScripts(): SerializedCompiledScriptsData = SerializedCompiledScriptsData.EMPTY
+
+    /**
+     * Get a cellId where a particular variable is declared
+     */
+    fun findVariableCell(variableName: String): Int
+
+    /**
+     * Returns a set of unaffected variables after execution
+     */
+    fun getUnchangedVariables(): Set<String>
 }

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

@@ -52,6 +52,10 @@ internal class InternalEvaluatorImpl(
         return SerializedCompiledScriptsData(scripts)
     }
 
+    override fun findVariableCell(variableName: String): Int {
+        return variablesWatcher.findDeclarationAddress(variableName) ?: -1
+    }
+
     override var writeCompiledClasses: Boolean
         get() = classWriter != null
         set(value) {

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

@@ -1,8 +1,14 @@
 package org.jetbrains.kotlinx.jupyter
 
+import kotlinx.serialization.Serializable
+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.SerializedVariablesState
 import java.lang.reflect.Field
+import kotlin.contracts.ExperimentalContracts
+import kotlin.contracts.contract
 import kotlin.reflect.KClass
 import kotlin.reflect.KProperty
 import kotlin.reflect.KProperty1
@@ -23,6 +29,16 @@ enum class PropertiesType {
     MIXED
 }
 
+@Serializable
+data class SerializedCommMessageContent(
+    val topLevelDescriptorName: String,
+    val descriptorsState: Map<String, SerializedVariablesState>
+)
+
+fun getVariablesDescriptorsFromJson(json: JsonObject): SerializedCommMessageContent {
+    return Json.decodeFromJsonElement<SerializedCommMessageContent>(json)
+}
+
 class ProcessedSerializedVarsState(
     val serializedVariablesState: SerializedVariablesState,
     val propertiesData: PropertiesData? = null,
@@ -94,9 +110,11 @@ class VariablesSerializer(private val serializationDepth: Int = 2, private val s
                 } == true
             }
 
-            val kProperties = if (value != null) value::class.declaredMemberProperties else {
-                null
-            }
+            val kProperties = try {
+                if (value != null) value::class.declaredMemberProperties else {
+                    null
+                }
+            } catch (ex: Exception) {null}
             val serializedVersion = SerializedVariablesState(simpleTypeName, getProperString(value), true)
             val descriptors = serializedVersion.fieldDescriptor
             if (isDescriptorsNeeded) {
@@ -198,6 +216,11 @@ class VariablesSerializer(private val serializationDepth: Int = 2, private val s
 
     private val isSerializationActive: Boolean = System.getProperty(serializationSystemProperty)?.toBooleanStrictOrNull() ?: true
 
+    /**
+     * Cache for not recomputing unchanged variables
+     */
+    val serializedVariablesCache: MutableMap<String, SerializedVariablesState> = mutableMapOf()
+
     fun serializeVariables(cellId: Int, variablesState: Map<String, VariableState>): Map<String, SerializedVariablesState> {
         if (!isSerializationActive) return emptyMap()
 
@@ -375,6 +398,7 @@ class VariablesSerializer(private val serializationDepth: Int = 2, private val s
      * Really wanted to use contracts here, but all usages should be provided with this annotation and,
      * perhaps, it may be a big overhead
      */
+    @OptIn(ExperimentalContracts::class)
     private fun iterateThrough(
         elem: Any,
         seenObjectsPerCell: MutableMap<RuntimeObjectWrapper, SerializedVariablesState>?,
@@ -383,6 +407,10 @@ class VariablesSerializer(private val serializationDepth: Int = 2, private val s
         instancesPerState: MutableMap<SerializedVariablesState, Any?>,
         callInstance: Any
     ) {
+        contract {
+            returns() implies (elem is Field || elem is KProperty1<*, *>)
+        }
+
         val name = if (elem is Field) elem.name else (elem as KProperty1<Any, *>).name
         val value = if (elem is Field) tryGetValueFromProperty(elem, callInstance).toObjectWrapper()
         else {
@@ -491,6 +519,8 @@ class VariablesSerializer(private val serializationDepth: Int = 2, private val s
         return value
     }
 
+    // use of Java 9 required
+    @SuppressWarnings("DEPRECATION")
     private fun tryGetValueFromProperty(property: Field, callInstance: Any): Any? {
         // some fields may be optimized out like array size. Thus, calling it.isAccessible would return error
         val canAccess = try {

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

@@ -73,7 +73,7 @@ fun ResultsRenderersProcessor.registerDefaultRenderers() {
  * Stores info about where a variable Y was declared and info about what are they at the address X.
  * K: key, stands for a way of addressing variables, e.g. address.
  * V: value, from Variable, choose any suitable type for your variable reference.
- * Default: T=Int, V=String
+ * Default: K=Int, V=String
  */
 class VariablesUsagesPerCellWatcher<K : Any, V : Any> {
     val cellVariables = mutableMapOf<K, MutableSet<V>>()
@@ -106,5 +106,7 @@ class VariablesUsagesPerCellWatcher<K : Any, V : Any> {
         }
     }
 
+    fun findDeclarationAddress(variableRef: V) = variablesDeclarationInfo[variableRef]
+
     fun ensureStorageCreation(address: K) = cellVariables.putIfAbsent(address, mutableSetOf())
 }