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

@@ -3,20 +3,20 @@ package org.jetbrains.kotlinx.jupyter.api
 import java.lang.reflect.Field
 
 interface VariableState {
-//    val property: KProperty<*>
     val property: Field
     val scriptInstance: Any?
     val stringValue: String?
     val value: Any?
 }
 
 data class VariableStateImpl(
-//    override val property: KProperty1<Any, *>,
     override val property: Field,
     override val scriptInstance: Any,
 ) : VariableState {
     private var cachedValue: Any? = null
 
+    // use of Java 9 required
+    @SuppressWarnings("DEPRECATION")
     fun update() {
         val wasAccessible = property.isAccessible
         property.isAccessible = true

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/message_types.kt

@@ -92,9 +92,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()
 }
@@ -580,12 +577,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

@@ -305,6 +305,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 -> {
             GlobalScope.launch(Dispatchers.Default) {
                 repl.complete(content.code, content.cursorPos) { result ->
@@ -321,8 +336,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

@@ -120,6 +120,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>
@@ -520,15 +522,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)
     }
 
 
@@ -565,8 +578,9 @@ 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>
 

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

@@ -30,4 +30,9 @@ interface InternalEvaluator {
      * returns empty data or null
      */
     fun popAddedCompiledScripts(): SerializedCompiledScriptsData = SerializedCompiledScriptsData.EMPTY
+
+    /**
+     * Get a cellId where particular variable is declared
+     */
+    fun findVariableCell(variableName: String): Int
 }

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

@@ -48,6 +48,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,
@@ -375,6 +391,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 +400,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 +512,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())
 }

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

@@ -805,6 +805,9 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
         assertEquals(2, actualContainer.fieldDescriptor.size)
         assertTrue(actualContainer.isContainer)
         assertEquals(listOf(1, 2, 3, 4).toString().substring(1, actualContainer.value!!.length + 1), actualContainer.value)
+
+        val serializer = repl.variablesSerializer
+        val newData = serializer.doIncrementalSerialization(0, "data", actualContainer)
     }
 
     @Test
@@ -904,6 +907,9 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
             assertTrue(state.isContainer)
             assertEquals("${values++}", state.value)
         }
+
+        val depthMostNode = actualContainer.fieldDescriptor.entries.first { it.value!!.isContainer }
+        val serializationAns = serializer.doIncrementalSerialization(0, depthMostNode.key, depthMostNode.value!!)
     }
 
     @Test
@@ -982,5 +988,27 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
                 }
             }
         }
+
+        runBlocking {
+            repl.serializeVariables("x", mapOf(propertyName to actualContainer)) { result ->
+                val data = result.descriptorsState
+                assertTrue(data.isNotEmpty())
+
+                val innerList = data.entries.last().value
+                assertTrue(innerList.isContainer)
+                var receivedDescriptor = innerList.fieldDescriptor
+                assertEquals(2, receivedDescriptor.size)
+                receivedDescriptor = receivedDescriptor.entries.last().value!!.fieldDescriptor
+
+                assertEquals(4, receivedDescriptor.size)
+                var values = 1
+                receivedDescriptor.forEach { (_, state) ->
+                    val fieldDescriptor = state!!.fieldDescriptor
+                    assertEquals(0, fieldDescriptor.size)
+                    assertTrue(state.isContainer)
+                    assertEquals("${values++}", state.value)
+                }
+            }
+        }
     }
 }