Add support for set container; make standard Arrays have values descriptors right away, not through 'size' and 'data' fields

Author: nikolay-egorov

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

@@ -424,7 +424,7 @@ class ReplForJupyterImpl(
             notebook.updateVariablesState(internalEvaluator)
             // printVars()
             // printUsagesInfo(jupyterId, cellVariables[jupyterId - 1])
-            val serializedData = variablesSerializer.serializeVariables(jupyterId - 1, notebook.variablesState)
+            val serializedData = variablesSerializer.serializeVariables(jupyterId - 1, notebook.variablesState, notebook.unchangedVariables())
 
             EvalResult(
                 result.result.value,

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

@@ -117,6 +117,17 @@ class VariablesSerializer(private val serializationDepth: Int = 2, private val s
             } catch (ex: Exception) {null}
             val serializedVersion = SerializedVariablesState(simpleTypeName, getProperString(value), true)
             val descriptors = serializedVersion.fieldDescriptor
+
+            // only for set case
+            if (simpleTypeName == "Set" && kProperties == null) {
+                value as Set<*>
+                val size = value.size
+                descriptors["size"] = createSerializeVariableState(
+                    "size", "Int", size
+                ).serializedVariablesState
+                descriptors.addDescriptor(value, "data")
+            }
+
             if (isDescriptorsNeeded) {
                 kProperties?.forEach { prop ->
                     val name = prop.name
@@ -206,7 +217,8 @@ class VariablesSerializer(private val serializationDepth: Int = 2, private val s
         Float::class.java,
         Double::class.java,
         Char::class.java,
-        Boolean::class.java
+        Boolean::class.java,
+        String::class.java
     )
 
     /**
@@ -219,9 +231,9 @@ class VariablesSerializer(private val serializationDepth: Int = 2, private val s
     /**
      * Cache for not recomputing unchanged variables
      */
-    val serializedVariablesCache: MutableMap<String, SerializedVariablesState> = mutableMapOf()
+    private val serializedVariablesCache: MutableMap<String, SerializedVariablesState> = mutableMapOf()
 
-    fun serializeVariables(cellId: Int, variablesState: Map<String, VariableState>): Map<String, SerializedVariablesState> {
+    fun serializeVariables(cellId: Int, variablesState: Map<String, VariableState>, unchangedVariables: Set<String>): Map<String, SerializedVariablesState> {
         if (!isSerializationActive) return emptyMap()
 
         if (seenObjectsPerCell.containsKey(cellId)) {
@@ -231,7 +243,12 @@ class VariablesSerializer(private val serializationDepth: Int = 2, private val s
             return emptyMap()
         }
         currentSerializeCount = 0
-        return variablesState.mapValues { serializeVariableState(cellId, it.key, it.value) }
+
+        val neededEntries = variablesState.filterKeys { unchangedVariables.contains(it) }
+
+        val serializedData = neededEntries.mapValues { serializeVariableState(cellId, it.key, it.value) }
+        serializedVariablesCache.putAll(serializedData)
+        return serializedVariablesCache
     }
 
     fun doIncrementalSerialization(cellId: Int, propertyName: String, serializedVariablesState: SerializedVariablesState): SerializedVariablesState {
@@ -307,6 +324,7 @@ class VariablesSerializer(private val serializationDepth: Int = 2, private val s
             }
             iterateThroughContainerMembers(cellId, value.objectInstance, serializedVersion.fieldDescriptor, currentCellDescriptors.processedSerializedVarsToJavaProperties[serializedVersion])
         }
+
         return processedData.serializedVariablesState
     }
 
@@ -318,7 +336,19 @@ class VariablesSerializer(private val serializationDepth: Int = 2, private val s
         kProperties: KPropertiesData? = null,
         currentDepth: Int = 0
     ) {
-        if ((properties == null && kProperties == null) || callInstance == null || currentDepth >= serializationDepth) return
+        fun iterateAndStoreValues(callInstance: Any, descriptorsState: MutableMap<String, SerializedVariablesState?>) {
+            if (callInstance is Collection<*>) {
+                callInstance.forEach {
+                    descriptorsState.addDescriptor(it)
+                }
+            } else if (callInstance is Array<*>) {
+                callInstance.forEach {
+                    descriptorsState.addDescriptor(it)
+                }
+            }
+        }
+
+        if ((properties == null && kProperties == null && callInstance !is Set<*>) || callInstance == null || currentDepth >= serializationDepth) return
 
         val serializedIteration = mutableMapOf<String, ProcessedSerializedVarsState>()
 
@@ -353,6 +383,17 @@ class VariablesSerializer(private val serializationDepth: Int = 2, private val s
         val isArrayType = checkForPossibleArray(callInstance)
         computedDescriptorsPerCell[cellId]!!.instancesPerState += instancesPerState
 
+        if (descriptor.size == 2 && descriptor.containsKey("data")) {
+            val listData = descriptor["data"]?.fieldDescriptor ?: return
+            if (descriptor.containsKey("size") && descriptor["size"]?.value == "null") {
+                descriptor.remove("size")
+                descriptor.remove("data")
+                iterateAndStoreValues(callInstance, descriptor)
+            } else {
+                iterateAndStoreValues(callInstance, listData)
+            }
+        }
+
         serializedIteration.forEach {
             val serializedVariablesState = it.value.serializedVariablesState
             val name = it.key
@@ -379,7 +420,7 @@ class VariablesSerializer(private val serializationDepth: Int = 2, private val s
                 )
             }
         }
-
+        /*
         if (descriptor.size == 2 && descriptor.containsKey("data")) {
             val listData = descriptor["data"]?.fieldDescriptor ?: return
             if (callInstance is Collection<*>) {
@@ -391,7 +432,7 @@ class VariablesSerializer(private val serializationDepth: Int = 2, private val s
                     listData.addDescriptor(it)
                 }
             }
-        }
+        }*/
     }
 
     /**
@@ -480,7 +521,7 @@ class VariablesSerializer(private val serializationDepth: Int = 2, private val s
         }
 
         val isContainer = if (membersProperties != null) (
-            !primitiveWrappersSet.contains(javaClass) && membersProperties.isNotEmpty() || value::class.java.isArray || (javaClass.isMemberClass)
+            !primitiveWrappersSet.contains(javaClass) && membersProperties.isNotEmpty() || value is Set<*> || value::class.java.isArray || javaClass.isMemberClass
             ) else false
         val type = if (value != null && value::class.java.isArray) {
             "Array"

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

@@ -668,18 +668,23 @@ class ReplVarsTest : AbstractSingleReplTest() {
             """
             val x = 124
             private var f = "abcd"
-            """.trimIndent()
+            """.trimIndent(),
+            jupyterId = 1
         )
         val state = repl.notebook.cellVariables
         assertTrue(state.isNotEmpty())
+
+        // f is not accessible from here
         eval(
             """
             private var z = 1
             z += x
-            """.trimIndent()
+            """.trimIndent(),
+            jupyterId = 1
         )
         assertTrue(state.isNotEmpty())
 
+        // TODO discuss if we really want this
         val setOfCell = setOf("z", "f", "x")
         assertTrue(state.containsValue(setOfCell))
     }
@@ -809,6 +814,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
 
         val serializer = repl.variablesSerializer
         val newData = serializer.doIncrementalSerialization(0, "data", actualContainer)
+        val a = 1
     }
 
     @Test
@@ -892,13 +898,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
         val serializer = repl.variablesSerializer
 
         val newData = serializer.doIncrementalSerialization(0, listData.fieldDescriptor.entries.first().key, actualContainer)
-        var receivedDescriptor = newData.fieldDescriptor
-        assertEquals(2, receivedDescriptor.size)
-        assertTrue(receivedDescriptor.containsKey("size"))
-
-        val innerList = receivedDescriptor.entries.last().value!!
-        assertTrue(innerList.isContainer)
-        receivedDescriptor = innerList.fieldDescriptor
+        val receivedDescriptor = newData.fieldDescriptor
         assertEquals(4, receivedDescriptor.size)
 
         var values = 1
@@ -946,13 +946,61 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
         assertEquals(3, newDescriptor["data"]!!.fieldDescriptor.size)
         val ansSet = mutableSetOf("a", "b", "c")
         newDescriptor["data"]!!.fieldDescriptor.forEach { (_, state) ->
-            assertTrue(state!!.isContainer)
+            assertFalse(state!!.isContainer)
             assertTrue(ansSet.contains(state.value))
             ansSet.remove(state.value)
         }
         assertTrue(ansSet.isEmpty())
     }
 
+
+    @Test
+    fun testSetContainer() {
+        var res = eval(
+            """
+            val x = setOf("a", "b", "cc", "c")
+            """.trimIndent(),
+            jupyterId = 1
+        )
+        var varsData = res.metadata.evaluatedVariablesState
+        assertEquals(1, varsData.size)
+        assertTrue(varsData.containsKey("x"))
+
+        var setData = varsData["x"]!!
+        assertTrue(setData.isContainer)
+        assertEquals(2, setData.fieldDescriptor.size)
+        var setDescriptors = setData.fieldDescriptor
+        assertEquals("4", setDescriptors["size"]!!.value)
+        assertTrue(setDescriptors["data"]!!.isContainer)
+        assertEquals(4, setDescriptors["data"]!!.fieldDescriptor.size)
+        assertEquals("a", setDescriptors["data"]!!.fieldDescriptor["a"]!!.value)
+        assertTrue(setDescriptors["data"]!!.fieldDescriptor.containsKey("b"))
+        assertTrue(setDescriptors["data"]!!.fieldDescriptor.containsKey("cc"))
+        assertTrue(setDescriptors["data"]!!.fieldDescriptor.containsKey("c"))
+
+        res = eval(
+            """
+            val c = mutableSetOf("a", "b", "cc", "c")
+            """.trimIndent(),
+            jupyterId = 2
+        )
+        varsData = res.metadata.evaluatedVariablesState
+        assertEquals(2, varsData.size)
+        assertTrue(varsData.containsKey("c"))
+
+        setData = varsData["c"]!!
+        assertTrue(setData.isContainer)
+        assertEquals(2, setData.fieldDescriptor.size)
+        setDescriptors = setData.fieldDescriptor
+        assertEquals("4", setDescriptors["size"]!!.value)
+        assertTrue(setDescriptors["data"]!!.isContainer)
+        assertEquals(4, setDescriptors["data"]!!.fieldDescriptor.size)
+        assertEquals("a", setDescriptors["data"]!!.fieldDescriptor["a"]!!.value)
+        assertTrue(setDescriptors["data"]!!.fieldDescriptor.containsKey("b"))
+        assertTrue(setDescriptors["data"]!!.fieldDescriptor.containsKey("cc"))
+        assertTrue(setDescriptors["data"]!!.fieldDescriptor.containsKey("c"))
+    }
+
     @Test
     fun testSerializationMessage() {
         val res = eval(
@@ -975,9 +1023,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
 
                 val innerList = data.entries.last().value
                 assertTrue(innerList.isContainer)
-                var receivedDescriptor = innerList.fieldDescriptor
-                assertEquals(2, receivedDescriptor.size)
-                receivedDescriptor = receivedDescriptor.entries.last().value!!.fieldDescriptor
+                val receivedDescriptor = innerList.fieldDescriptor
 
                 assertEquals(4, receivedDescriptor.size)
                 var values = 1
@@ -997,9 +1043,8 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
 
                 val innerList = data.entries.last().value
                 assertTrue(innerList.isContainer)
-                var receivedDescriptor = innerList.fieldDescriptor
-                assertEquals(2, receivedDescriptor.size)
-                receivedDescriptor = receivedDescriptor.entries.last().value!!.fieldDescriptor
+                val receivedDescriptor = innerList.fieldDescriptor
+
 
                 assertEquals(4, receivedDescriptor.size)
                 var values = 1
@@ -1050,5 +1095,26 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
         )
         assertTrue(state.isNotEmpty())
         assertEquals(state, setOfPrevCell)
+
+
+        eval(
+            """
+            private val x = 341
+            protected val z = "abcd"
+            """.trimIndent(),
+            jupyterId = 1
+        )
+        assertTrue(state.isEmpty())
+
+        eval(
+            """
+            private val x = "abcd"
+            var f = 47
+            internal val z = 47
+            """.trimIndent(),
+            jupyterId = 1
+        )
+        assertTrue(state.isNotEmpty())
+        assertEquals(setOfPrevCell, state)
     }
 }