Get unchangedVariables from notebook to reflect after execution state properly

Author: nikolay-egorov

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

@@ -123,6 +123,7 @@ class NotebookImpl(
     private var currentCellVariables = mapOf<Int, Set<String>>()
     private val history = arrayListOf<CodeCellImpl>()
     private var mainCellCreated = false
+    private val unchangedVariables: MutableSet<String> = mutableSetOf()
 
     val displays = DisplayContainerImpl()
 
@@ -142,12 +143,16 @@ class NotebookImpl(
     fun updateVariablesState(evaluator: InternalEvaluator) {
         variablesState += evaluator.variablesHolder
         currentCellVariables = evaluator.cellVariables
+        unchangedVariables.clear()
+        unchangedVariables.addAll(evaluator.getUnchangedVariables())
     }
 
     fun updateVariablesState(varsStateUpdate: Map<String, VariableState>) {
         variablesState += varsStateUpdate
     }
 
+    fun unchangedVariables(): Set<String> = unchangedVariables
+
     fun variablesReportAsHTML(): String {
         return generateHTMLVarsReport(variablesState)
     }

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

@@ -32,7 +32,12 @@ interface InternalEvaluator {
     fun popAddedCompiledScripts(): SerializedCompiledScriptsData = SerializedCompiledScriptsData.EMPTY
 
     /**
-     * Get a cellId where particular variable is declared
+     * Get a cellId where a particular variable is declared
      */
     fun findVariableCell(variableName: String): Int
+
+    /**
+     * Returns a set of unaffected variables after execution
+     */
+    fun getUnchangedVariables(): Set<String>
 }

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) {
@@ -170,7 +174,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/serializationUtils.kt

@@ -110,9 +110,11 @@ class VariablesSerializer(private val serializationDepth: Int = 2, private val s
                 } == true
             }
 
-            val kProperties = if (value != null) value::class.declaredMemberProperties else {
-                null
-            }
+            val kProperties = try {
+                if (value != null) value::class.declaredMemberProperties else {
+                    null
+                }
+            } catch (ex: Exception) {null}
             val serializedVersion = SerializedVariablesState(simpleTypeName, getProperString(value), true)
             val descriptors = serializedVersion.fieldDescriptor
             if (isDescriptorsNeeded) {
@@ -214,6 +216,11 @@ class VariablesSerializer(private val serializationDepth: Int = 2, private val s
 
     private val isSerializationActive: Boolean = System.getProperty(serializationSystemProperty)?.toBooleanStrictOrNull() ?: true
 
+    /**
+     * Cache for not recomputing unchanged variables
+     */
+    val serializedVariablesCache: MutableMap<String, SerializedVariablesState> = mutableMapOf()
+
     fun serializeVariables(cellId: Int, variablesState: Map<String, VariableState>): Map<String, SerializedVariablesState> {
         if (!isSerializationActive) return emptyMap()
 

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

@@ -23,6 +23,7 @@ import kotlin.script.experimental.api.SourceCode
 import kotlin.test.assertEquals
 import kotlin.test.assertFails
 import kotlin.test.assertFalse
+import kotlin.test.assertNotEquals
 import kotlin.test.fail
 
 class ReplTests : AbstractSingleReplTest() {
@@ -1011,4 +1012,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)
+    }
 }

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

@@ -1,6 +1,7 @@
 package org.jetbrains.kotlinx.jupyter.test.repl
 
 import org.jetbrains.kotlinx.jupyter.ReplForJupyterImpl
+import org.jetbrains.kotlinx.jupyter.VariablesUsagesPerCellWatcher
 import org.jetbrains.kotlinx.jupyter.api.Code
 import org.jetbrains.kotlinx.jupyter.api.FieldValue
 import org.jetbrains.kotlinx.jupyter.api.VariableState
@@ -48,6 +49,8 @@ internal class MockedInternalEvaluator : TrackedInternalEvaluator {
     override val variablesHolder = mutableMapOf<String, VariableState>()
     override val cellVariables = mutableMapOf<Int, MutableSet<String>>()
 
+    val variablesWatcher: VariablesUsagesPerCellWatcher<Int, String> = VariablesUsagesPerCellWatcher()
+
     override val results: List<Any?>
         get() = executedCodes.map { null }
 
@@ -64,6 +67,10 @@ internal class MockedInternalEvaluator : TrackedInternalEvaluator {
         }
         return -1
     }
+
+    override fun getUnchangedVariables(): Set<String> {
+        return variablesWatcher.getUnchangedVariables()
+    }
 }
 
 internal class TrackedInternalEvaluatorImpl(private val baseEvaluator: InternalEvaluator) : TrackedInternalEvaluator, InternalEvaluator by baseEvaluator {