Add gradle property jupyter.serialization.enabled;fix recursive structures with jvmFields

nikolay-egorov · nikolay-egorov · commit e050cb6f2b1e · 2021-07-28T17:31:15.000+03:00
diff --git a/build.gradle.kts b/build.gradle.kts
@@ -12,6 +12,8 @@ plugins {
 
 val deploy: Configuration by configurations.creating
 
+val serializationFlagProperty = "jupyter.serialization.enabled"
+
 deploy.apply {
     exclude("org.jetbrains.kotlinx", "kotlinx-serialization-json-jvm")
     exclude("org.jetbrains.kotlinx", "kotlinx-serialization-core-jvm")
@@ -108,7 +110,8 @@ tasks {
 
             "junit.jupiter.execution.parallel.enabled" to doParallelTesting.toString() as Any,
             "junit.jupiter.execution.parallel.mode.default" to "concurrent",
-            "junit.jupiter.execution.parallel.mode.classes.default" to "concurrent"
+            "junit.jupiter.execution.parallel.mode.classes.default" to "concurrent",
+            serializationFlagProperty to "true"
         )
     }
 
diff --git a/jupyter-lib/shared-compiler/src/main/kotlin/org/jetbrains/kotlinx/jupyter/compiler/util/serializedCompiledScript.kt b/jupyter-lib/shared-compiler/src/main/kotlin/org/jetbrains/kotlinx/jupyter/compiler/util/serializedCompiledScript.kt
@@ -22,7 +22,6 @@ data class SerializedCompiledScriptsData(
 
 @Serializable
 data class SerializedVariablesState(
-    val name: String = "",
     val type: String = "",
     val value: String? = null,
     val isContainer: Boolean = false
diff --git a/src/main/kotlin/org/jetbrains/kotlinx/jupyter/repl.kt b/src/main/kotlin/org/jetbrains/kotlinx/jupyter/repl.kt
@@ -424,8 +424,6 @@ class ReplForJupyterImpl(
             notebook.updateVariablesState(internalEvaluator)
             // printVars()
             // printUsagesInfo(jupyterId, cellVariables[jupyterId - 1])
-            val entry = notebook.variablesState.entries.lastOrNull()
-            val serializedVarsState =  variablesSerializer.serializeVariableState(jupyterId - 1, entry?.key, entry?.value)
             val serializedData = variablesSerializer.serializeVariables(jupyterId - 1, notebook.variablesState)
 
             EvalResult(rendered, EvaluatedSnippetMetadata(newClasspath, compiledData, newImports, serializedData))
diff --git a/src/main/kotlin/org/jetbrains/kotlinx/jupyter/serializationUtils.kt b/src/main/kotlin/org/jetbrains/kotlinx/jupyter/serializationUtils.kt
@@ -15,39 +15,16 @@ typealias FieldDescriptor = Map<String, SerializedVariablesState?>
 typealias MutableFieldDescriptor = MutableMap<String, SerializedVariablesState?>
 typealias PropertiesData = Collection<KProperty1<out Any, *>>
 
-data class ProcessedSerializedVarsState(
+class ProcessedSerializedVarsState(
     val serializedVariablesState: SerializedVariablesState,
     val propertiesData: PropertiesData?,
     val jvmOnlyFields: Array<Field>? = null
-) {
-    // autogenerated
-    override fun equals(other: Any?): Boolean {
-        if (this === other) return true
-        if (javaClass != other?.javaClass) return false
-
-        other as ProcessedSerializedVarsState
-
-        if (serializedVariablesState != other.serializedVariablesState) return false
-        if (propertiesData != other.propertiesData) return false
-        if (jvmOnlyFields != null) {
-            if (other.jvmOnlyFields == null) return false
-            if (!jvmOnlyFields.contentEquals(other.jvmOnlyFields)) return false
-        } else if (other.jvmOnlyFields != null) return false
-
-        return true
-    }
-
-    override fun hashCode(): Int {
-        var result = serializedVariablesState.hashCode()
-        result = 31 * result + (propertiesData?.hashCode() ?: 0)
-        result = 31 * result + (jvmOnlyFields?.contentHashCode() ?: 0)
-        return result
-    }
-}
+)
 
 data class ProcessedDescriptorsState(
-    // perhaps, better tp make SerializedVariablesState -> PropertiesData?
-    val processedSerializedVarsState: MutableMap<SerializedVariablesState, PropertiesData?> = mutableMapOf(),
+    val processedSerializedVarsToKProperties: MutableMap<SerializedVariablesState, PropertiesData?> = mutableMapOf(),
+    // do we need this? Probably, not
+    // val processedSerializedVarsToJvmFields: MutableMap<SerializedVariablesState, Array<Field>?> = mutableMapOf(),
     val instancesPerState: MutableMap<SerializedVariablesState, Any?> = mutableMapOf()
 )
 
@@ -68,7 +45,11 @@ class VariablesSerializer(private val serializationStep: Int = 2, private val se
      */
     private val computedDescriptorsPerCell: MutableMap<Int, ProcessedDescriptorsState> = mutableMapOf()
 
+    private val isSerializationActive: Boolean = System.getProperty(serializationEnvProperty)?.toBooleanStrictOrNull() ?: true
+
     fun serializeVariables(cellId: Int, variablesState: Map<String, VariableState>): Map<String, SerializedVariablesState> {
+        if (!isSerializationActive) return emptyMap()
+
         if (seenObjectsPerCell.containsKey(cellId)) {
             seenObjectsPerCell[cellId]!!.clear()
         }
@@ -80,6 +61,8 @@ class VariablesSerializer(private val serializationStep: Int = 2, private val se
     }
 
     fun doIncrementalSerialization(cellId: Int, propertyName: String, serializedVariablesState: SerializedVariablesState): SerializedVariablesState {
+        if (!isSerializationActive) return serializedVariablesState
+
         val cellDescriptors = computedDescriptorsPerCell[cellId] ?: return serializedVariablesState
         return updateVariableState(cellId, propertyName, cellDescriptors, serializedVariablesState)
     }
@@ -94,9 +77,8 @@ class VariablesSerializer(private val serializationStep: Int = 2, private val se
         evaluatedDescriptorsState: ProcessedDescriptorsState,
         serializedVariablesState: SerializedVariablesState
     ): SerializedVariablesState {
-        // TODO: consider anonymous class as well and fix bug with state
         val value = evaluatedDescriptorsState.instancesPerState[serializedVariablesState]
-        val propertiesData = evaluatedDescriptorsState.processedSerializedVarsState[serializedVariablesState]
+        val propertiesData = evaluatedDescriptorsState.processedSerializedVarsToKProperties[serializedVariablesState]
         if (propertiesData == null && value != null && (value::class.java.isArray || value::class.java.isMemberClass)) {
             return serializeVariableState(cellId, propertyName, propertiesData, value, false)
         }
@@ -108,21 +90,25 @@ class VariablesSerializer(private val serializationStep: Int = 2, private val se
     }
 
     fun serializeVariableState(cellId: Int, name: String?, variableState: VariableState?, isOverride: Boolean = true): SerializedVariablesState {
-        if (variableState == null || name == null) return SerializedVariablesState()
+        if (!isSerializationActive || variableState == null || name == null) return SerializedVariablesState()
         return serializeVariableState(cellId, name, variableState.property, variableState.value, isOverride)
     }
 
-    fun serializeVariableState(cellId: Int, name: String, property: Field?, value: Any?, isOverride: Boolean = true): SerializedVariablesState {
+    private fun serializeVariableState(cellId: Int, name: String, property: Field?, value: Any?, isOverride: Boolean = true): SerializedVariablesState {
         val processedData = createSerializeVariableState(name, getSimpleTypeNameFrom(property, value), value)
         return doActualSerialization(cellId, processedData, value, isOverride)
     }
 
-    fun serializeVariableState(cellId: Int, name: String, property: KProperty<*>, value: Any?, isOverride: Boolean = true): SerializedVariablesState {
+    private fun serializeVariableState(cellId: Int, name: String, property: KProperty<*>, value: Any?, isOverride: Boolean = true): SerializedVariablesState {
         val processedData = createSerializeVariableState(name, getSimpleTypeNameFrom(property, value), value)
         return doActualSerialization(cellId, processedData, value, isOverride)
     }
 
     private fun doActualSerialization(cellId: Int, processedData: ProcessedSerializedVarsState, value: Any?, isOverride: Boolean = true): SerializedVariablesState {
+        fun isCanBeComputed(fieldDescriptors: MutableMap<String, SerializedVariablesState?>): Boolean {
+            return (fieldDescriptors.isEmpty() || (fieldDescriptors.isNotEmpty() && fieldDescriptors.entries.first().value?.fieldDescriptor!!.isEmpty()))
+        }
+
         val serializedVersion = processedData.serializedVariablesState
 
         seenObjectsPerCell.putIfAbsent(cellId, mutableMapOf())
@@ -131,13 +117,27 @@ class VariablesSerializer(private val serializationStep: Int = 2, private val se
             computedDescriptorsPerCell[cellId] = ProcessedDescriptorsState()
         }
         val currentCellDescriptors = computedDescriptorsPerCell[cellId]
-        currentCellDescriptors!!.processedSerializedVarsState[serializedVersion] = processedData.propertiesData
+        currentCellDescriptors!!.processedSerializedVarsToKProperties[serializedVersion] = processedData.propertiesData
+//        currentCellDescriptors.processedSerializedVarsToJvmFields[serializedVersion] = processedData.jvmOnlyFields
 
         if (value != null) {
-            seenObjectsPerCell[cellId]!![value] = serializedVersion
+            seenObjectsPerCell[cellId]!!.putIfAbsent(value, serializedVersion)
         }
         if (serializedVersion.isContainer) {
-            iterateThroughContainerMembers(cellId, value, serializedVersion.fieldDescriptor, currentCellDescriptors.processedSerializedVarsState[serializedVersion])
+            // check for seen
+            if (seenObjectsPerCell[cellId]!!.containsKey(value)) {
+                val previouslySerializedState = seenObjectsPerCell[cellId]!![value] ?: return processedData.serializedVariablesState
+                serializedVersion.fieldDescriptor += previouslySerializedState.fieldDescriptor
+                if (isCanBeComputed(serializedVersion.fieldDescriptor)) {
+                    iterateThroughContainerMembers(cellId, value, serializedVersion.fieldDescriptor, currentCellDescriptors.processedSerializedVarsToKProperties[serializedVersion])
+                }
+            } else {
+                // add jvm descriptors
+                processedData.jvmOnlyFields?.forEach {
+                    serializedVersion.fieldDescriptor[it.name] = serializeVariableState(cellId, it.name, it, value)
+                }
+                iterateThroughContainerMembers(cellId, value, serializedVersion.fieldDescriptor, currentCellDescriptors.processedSerializedVarsToKProperties[serializedVersion])
+            }
         }
         return processedData.serializedVariablesState
     }
@@ -150,6 +150,7 @@ class VariablesSerializer(private val serializationStep: Int = 2, private val se
         seenObjectsPerCell.putIfAbsent(cellId, mutableMapOf())
         val seenObjectsPerCell = seenObjectsPerCell[cellId]
         val currentCellDescriptors = computedDescriptorsPerCell[cellId]!!
+        // ok, it's a copy on the left for some reason
         val instancesPerState = currentCellDescriptors.instancesPerState
 
         for (it in properties) {
@@ -181,6 +182,26 @@ class VariablesSerializer(private val serializationStep: Int = 2, private val se
         }
 
         val isArrayType = checkCreateForPossibleArray(callInstance, descriptor, serializedIteration)
+        computedDescriptorsPerCell[cellId]!!.instancesPerState += instancesPerState
+
+        // check for seen
+        // for now it's O(c*n)
+        if (serializedIteration.isEmpty()) {
+            val processedVars = computedDescriptorsPerCell[cellId]!!.processedSerializedVarsToKProperties
+            descriptor.forEach { (_, state) ->
+                if (processedVars.containsKey(state)) {
+                    processedVars.entries.firstOrNull {
+                        val itValue = it.key
+                        if (itValue.value == state?.value && itValue.type == state?.value) {
+                            state?.fieldDescriptor?.put(itValue.type, itValue)
+                            true
+                        } else {
+                            false
+                        }
+                    }
+                }
+            }
+        }
 
         serializedIteration.forEach {
             val serializedVariablesState = it.value.serializedVariablesState
@@ -227,8 +248,12 @@ class VariablesSerializer(private val serializationStep: Int = 2, private val se
                 // update state with JVMFields
                 it.value.jvmOnlyFields?.forEach { field ->
                     serializedVariablesState.fieldDescriptor[field.name] = serializeVariableState(cellId, field.name, field, neededCallInstance)
-                    instancesPerState[serializedVariablesState] = neededCallInstance
+                    val properInstance = serializedVariablesState.fieldDescriptor[field.name]
+                    instancesPerState[properInstance!!] = neededCallInstance
+                    seenObjectsPerCell?.set(neededCallInstance!!, serializedVariablesState)
                 }
+                computedDescriptorsPerCell[cellId]!!.instancesPerState += instancesPerState
+//                computedDescriptorsPerCell[cellId]!!.processedSerializedVarsToJvmFields[serializedVariablesState] = it.value.jvmOnlyFields
                 iterateThroughContainerMembers(
                     cellId,
                     neededCallInstance,
@@ -286,7 +311,7 @@ class VariablesSerializer(private val serializationStep: Int = 2, private val se
             simpleTypeName.toString()
         }
 
-        val serializedVariablesState = SerializedVariablesState(name, type, getProperString(value), isContainer)
+        val serializedVariablesState = SerializedVariablesState(type, getProperString(value), isContainer)
 
         return ProcessedSerializedVarsState(serializedVariablesState, membersProperties, jvmFields)
     }
@@ -324,6 +349,10 @@ class VariablesSerializer(private val serializationStep: Int = 2, private val se
             false
         }
     }
+
+    companion object {
+        const val serializationEnvProperty = "jupyter.serialization.enabled"
+    }
 }
 
 fun getProperString(value: Any?): String {
diff --git a/src/main/kotlin/org/jetbrains/kotlinx/jupyter/util.kt b/src/main/kotlin/org/jetbrains/kotlinx/jupyter/util.kt
@@ -2,6 +2,7 @@ package org.jetbrains.kotlinx.jupyter
 
 import org.jetbrains.kotlinx.jupyter.api.bufferedImageRenderer
 import org.jetbrains.kotlinx.jupyter.codegen.ResultsRenderersProcessor
+import org.jetbrains.kotlinx.jupyter.compiler.util.SerializedVariablesState
 import org.jetbrains.kotlinx.jupyter.compiler.util.SourceCodeImpl
 import kotlin.script.experimental.api.ScriptDiagnostic
 import kotlin.script.experimental.api.SourceCode
@@ -69,6 +70,10 @@ fun ResultsRenderersProcessor.registerDefaultRenderers() {
     register(bufferedImageRenderer)
 }
 
+fun Map<String, SerializedVariablesState>.getValuesToString(): Map<String, String?> {
+    return this.mapValues { it.value.value }
+}
+
 /**
  * Stores info about where a variable Y was declared and info about what are they at the address X.
  * K: key, stands for a way of addressing variables, e.g. address.
diff --git a/src/test/kotlin/org/jetbrains/kotlinx/jupyter/test/repl/ReplTests.kt b/src/test/kotlin/org/jetbrains/kotlinx/jupyter/test/repl/ReplTests.kt
@@ -9,6 +9,7 @@ import org.jetbrains.kotlinx.jupyter.api.VariableStateImpl
 import org.jetbrains.kotlinx.jupyter.exceptions.ReplCompilerException
 import org.jetbrains.kotlinx.jupyter.generateDiagnostic
 import org.jetbrains.kotlinx.jupyter.generateDiagnosticFromAbsolute
+import org.jetbrains.kotlinx.jupyter.getValuesToString
 import org.jetbrains.kotlinx.jupyter.repl.CompletionResult
 import org.jetbrains.kotlinx.jupyter.repl.ListErrorsResult
 import org.jetbrains.kotlinx.jupyter.test.getOrFail
@@ -463,7 +464,7 @@ class ReplVarsTest : AbstractSingleReplTest() {
             "y" to "0",
             "z" to "47"
         )
-        assertEquals(res.metadata.evaluatedVariablesState.mapValues { it.value.value }, varsUpdate)
+        assertEquals(res.metadata.evaluatedVariablesState.getValuesToString(), varsUpdate)
         assertFalse(repl.notebook.variablesState.isEmpty())
         val varsState = repl.notebook.variablesState
         assertEquals("1", varsState.getStringValue("x"))
@@ -866,7 +867,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
 
         val serializer = repl.variablesSerializer
 
-        val newData = serializer.doIncrementalSerialization(0, descriptor["i"]!!.name, descriptor["i"]!!)
+        val newData = serializer.doIncrementalSerialization(0, "i", descriptor["i"]!!)
         val a = 1
     }
 
@@ -887,7 +888,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
         val actualContainer = listData.fieldDescriptor.entries.first().value!!
         val serializer = repl.variablesSerializer
 
-        val newData = serializer.doIncrementalSerialization(0, actualContainer.name, actualContainer)
+        val newData = serializer.doIncrementalSerialization(0, listData.fieldDescriptor.entries.first().key, actualContainer)
         var receivedDescriptor = newData.fieldDescriptor
         assertEquals(2, receivedDescriptor.size)
         assertTrue(receivedDescriptor.containsKey("size"))

Original file line number	Diff line number	Diff line change
`@@ -12,6 +12,8 @@ plugins {`
`12`	`12`
`13`	`13`	`val deploy: Configuration by configurations.creating`
`14`	`14`
	`15`	`+val serializationFlagProperty = "jupyter.serialization.enabled"`
	`16`	`+`
`15`	`17`	`deploy.apply {`
`16`	`18`	`exclude("org.jetbrains.kotlinx", "kotlinx-serialization-json-jvm")`
`17`	`19`	`exclude("org.jetbrains.kotlinx", "kotlinx-serialization-core-jvm")`
`@@ -108,7 +110,8 @@ tasks {`
`108`	`110`
`109`	`111`	`"junit.jupiter.execution.parallel.enabled" to doParallelTesting.toString() as Any,`
`110`	`112`	`"junit.jupiter.execution.parallel.mode.default" to "concurrent",`
`111`		`- "junit.jupiter.execution.parallel.mode.classes.default" to "concurrent"`
	`113`	`+ "junit.jupiter.execution.parallel.mode.classes.default" to "concurrent",`
	`114`	`+ serializationFlagProperty to "true"`
`112`	`115`	`)`
`113`	`116`	`}`
`114`	`117`