Get unchangedVariables from notebook to reflect after execution state properly

Author: nikolay-egorov

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

@@ -52,6 +52,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) {
@@ -168,7 +172,6 @@ internal class InternalEvaluatorImpl(
         }.toHashSet()
         val ans = mutableMapOf<String, VariableStateImpl>()
         fields.forEach { property ->
-//            if (property.name.startsWith("script$")) return@forEach
             if (!memberKPropertiesNames.contains(property.name)) return@forEach
 
             val state = VariableStateImpl(property, cellClassInstance)

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

@@ -26,6 +26,7 @@ import kotlin.test.assertEquals
 import kotlin.test.assertFails
 import kotlin.test.assertFailsWith
 import kotlin.test.assertFalse
+import kotlin.test.assertNotEquals
 import kotlin.test.assertNotNull
 import kotlin.test.fail
 
@@ -1056,4 +1057,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)
+    }
 }