Get unchangedVariables from notebook to reflect after execution state properly

Author: nikolay-egorov

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

@@ -56,6 +56,10 @@ internal class InternalEvaluatorImpl(
         return variablesWatcher.findDeclarationAddress(variableName) ?: -1
     }
 
+    override fun getUnchangedVariables(): Set<String> {
+        return variablesWatcher.getUnchangedVariables()
+    }
+
     override var writeCompiledClasses: Boolean
         get() = classWriter != null
         set(value) {

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

@@ -83,6 +83,9 @@ class VariablesUsagesPerCellWatcher<K : Any, V : Any> {
      */
     private val variablesDeclarationInfo: MutableMap<V, K> = mutableMapOf()
 
+    private val unchangedVariables: MutableSet<V> = mutableSetOf()
+//    private val unchangedVariables: MutableSet<V> = mutableSetOf()
+
     fun addDeclaration(address: K, variableRef: V) {
         ensureStorageCreation(address)
 
@@ -91,21 +94,35 @@ class VariablesUsagesPerCellWatcher<K : Any, V : Any> {
             val oldCellId = variablesDeclarationInfo[variableRef]
             if (oldCellId != address) {
                 cellVariables[oldCellId]?.remove(variableRef)
+                unchangedVariables.remove(variableRef)
             }
+        } else {
+            unchangedVariables.add(variableRef)
         }
         variablesDeclarationInfo[variableRef] = address
         cellVariables[address]?.add(variableRef)
     }
 
-    fun addUsage(address: K, variableRef: V) = cellVariables[address]?.add(variableRef)
+    fun addUsage(address: K, variableRef: V) {
+        cellVariables[address]?.add(variableRef)
+        if (variablesDeclarationInfo[variableRef] != address) {
+            unchangedVariables.remove(variableRef)
+        }
+    }
 
     fun removeOldUsages(newAddress: K) {
         // remove known modifying usages in this cell
         cellVariables[newAddress]?.removeIf {
-            variablesDeclarationInfo[it] != newAddress
+            val predicate = variablesDeclarationInfo[it] != newAddress
+            if (predicate) {
+                unchangedVariables.add(it)
+            }
+            predicate
         }
     }
 
+    fun getUnchangedVariables(): Set<V> = unchangedVariables
+
     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

@@ -696,6 +696,93 @@ class ReplVarsTest : AbstractSingleReplTest() {
     fun testOutVarRendering() {
         eval("Out").resultValue.shouldNotBeNull()
     }
+
+
+    @Test
+    fun testSeparatePrivateDefsUsage() {
+        eval(
+            """
+            private val x = "abcd"
+            private var f = 47
+            """.trimIndent(),
+            jupyterId = 1
+        )
+        val state = repl.notebook.cellVariables
+        assertTrue(state[0]!!.contains("x"))
+
+        eval(
+            """
+            val x = 341
+            private var f = "abcd"
+            """.trimIndent(),
+            jupyterId = 2
+        )
+        assertTrue(state.isNotEmpty())
+        assertTrue(state[0]!!.isEmpty())
+        assertTrue(state[1]!!.contains("x"))
+
+        val setOfPrevCell = setOf<String>()
+        val setOfNextCell = setOf("x", "f")
+        assertEquals(state[0], setOfPrevCell)
+        assertEquals(state[1], setOfNextCell)
+    }
+
+    @Test
+    fun testRecursiveVarsState() {
+        eval(
+            """
+            val l = mutableListOf<Any>()
+            l.add(listOf(l))
+            
+            val m = mapOf(1 to l)
+            
+            val z = setOf(1, 2, 4)
+            """.trimIndent(),
+            jupyterId = 1
+        )
+        val state = repl.notebook.variablesState
+        assertTrue(state.contains("l"))
+        assertTrue(state.contains("m"))
+        assertTrue(state.contains("z"))
+
+        assertEquals("ArrayList: recursive structure", state["l"]!!.stringValue)
+        assertTrue(state["m"]!!.stringValue!!.contains(" recursive structure"))
+        assertEquals("[1, 2, 4]", state["z"]!!.stringValue)
+    }
+
+    @Test
+    fun testSeparatePrivateCellsUsage() {
+        eval(
+            """
+            private val x = "abcd"
+            var f = 47
+            internal val z = 47
+            """.trimIndent(),
+            jupyterId = 1
+        )
+        val state = repl.notebook.cellVariables
+        assertTrue(state[0]!!.contains("x"))
+        assertTrue(state[0]!!.contains("z"))
+
+        eval(
+            """
+            private val x = 341
+            f += x
+            protected val z = "abcd"
+            """.trimIndent(),
+            jupyterId = 2
+        )
+        assertTrue(state.isNotEmpty())
+        assertTrue(state[0]!!.isNotEmpty())
+        assertFalse(state[0]!!.contains("x"))
+        assertFalse(state[0]!!.contains("z"))
+        assertTrue(state[1]!!.contains("x"))
+
+        val setOfPrevCell = setOf("f")
+        val setOfNextCell = setOf("x", "f", "z")
+        assertEquals(state[0], setOfPrevCell)
+        assertEquals(state[1], setOfNextCell)
+    }
 }
 
 class ReplVarsSerializationTest : AbstractSingleReplTest() {
@@ -933,4 +1020,43 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
             }
         }
     }
+
+
+    @Test
+    fun testUnchangedVariables() {
+        eval(
+            """
+            private val x = "abcd"
+            var f = 47
+            internal val z = 47
+            """.trimIndent(),
+            jupyterId = 1
+        )
+        val state = repl.notebook.unchangedVariables()
+        val setOfCell = setOf("x", "f", "z")
+        assertTrue(state.isNotEmpty())
+        assertEquals(setOfCell, state)
+
+        eval(
+            """
+            private val x = 341
+            f += x
+            protected val z = "abcd"
+            """.trimIndent(),
+            jupyterId = 2
+        )
+        assertTrue(state.isEmpty())
+        val setOfPrevCell = setOf("f")
+        assertNotEquals(setOfCell, setOfPrevCell)
+
+        eval(
+            """
+            private val x = 341
+            protected val z = "abcd"
+            """.trimIndent(),
+            jupyterId = 2
+        )
+        assertTrue(state.isNotEmpty())
+        assertEquals(state, setOfPrevCell)
+    }
 }