Skip to content

Commit 4d3df57

Browse files
nikolay-egorovnikolay-egorov
committed
Store all meta-data related to a top-level variable name
1 parent d9a6403 commit 4d3df57

File tree

4 files changed

+64
-53
lines changed

4 files changed

+64
-53
lines changed

‎src/main/kotlin/org/jetbrains/kotlinx/jupyter/protocol.kt

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -346,9 +346,7 @@ fun JupyterConnection.Socket.shellMessagesHandler(msg: Message, repl: ReplForJup
346346
sendWrapped(msg, makeReplyMessage(msg, MessageType.SERIALIZATION_REPLY, content = result))
347347
}
348348
} else {
349-
repl.serializeVariables(content.cellId, content.descriptorsState) { result ->
350-
sendWrapped(msg, makeReplyMessage(msg, MessageType.SERIALIZATION_REPLY, content = result))
351-
}
349+
sendWrapped(msg, makeReplyMessage(msg, MessageType.SERIALIZATION_REPLY, content = null))
352350
}
353351
}
354352
}

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -132,7 +132,7 @@ interface ReplForJupyter {
132132

133133
suspend fun listErrors(code: Code, callback: (ListErrorsResult) -> Unit)
134134

135-
suspend fun serializeVariables(cellId: Int, descriptorsState: Map<String, SerializedVariablesState>, callback: (SerializationReply) -> Unit)
135+
suspend fun serializeVariables(cellId: Int, topLevelVarName: String, descriptorsState: Map<String, SerializedVariablesState>, callback: (SerializationReply) -> Unit)
136136

137137
suspend fun serializeVariables(topLevelVarName: String, descriptorsState: Map<String, SerializedVariablesState>, pathToDescriptor: List<String> = emptyList(),
138138
callback: (SerializationReply) -> Unit)
@@ -546,8 +546,8 @@ class ReplForJupyterImpl(
546546
}
547547

548548
private val serializationQueue = LockQueue<SerializationReply, SerializationArgs>()
549-
override suspend fun serializeVariables(cellId: Int, descriptorsState: Map<String, SerializedVariablesState>, callback: (SerializationReply) -> Unit) {
550-
doWithLock(SerializationArgs(descriptorsState, cellId = cellId, callback = callback), serializationQueue, SerializationReply(cellId, descriptorsState), ::doSerializeVariables)
549+
override suspend fun serializeVariables(cellId: Int, topLevelVarName: String, descriptorsState: Map<String, SerializedVariablesState>, callback: (SerializationReply) -> Unit) {
550+
doWithLock(SerializationArgs(descriptorsState, cellId = cellId, topLevelVarName = topLevelVarName, callback = callback), serializationQueue, SerializationReply(cellId, descriptorsState), ::doSerializeVariables)
551551
}
552552

553553
override suspend fun serializeVariables(topLevelVarName: String, descriptorsState: Map<String, SerializedVariablesState>, pathToDescriptor: List<String>,
@@ -563,7 +563,7 @@ class ReplForJupyterImpl(
563563
finalAns
564564
}
565565
args.descriptorsState.forEach { (name, state) ->
566-
resultMap[name] = variablesSerializer.doIncrementalSerialization(cellId - 1, name, state, args.pathToDescriptor)
566+
resultMap[name] = variablesSerializer.doIncrementalSerialization(cellId - 1, args.topLevelVarName ,name, state, args.pathToDescriptor)
567567
}
568568
log.debug("Serialization cellID: $cellId")
569569
log.debug("Serialization answer: ${resultMap.entries.first().value.fieldDescriptor}")

‎src/main/kotlin/org/jetbrains/kotlinx/jupyter/serializationUtils.kt

Lines changed: 29 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,8 @@ class ProcessedSerializedVarsState(
5656
data class ProcessedDescriptorsState(
5757
val processedSerializedVarsToJavaProperties: MutableMap<SerializedVariablesState, PropertiesData?> = mutableMapOf(),
5858
val processedSerializedVarsToKTProperties: MutableMap<SerializedVariablesState, KPropertiesData?> = mutableMapOf(),
59-
val instancesPerState: MutableMap<SerializedVariablesState, Any?> = mutableMapOf()
59+
val instancesPerState: MutableMap<SerializedVariablesState, Any?> = mutableMapOf(),
60+
val parent: ProcessedDescriptorsState? = null
6061
)
6162

6263
data class RuntimeObjectWrapper(
@@ -276,16 +277,6 @@ class VariablesSerializer(
276277
*/
277278
if (descriptors.size == 1 && descriptors.entries.first().key == "size") {
278279
descriptors.addDescriptor(value, "data")
279-
/*
280-
if (value is Collection<*>) {
281-
value.forEach {
282-
iterateThrough(descriptors, it)
283-
}
284-
} else if (value is Array<*>) {
285-
value.forEach {
286-
iterateThrough(descriptors, it)
287-
}
288-
}*/
289280
}
290281
}
291282

@@ -319,9 +310,9 @@ class VariablesSerializer(
319310
)
320311

321312
/**
322-
* Stores info computed descriptors in a cell
313+
* Stores info computed descriptors in a cell starting from the very variable as a root
323314
*/
324-
private val computedDescriptorsPerCell: MutableMap<Int, ProcessedDescriptorsState> = mutableMapOf()
315+
private val computedDescriptorsPerCell: MutableMap<Int, MutableMap<String, ProcessedDescriptorsState>> = mutableMapOf()
325316

326317
private val isSerializationActive: Boolean = System.getProperty(serializationSystemProperty)?.toBooleanStrictOrNull() ?: true
327318

@@ -409,7 +400,7 @@ class VariablesSerializer(
409400
log.debug("Unchanged variables: ${unchangedVariables - neededEntries.keys}")
410401

411402
// remove previous data
412-
computedDescriptorsPerCell[cellId]?.instancesPerState?.clear()
403+
// computedDescriptorsPerCell[cellId]?.instancesPerState?.clear()
413404
val serializedData = neededEntries.mapValues {
414405
val actualCell = variablesCells[it.key] ?: cellId
415406
serializeVariableState(actualCell, it.key, it.value)
@@ -424,14 +415,15 @@ class VariablesSerializer(
424415

425416
fun doIncrementalSerialization(
426417
cellId: Int,
418+
topLevelName: String,
427419
propertyName: String,
428420
serializedVariablesState: SerializedVariablesState,
429421
pathToDescriptor: List<String> = emptyList()
430422
): SerializedVariablesState {
431423
if (!isSerializationActive) return serializedVariablesState
432424

433425
val cellDescriptors = computedDescriptorsPerCell[cellId] ?: return serializedVariablesState
434-
return updateVariableState(cellId, propertyName, cellDescriptors, serializedVariablesState)
426+
return updateVariableState(cellId, propertyName, cellDescriptors[topLevelName]!!, serializedVariablesState)
435427
}
436428

437429
/**
@@ -456,38 +448,39 @@ class VariablesSerializer(
456448
return serializeVariableState(cellId, propertyName, property, value, isRecursive = false, false)
457449
}
458450

459-
private fun serializeVariableState(cellId: Int, name: String?, variableState: VariableState?, isOverride: Boolean = true): SerializedVariablesState {
460-
if (!isSerializationActive || variableState == null || name == null) return SerializedVariablesState()
451+
private fun serializeVariableState(cellId: Int, topLevelName: String?, variableState: VariableState?, isOverride: Boolean = true): SerializedVariablesState {
452+
if (!isSerializationActive || variableState == null || topLevelName == null) return SerializedVariablesState()
461453
// force recursive check
462454
variableState.stringValue
463-
return serializeVariableState(cellId, name, variableState.property, variableState.value.getOrNull(), variableState.isRecursive, isOverride)
455+
return serializeVariableState(cellId, topLevelName, variableState.property, variableState.value.getOrNull(), variableState.isRecursive, isOverride)
464456
}
465457

466-
private fun serializeVariableState(cellId: Int, name: String, property: Field?, value: Any?, isRecursive: Boolean, isOverride: Boolean = true): SerializedVariablesState {
458+
private fun serializeVariableState(cellId: Int, topLevelName: String, property: Field?, value: Any?, isRecursive: Boolean, isOverride: Boolean = true): SerializedVariablesState {
467459
val wrapper = value.toObjectWrapper(isRecursive)
468-
val processedData = createSerializeVariableState(name, getSimpleTypeNameFrom(property, value), wrapper)
469-
return doActualSerialization(cellId, processedData, wrapper, isRecursive, isOverride)
460+
val processedData = createSerializeVariableState(topLevelName, getSimpleTypeNameFrom(property, value), wrapper)
461+
return doActualSerialization(cellId, topLevelName, processedData, wrapper, isRecursive, isOverride)
470462
}
471463

472-
private fun serializeVariableState(cellId: Int, name: String, property: KProperty<*>, value: Any?, isRecursive: Boolean, isOverride: Boolean = true): SerializedVariablesState {
464+
private fun serializeVariableState(cellId: Int, topLevelName: String, property: KProperty<*>, value: Any?, isRecursive: Boolean, isOverride: Boolean = true): SerializedVariablesState {
473465
val wrapper = value.toObjectWrapper(isRecursive)
474-
val processedData = createSerializeVariableState(name, getSimpleTypeNameFrom(property, value), wrapper)
475-
return doActualSerialization(cellId, processedData, wrapper, isRecursive, isOverride)
466+
val processedData = createSerializeVariableState(topLevelName, getSimpleTypeNameFrom(property, value), wrapper)
467+
return doActualSerialization(cellId, topLevelName, processedData, wrapper, isRecursive, isOverride)
476468
}
477469

478-
private fun doActualSerialization(cellId: Int, processedData: ProcessedSerializedVarsState, value: RuntimeObjectWrapper, isRecursive: Boolean, isOverride: Boolean = true): SerializedVariablesState {
470+
private fun doActualSerialization(cellId: Int, topLevelName:String, processedData: ProcessedSerializedVarsState, value: RuntimeObjectWrapper, isRecursive: Boolean, isOverride: Boolean = true): SerializedVariablesState {
479471
val serializedVersion = processedData.serializedVariablesState
480472

481473
seenObjectsPerCell.putIfAbsent(cellId, mutableMapOf())
474+
computedDescriptorsPerCell.putIfAbsent(cellId, mutableMapOf())
482475

483476
if (isOverride) {
484-
val instances = computedDescriptorsPerCell[cellId]?.instancesPerState
485-
computedDescriptorsPerCell[cellId] = ProcessedDescriptorsState()
477+
val instances = computedDescriptorsPerCell[cellId]?.get(topLevelName)?.instancesPerState
478+
computedDescriptorsPerCell[cellId]!![topLevelName] = ProcessedDescriptorsState()
486479
if (instances != null) {
487-
computedDescriptorsPerCell[cellId]!!.instancesPerState += instances
480+
computedDescriptorsPerCell[cellId]!![topLevelName]!!.instancesPerState += instances
488481
}
489482
}
490-
val currentCellDescriptors = computedDescriptorsPerCell[cellId]
483+
val currentCellDescriptors = computedDescriptorsPerCell[cellId]?.get(topLevelName)
491484
// TODO should we stack?
492485
currentCellDescriptors!!.processedSerializedVarsToJavaProperties[serializedVersion] = processedData.propertiesData
493486
currentCellDescriptors.processedSerializedVarsToKTProperties[serializedVersion] = processedData.kPropertiesData
@@ -507,9 +500,9 @@ class VariablesSerializer(
507500
if (kProperties?.size == 1 && kProperties.first().name == "size") {
508501
serializedVersion.fieldDescriptor.addDescriptor(value.objectInstance, "data")
509502
}
510-
iterateThroughContainerMembers(cellId, value.objectInstance, serializedVersion.fieldDescriptor, isRecursive = isRecursive, kProperties = currentCellDescriptors.processedSerializedVarsToKTProperties[serializedVersion])
503+
iterateThroughContainerMembers(cellId, topLevelName, value.objectInstance, serializedVersion.fieldDescriptor, isRecursive = isRecursive, kProperties = currentCellDescriptors.processedSerializedVarsToKTProperties[serializedVersion])
511504
} else {
512-
iterateThroughContainerMembers(cellId, value.objectInstance, serializedVersion.fieldDescriptor, isRecursive = isRecursive, currentCellDescriptors.processedSerializedVarsToJavaProperties[serializedVersion])
505+
iterateThroughContainerMembers(cellId, topLevelName, value.objectInstance, serializedVersion.fieldDescriptor, isRecursive = isRecursive, currentCellDescriptors.processedSerializedVarsToJavaProperties[serializedVersion])
513506
}
514507
}
515508

@@ -518,6 +511,7 @@ class VariablesSerializer(
518511

519512
private fun iterateThroughContainerMembers(
520513
cellId: Int,
514+
topLevelName: String,
521515
callInstance: Any?,
522516
descriptor: MutableFieldDescriptor,
523517
isRecursive: Boolean = false,
@@ -543,7 +537,7 @@ class VariablesSerializer(
543537

544538
seenObjectsPerCell.putIfAbsent(cellId, mutableMapOf())
545539
val seenObjectsPerCell = seenObjectsPerCell[cellId]
546-
val currentCellDescriptors = computedDescriptorsPerCell[cellId]!!
540+
val currentCellDescriptors = computedDescriptorsPerCell[cellId]!![topLevelName]!!
547541
// ok, it's a copy on the left for some reason
548542
val instancesPerState = currentCellDescriptors.instancesPerState
549543

@@ -570,7 +564,7 @@ class VariablesSerializer(
570564
}
571565

572566
val isArrayType = checkForPossibleArray(callInstance)
573-
computedDescriptorsPerCell[cellId]!!.instancesPerState += instancesPerState
567+
computedDescriptorsPerCell[cellId]!![topLevelName]!!.instancesPerState += instancesPerState
574568

575569
if (descriptor.size == 2 && (descriptor.containsKey("data") || descriptor.containsKey("element"))) {
576570
val singleElemMode = descriptor.containsKey("element")
@@ -606,9 +600,10 @@ class VariablesSerializer(
606600
}
607601
}.toObjectWrapper(isRecursive)
608602

609-
computedDescriptorsPerCell[cellId]!!.instancesPerState += instancesPerState
603+
computedDescriptorsPerCell[cellId]!![topLevelName]!!.instancesPerState += instancesPerState
610604
iterateThroughContainerMembers(
611605
cellId,
606+
topLevelName,
612607
neededCallInstance.objectInstance,
613608
serializedVariablesState.fieldDescriptor,
614609
isRecursive = isRecursive,

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

Lines changed: 30 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -794,7 +794,7 @@ class ReplVarsTest : AbstractSingleReplTest() {
794794
val serializer = repl.variablesSerializer
795795
val descriptor = res.evaluatedVariablesState["l"]!!.fieldDescriptor
796796
val innerList = descriptor["elementData"]!!.fieldDescriptor["data"]
797-
val newData = serializer.doIncrementalSerialization(0, "data", innerList!!)
797+
val newData = serializer.doIncrementalSerialization(0, "l", "data", innerList!!)
798798
assertEquals(2, newData.fieldDescriptor.size)
799799
}
800800

@@ -835,7 +835,7 @@ class ReplVarsTest : AbstractSingleReplTest() {
835835
jupyterId = 2
836836
).metadata.evaluatedVariablesState
837837
val innerList = res["l"]!!.fieldDescriptor["elementData"]!!.fieldDescriptor["data"]
838-
val newData = serializer.doIncrementalSerialization(0, "data", innerList!!)
838+
val newData = serializer.doIncrementalSerialization(0, "l","data", innerList!!)
839839
assertTrue(newData.isContainer)
840840
assertTrue(newData.fieldDescriptor.size > 4)
841841
}
@@ -951,7 +951,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
951951
assertEquals(listOf(1, 2, 3, 4).toString().substring(1, actualContainer.value!!.length + 1), actualContainer.value)
952952

953953
val serializer = repl.variablesSerializer
954-
val newData = serializer.doIncrementalSerialization(0, "data", actualContainer)
954+
val newData = serializer.doIncrementalSerialization(0, "x","data", actualContainer)
955955
}
956956

957957
@Test
@@ -1040,7 +1040,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
10401040

10411041
val serializer = repl.variablesSerializer
10421042

1043-
val newData = serializer.doIncrementalSerialization(0, "i", descriptor["i"]!!)
1043+
val newData = serializer.doIncrementalSerialization(0, "c", "i", descriptor["i"]!!)
10441044
}
10451045

10461046
@Test
@@ -1060,7 +1060,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
10601060
val actualContainer = listData.fieldDescriptor.entries.first().value!!
10611061
val serializer = repl.variablesSerializer
10621062

1063-
val newData = serializer.doIncrementalSerialization(0, listData.fieldDescriptor.entries.first().key, actualContainer)
1063+
val newData = serializer.doIncrementalSerialization(0, "x", listData.fieldDescriptor.entries.first().key, actualContainer)
10641064
val receivedDescriptor = newData.fieldDescriptor
10651065
assertEquals(4, receivedDescriptor.size)
10661066

@@ -1073,7 +1073,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
10731073
}
10741074

10751075
val depthMostNode = actualContainer.fieldDescriptor.entries.first { it.value!!.isContainer }
1076-
val serializationAns = serializer.doIncrementalSerialization(0, depthMostNode.key, depthMostNode.value!!)
1076+
val serializationAns = serializer.doIncrementalSerialization(0, "x", depthMostNode.key, depthMostNode.value!!)
10771077
}
10781078

10791079
@Test
@@ -1091,7 +1091,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
10911091
val serializer = repl.variablesSerializer
10921092
val path = listOf("x", "a")
10931093

1094-
val newData = serializer.doIncrementalSerialization(0, listData.fieldDescriptor.entries.first().key, actualContainer, path)
1094+
val newData = serializer.doIncrementalSerialization(0, "x", listData.fieldDescriptor.entries.first().key, actualContainer, path)
10951095
val receivedDescriptor = newData.fieldDescriptor
10961096
assertEquals(4, receivedDescriptor.size)
10971097

@@ -1132,7 +1132,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
11321132

11331133
val serializer = repl.variablesSerializer
11341134

1135-
var newData = serializer.doIncrementalSerialization(0, "values", valuesDescriptor)
1135+
var newData = serializer.doIncrementalSerialization(0, "x", "values", valuesDescriptor)
11361136
var newDescriptor = newData.fieldDescriptor
11371137
assertEquals("4", newDescriptor["size"]!!.value)
11381138
assertEquals(3, newDescriptor["data"]!!.fieldDescriptor.size)
@@ -1147,7 +1147,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
11471147
val entriesDescriptor = listDescriptors["entries"]!!
11481148
assertEquals("4", valuesDescriptor.fieldDescriptor["size"]!!.value)
11491149
assertTrue(valuesDescriptor.fieldDescriptor["data"]!!.isContainer)
1150-
newData = serializer.doIncrementalSerialization(0, "entries", entriesDescriptor)
1150+
newData = serializer.doIncrementalSerialization(0, "x", "entries", entriesDescriptor)
11511151
newDescriptor = newData.fieldDescriptor
11521152
assertEquals("4", newDescriptor["size"]!!.value)
11531153
assertEquals(4, newDescriptor["data"]!!.fieldDescriptor.size)
@@ -1227,7 +1227,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
12271227
val propertyName = listData.fieldDescriptor.entries.first().key
12281228

12291229
runBlocking {
1230-
repl.serializeVariables(1, mapOf(propertyName to actualContainer)) { result ->
1230+
repl.serializeVariables(1, "x", mapOf(propertyName to actualContainer)) { result ->
12311231
val data = result.descriptorsState
12321232
assertTrue(data.isNotEmpty())
12331233

@@ -1288,7 +1288,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
12881288
val propertyName = listData.fieldDescriptor.entries.first().key
12891289

12901290
runBlocking {
1291-
repl.serializeVariables(1, mapOf(propertyName to actualContainer)) { result ->
1291+
repl.serializeVariables(1, "c", mapOf(propertyName to actualContainer)) { result ->
12921292
val data = result.descriptorsState
12931293
assertTrue(data.isNotEmpty())
12941294

@@ -1303,7 +1303,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
13031303

13041304
val anotherI = originalClass.fieldDescriptor["i"]!!
13051305
runBlocking {
1306-
repl.serializeVariables(1, mapOf(propertyName to anotherI)) { res ->
1306+
repl.serializeVariables(1, "c", mapOf(propertyName to anotherI)) { res ->
13071307
val data = res.descriptorsState
13081308
val innerList = data.entries.last().value
13091309
assertTrue(innerList.isContainer)
@@ -1395,4 +1395,22 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
13951395
state = repl.notebook.unchangedVariables()
13961396
assertTrue(state.isEmpty())
13971397
}
1398+
1399+
@Test
1400+
fun testSerializationClearInfo() {
1401+
var res = eval(
1402+
"""
1403+
val x = listOf(1, 2, 3, 4)
1404+
""".trimIndent(),
1405+
jupyterId = 1
1406+
).metadata.evaluatedVariablesState
1407+
var state = repl.notebook.unchangedVariables()
1408+
res = eval(
1409+
"""
1410+
val x = listOf(1, 2, 3, 4)
1411+
""".trimIndent(),
1412+
jupyterId = 2
1413+
).metadata.evaluatedVariablesState
1414+
val a = 1
1415+
}
13981416
}

0 commit comments

Comments
 (0)