Delete meta-data of redeclared variables

Author: nikolay-egorov

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

@@ -402,6 +402,8 @@ class ReplForJupyterImpl(
 
             val compiledData: SerializedCompiledScriptsData
             val newImports: List<String>
+            val oldDeclarations: MutableMap<String, Int> = mutableMapOf()
+            oldDeclarations.putAll(internalEvaluator.getVariablesDeclarationInfo())
             val result = try {
                 log.debug("Current cell id: $jupyterId")
                 executor.execute(code, displayHandler, currentCellId = jupyterId - 1) { internalId, codeToExecute ->
@@ -431,7 +433,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, 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/repl/InternalEvaluator.kt

@@ -36,6 +36,8 @@ interface InternalEvaluator {
      */
     fun findVariableCell(variableName: String): Int
 
+    fun getVariablesDeclarationInfo(): Map<String, Int>
+
     /**
      * Returns a set of unaffected variables after execution
      */

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

@@ -52,6 +52,8 @@ internal class InternalEvaluatorImpl(
         return variablesWatcher.findDeclarationAddress(variableName) ?: -1
     }
 
+    override fun getVariablesDeclarationInfo(): Map<String, Int> = variablesWatcher.variablesDeclarationInfo
+
     override fun getUnchangedVariables(): Set<String> {
         return variablesWatcher.getUnchangedVariables()
     }

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

@@ -56,8 +56,7 @@ class ProcessedSerializedVarsState(
 data class ProcessedDescriptorsState(
     val processedSerializedVarsToJavaProperties: MutableMap<SerializedVariablesState, PropertiesData?> = mutableMapOf(),
     val processedSerializedVarsToKTProperties: MutableMap<SerializedVariablesState, KPropertiesData?> = mutableMapOf(),
-    val instancesPerState: MutableMap<SerializedVariablesState, Any?> = mutableMapOf(),
-    val parent: ProcessedDescriptorsState? = null
+    val instancesPerState: MutableMap<SerializedVariablesState, Any?> = mutableMapOf()
 )
 
 data class RuntimeObjectWrapper(
@@ -73,6 +72,7 @@ data class RuntimeObjectWrapper(
         return objectInstance === other
     }
 
+    // TODO: it's not changing after recreation
     override fun hashCode(): Int {
         return if (isRecursive) Random.nextInt() else objectInstance?.hashCode() ?: 0
     }
@@ -93,7 +93,7 @@ fun Any?.getToStringValue(isRecursive: Boolean = false): String {
 }
 
 fun Any?.getUniqueID(isRecursive: Boolean = false): String {
-    return if (this != null) {
+    return if (this != null && this !is Map.Entry<*, *>) {
         val hashCode = if (isRecursive) {
             Random.nextLong()
         } else {
@@ -197,8 +197,9 @@ class VariablesSerializer(
                 val isRecursive = stringedValue.contains(": recursive structure")
                 if (!isRecursive && simpleTypeName == "LinkedEntrySet") {
                     getProperEntrySetRepresentation(value)
-                } else
-                value.getUniqueID(isRecursive)
+                } else {
+                    value.getUniqueID(isRecursive)
+                }
             } else {
                 ""
             }
@@ -285,12 +286,12 @@ class VariablesSerializer(
     }
 
     /**
-     * Map of Map of seen objects.
-     * First Key: cellId
+     * Map of Map of seen objects related to a particular variable serialization
+     * First Key: topLevel variable Name
      * Second Key: actual value
      * Value: serialized VariableState
      */
-    private val seenObjectsPerCell: MutableMap<Int, MutableMap<RuntimeObjectWrapper, SerializedVariablesState>> = mutableMapOf()
+    private val seenObjectsPerVariable: MutableMap<String, MutableMap<RuntimeObjectWrapper, SerializedVariablesState>> = mutableMapOf()
 
     private var currentSerializeCount: Int = 0
 
@@ -355,18 +356,19 @@ class VariablesSerializer(
 
     override suspend fun clearStateInfo(currentState: Int) {
         computedDescriptorsPerCell.remove(currentState)
-        seenObjectsPerCell.remove(currentState)
+        // seenObjectsPerVariable.remove(currentState)
     }
 
     suspend fun tryValidateCache(currentCellId: Int, cellVariables: Map<Int, Set<String>>) {
         if (!isShouldRemove(currentCellId)) return
         clearOldData(currentCellId, cellVariables)
     }
 
-    fun serializeVariables(cellId: Int, variablesState: Map<String, VariableState>, variablesCells: Map<String, Int>, unchangedVariables: Set<String>): Map<String, SerializedVariablesState> {
+    fun serializeVariables(cellId: Int, variablesState: Map<String, VariableState>, oldDeclarations: Map<String, Int>, variablesCells: Map<String, Int>, unchangedVariables: Set<String>): Map<String, SerializedVariablesState> {
         fun removeNonExistingEntries() {
             val toRemoveSet = mutableSetOf<String>()
             serializedVariablesCache.forEach { (name, _) ->
+                // seems like this never gonna happen
                 if (!variablesState.containsKey(name)) {
                     toRemoveSet.add(name)
                 }
@@ -376,9 +378,6 @@ class VariablesSerializer(
 
         if (!isSerializationActive) return emptyMap()
 
-        if (seenObjectsPerCell.containsKey(cellId)) {
-            seenObjectsPerCell[cellId]!!.clear()
-        }
         if (variablesState.isEmpty()) {
             return emptyMap()
         }
@@ -388,21 +387,26 @@ class VariablesSerializer(
             if (wasRedeclared) {
                 removedFromSightVariables.remove(it)
             }
-            // todo: might consider self-recursive elements always to recompute since it's non comparable via strings
+            // TODO: might consider self-recursive elements always to recompute since it's non comparable via strings
             if (serializedVariablesCache.isEmpty()) {
                 true
-            } else
-            (!unchangedVariables.contains(it) || serializedVariablesCache[it]?.value != variablesState[it]?.stringValue) &&
-                !removedFromSightVariables.contains(it)
+            } else {
+                (!unchangedVariables.contains(it) || serializedVariablesCache[it]?.value != variablesState[it]?.stringValue) &&
+                    !removedFromSightVariables.contains(it)
+            }
         }
         log.debug("Variables state as is: $variablesState")
         log.debug("Serializing variables after filter: $neededEntries")
         log.debug("Unchanged variables: ${unchangedVariables - neededEntries.keys}")
 
         // remove previous data
-        // computedDescriptorsPerCell[cellId]?.instancesPerState?.clear()
         val serializedData = neededEntries.mapValues {
             val actualCell = variablesCells[it.key] ?: cellId
+            if (oldDeclarations.containsKey(it.key)) {
+                val oldCell = oldDeclarations[it.key]!!
+                computedDescriptorsPerCell[oldCell]?.remove(it.key)
+                seenObjectsPerVariable.remove(it.key)
+            }
             serializeVariableState(actualCell, it.key, it.value)
         }
 
@@ -467,10 +471,13 @@ class VariablesSerializer(
         return doActualSerialization(cellId, topLevelName, processedData, wrapper, isRecursive, isOverride)
     }
 
-    private fun doActualSerialization(cellId: Int, topLevelName:String, processedData: ProcessedSerializedVarsState, value: RuntimeObjectWrapper, isRecursive: Boolean, isOverride: Boolean = true): SerializedVariablesState {
+    private fun doActualSerialization(cellId: Int, topLevelName: String, processedData: ProcessedSerializedVarsState, value: RuntimeObjectWrapper, isRecursive: Boolean, isOverride: Boolean = true): SerializedVariablesState {
+        fun checkIsNotStandardDescriptor(descriptor: MutableMap<String, SerializedVariablesState?>): Boolean {
+            return descriptor.isNotEmpty() && !descriptor.containsKey("size") && !descriptor.containsKey("data")
+        }
         val serializedVersion = processedData.serializedVariablesState
 
-        seenObjectsPerCell.putIfAbsent(cellId, mutableMapOf())
+        seenObjectsPerVariable.putIfAbsent(topLevelName, mutableMapOf())
         computedDescriptorsPerCell.putIfAbsent(cellId, mutableMapOf())
 
         if (isOverride) {
@@ -482,17 +489,21 @@ class VariablesSerializer(
         }
         val currentCellDescriptors = computedDescriptorsPerCell[cellId]?.get(topLevelName)
         // TODO should we stack?
+        // i guess, not
         currentCellDescriptors!!.processedSerializedVarsToJavaProperties[serializedVersion] = processedData.propertiesData
         currentCellDescriptors.processedSerializedVarsToKTProperties[serializedVersion] = processedData.kPropertiesData
 
         if (value.objectInstance != null) {
-            seenObjectsPerCell[cellId]!!.putIfAbsent(value, serializedVersion)
+            seenObjectsPerVariable[topLevelName]!!.putIfAbsent(value, serializedVersion)
         }
         if (serializedVersion.isContainer) {
             // check for seen
-            if (seenObjectsPerCell[cellId]!!.containsKey(value)) {
-                val previouslySerializedState = seenObjectsPerCell[cellId]!![value] ?: return processedData.serializedVariablesState
+            if (seenObjectsPerVariable[topLevelName]!!.containsKey(value)) {
+                val previouslySerializedState = seenObjectsPerVariable[topLevelName]!![value] ?: return processedData.serializedVariablesState
                 serializedVersion.fieldDescriptor += previouslySerializedState.fieldDescriptor
+                if (checkIsNotStandardDescriptor(serializedVersion.fieldDescriptor)) {
+                    return serializedVersion
+                }
             }
             val type = processedData.propertiesType
             if (type == PropertiesType.KOTLIN) {
@@ -535,8 +546,8 @@ class VariablesSerializer(
 
         val serializedIteration = mutableMapOf<String, ProcessedSerializedVarsState>()
 
-        seenObjectsPerCell.putIfAbsent(cellId, mutableMapOf())
-        val seenObjectsPerCell = seenObjectsPerCell[cellId]
+        seenObjectsPerVariable.putIfAbsent(topLevelName, mutableMapOf())
+        val seenObjectsPerCell = seenObjectsPerVariable[topLevelName]
         val currentCellDescriptors = computedDescriptorsPerCell[cellId]!![topLevelName]!!
         // ok, it's a copy on the left for some reason
         val instancesPerState = currentCellDescriptors.instancesPerState
@@ -645,7 +656,7 @@ class VariablesSerializer(
             getSimpleTypeNameFrom(elem, value.objectInstance) ?: "null"
         }
         serializedIteration[name] = if (standardContainersUtilizer.isStandardType(simpleType)) {
-            // todo might add isRecursive
+            // TODO might add isRecursive
             standardContainersUtilizer.serializeContainer(simpleType, value.objectInstance, true)
         } else {
             createSerializeVariableState(name, simpleType, value)

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

@@ -81,7 +81,7 @@ class VariablesUsagesPerCellWatcher<K : Any, V : Any> {
     /**
      * Tells in which cell a variable was declared
      */
-    private val variablesDeclarationInfo: MutableMap<V, K> = mutableMapOf()
+    val variablesDeclarationInfo: MutableMap<V, K> = mutableMapOf()
 
     private val unchangedVariables: MutableSet<V> = mutableSetOf()
 

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

@@ -666,7 +666,7 @@ class ReplVarsTest : AbstractSingleReplTest() {
             """.trimIndent(),
             jupyterId = 1
         )
-        val state = repl.notebook.cellVariables
+        var state = repl.notebook.cellVariables
         assertTrue(state.isNotEmpty())
 
         // f is not accessible from here
@@ -677,10 +677,10 @@ class ReplVarsTest : AbstractSingleReplTest() {
             """.trimIndent(),
             jupyterId = 2
         )
+        state = repl.notebook.cellVariables
         assertTrue(state.isNotEmpty())
-
-        // TODO discuss if we really want "z", "f", "x"
-        val setOfCell = setOf("z")
+        // ignore primitive references precise check for Java > 8
+        val setOfCell = setOf("z", "x")
         assertTrue(state.containsValue(setOfCell))
     }
 
@@ -808,7 +808,7 @@ class ReplVarsTest : AbstractSingleReplTest() {
             jupyterId = 2
         ).metadata.evaluatedVariablesState
         val innerList = res["l"]!!.fieldDescriptor["elementData"]!!.fieldDescriptor["data"]
-        val newData = serializer.doIncrementalSerialization(0, "l","data", innerList!!)
+        val newData = serializer.doIncrementalSerialization(0, "l", "data", innerList!!)
         assertTrue(newData.isContainer)
         assertTrue(newData.fieldDescriptor.size > 4)
     }
@@ -860,17 +860,6 @@ class ReplVarsTest : AbstractSingleReplTest() {
         var state = repl.notebook.unchangedVariables()
         assertEquals(3, state.size)
 
-        eval(
-            """
-            private val x = "abcd"
-            internal val z = 47
-            """.trimIndent(),
-            jupyterId = 2
-        )
-        state = repl.notebook.unchangedVariables()
-        assertEquals(1, state.size)
-        assertTrue(state.contains("f"))
-
         eval(
             """
             private val x = "abcd"
@@ -889,7 +878,7 @@ class ReplVarsTest : AbstractSingleReplTest() {
             jupyterId = 3
         )
         state = repl.notebook.unchangedVariables()
-        assertEquals(2, state.size)
+        assertEquals(1, state.size)
     }
 }
 
@@ -924,7 +913,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)
+        val newData = serializer.doIncrementalSerialization(0, "x", "data", actualContainer)
     }
 
     @Test
@@ -1382,8 +1371,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
             """
             val x = listOf(1, 2, 3, 4)
             """.trimIndent(),
-        jupyterId = 2
+            jupyterId = 2
         ).metadata.evaluatedVariablesState
-        val a = 1
     }
 }

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

@@ -68,6 +68,10 @@ internal class MockedInternalEvaluator : TrackedInternalEvaluator {
         return -1
     }
 
+    override fun getVariablesDeclarationInfo(): Map<String, Int> {
+        return variablesWatcher.variablesDeclarationInfo
+    }
+
     override fun getUnchangedVariables(): Set<String> {
         return variablesWatcher.getUnchangedVariables()
     }