Fixed bug with usageMap with var redefinition; some tests added

Author: nikolay-egorov

jupyter-lib/api/src/main/kotlin/org/jetbrains/kotlinx/jupyter/api/Notebook.kt

@@ -14,6 +14,12 @@ interface Notebook {
      */
     val variablesMap: MutableMap<String, VariableState>
 
+    /**
+     * Stores useful variables in a cell.
+     * Useful <==> declarations + modifying references
+     */
+    var usageMap: Map<Int, Set<String>>
+
     /**
      * Mapping allowing to get cell by execution number
      */

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

@@ -106,6 +106,8 @@ class NotebookImpl(
 
     override val variablesMap = mutableMapOf<String, VariableState>()
 
+    override var usageMap = mapOf<Int, Set<String>>()
+
     override fun getCell(id: Int): CodeCellImpl {
         return cells[id] ?: throw ArrayIndexOutOfBoundsException(
             "There is no cell with number '$id'"

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

@@ -381,6 +381,7 @@ class ReplForJupyterImpl(
             } ?: emptyList()
 
             notebook.updateVarsState(varsMap)
+            notebook.usageMap = usageMap
             printVars()
 //            printUsage(jupyterId, usageMap[jupyterId - 1]!!)
 

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

@@ -190,16 +190,22 @@ internal class InternalEvaluatorImpl(
             } else {
                 val stringVal = casted.get(cellKlassInstance).toString()
                 val state = VariableState(casted, cellKlassInstance, stringVal)
+
+                // redeclaration of any type
+                if (containingMap.containsKey(casted.name)) {
+                    val oldCellId = containingMap[casted.name]
+                    if (oldCellId != cellId) {
+                        usageMap[oldCellId]?.remove(casted.name)
+                    }
+                }
+                // it was val, now it's var
+                if (it is KMutableProperty1 && variablesCache.containsKey(casted.name)) {
+                    variablesCache.remove(casted.name)
+                }
+                containingMap[casted.name] = cellId
+                usageMap[cellId]?.add(casted.name)
                 // invariant with changes: cache --> varsMap, other way is not allowed
                 if (it !is KMutableProperty1) {
-                    // it is a replacement
-                    if (variablesCache.containsKey(casted.name)) {
-                        val oldCellID = containingMap[casted.name]
-                        usageMap[oldCellID]?.remove(casted.name)
-                    }
-                    containingMap[casted.name] = cellId
-                    usageMap[cellId]?.add(casted.name)
-
                     variablesCache[casted.name] = state
                     variablesMap[casted.name] = state
                     return@forEach

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

@@ -443,7 +443,7 @@ class ReplTests : AbstractSingleReplTest() {
     }
 
     @Test
-    fun testStateConsistency() {
+    fun testVarsStateConsistency() {
         assertTrue(repl.notebook.variablesMap.isEmpty())
         val res = eval(
             """
@@ -471,7 +471,7 @@ class ReplTests : AbstractSingleReplTest() {
     }
 
     @Test
-    fun testEmptyState() {
+    fun testVarsEmptyState() {
         val res = eval("3+2")
         val state = repl.notebook.variablesMap
         val strState = mutableMapOf<String, String>()
@@ -504,4 +504,151 @@ class ReplTests : AbstractSingleReplTest() {
         assertEquals("abc", varsState["y"]!!.stringValue)
         assertEquals("1", varsState["z"]!!.stringValue)
     }
+
+    @Test
+    fun testVarsCaptureSeparateCells() {
+        eval(
+            """
+            val x = 1 
+            val y = "abc"
+            val z = x
+            """.trimIndent()
+        )
+        val varsState = repl.notebook.variablesMap
+        assertTrue(varsState.isNotEmpty())
+        eval(
+            """
+            val x = "abc" 
+            var y = 123
+            val z = x
+            """.trimIndent(),
+            jupyterId = 1
+        )
+        assertTrue(varsState.isNotEmpty())
+        assertEquals(3, varsState.size)
+        assertEquals("abc", varsState["x"]!!.stringValue)
+        assertEquals("123", varsState["y"]!!.stringValue)
+        assertEquals("abc", varsState["z"]!!.stringValue)
+
+        eval(
+            """
+            val x = 1024 
+            y += 123
+            """.trimIndent(),
+            jupyterId = 2
+        )
+
+        assertTrue(varsState.isNotEmpty())
+        assertEquals(3, varsState.size)
+        assertEquals("1024", varsState["x"]!!.stringValue)
+        assertEquals("${123 * 2}", varsState["y"]!!.stringValue)
+        assertEquals("abc", varsState["z"]!!.stringValue)
+    }
+
+    @Test
+    fun testVarsUsageConsistency() {
+        eval("3+2")
+        val state = repl.notebook.usageMap
+        assertTrue(state.values.size == 1)
+        assertTrue(state.values.first().isEmpty())
+        val setOfNextCell = setOf<String>()
+        assertEquals(state.values.first(), setOfNextCell)
+    }
+
+    @Test
+    fun testVarsDefsUsage() {
+        eval(
+            """
+            val x = 1
+            val z = "abcd"
+            var f = 47
+            """.trimIndent()
+        )
+        val state = repl.notebook.usageMap
+        assertTrue(state.isNotEmpty())
+        assertTrue(state.values.first().isNotEmpty())
+        val setOfCell = setOf("z", "f", "x")
+        assertTrue(state.containsValue(setOfCell))
+    }
+
+    @Test
+    fun testVarsDefNRefUsage() {
+        eval(
+            """
+            val x = "abcd"
+            var f = 47
+            """.trimIndent()
+        )
+        val state = repl.notebook.usageMap
+        assertTrue(state.isNotEmpty())
+        eval(
+            """
+            val z = 1
+            f += f
+            """.trimIndent()
+        )
+        assertTrue(state.isNotEmpty())
+
+        val setOfCell = setOf("z", "f", "x")
+        assertTrue(state.containsValue(setOfCell))
+    }
+
+    @Test
+    fun testSeparateDefsUsage() {
+        eval(
+            """
+            val x = "abcd"
+            var f = 47
+            """.trimIndent(),
+            jupyterId = 1
+        )
+        val state = repl.notebook.usageMap
+        assertTrue(state[0]!!.contains("x"))
+
+        eval(
+            """
+            val x = 341
+            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 testSeparateCellsUsage() {
+        eval(
+            """
+            val x = "abcd"
+            var f = 47
+            """.trimIndent(),
+            jupyterId = 1
+        )
+        val state = repl.notebook.usageMap
+        assertTrue(state[0]!!.contains("x"))
+
+        eval(
+            """
+            val x = 341
+            f += x
+            """.trimIndent(),
+            jupyterId = 2
+        )
+        assertTrue(state.isNotEmpty())
+        assertTrue(state[0]!!.isNotEmpty())
+        assertFalse(state[0]!!.contains("x"))
+        assertTrue(state[1]!!.contains("x"))
+
+        val setOfPrevCell = setOf("f")
+        val setOfNextCell = setOf("x", "f")
+        assertEquals(state[0], setOfPrevCell)
+        assertEquals(state[1], setOfNextCell)
+    }
 }

src/test/kotlin/org/jetbrains/kotlinx/jupyter/test/testUtil.kt

@@ -168,6 +168,7 @@ object NotebookMock : Notebook {
     override val cellsList: Collection<CodeCell>
         get() = emptyList()
     override val variablesMap = mutableMapOf<String, VariableState>()
+    override var usageMap = mapOf<Int, Set<String>>()
 
     override fun getCell(id: Int): CodeCellImpl {
         return cells[id] ?: throw ArrayIndexOutOfBoundsException(