Skip to content

Commit 7ac866a

Browse files
Store all meta-data related to a top-level variable name
1 parent b4ac2ce commit 7ac866a

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
@@ -347,9 +347,7 @@ fun JupyterConnection.Socket.shellMessagesHandler(msg: Message, repl: ReplForJup
347347
sendWrapped(msg, makeReplyMessage(msg, MessageType.SERIALIZATION_REPLY, content = result))
348348
}
349349
} else {
350-
repl.serializeVariables(content.cellId, content.descriptorsState) { result ->
351-
sendWrapped(msg, makeReplyMessage(msg, MessageType.SERIALIZATION_REPLY, content = result))
352-
}
350+
sendWrapped(msg, makeReplyMessage(msg, MessageType.SERIALIZATION_REPLY, content = null))
353351
}
354352
}
355353
}

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

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

141141
suspend fun listErrors(code: Code, callback: (ListErrorsResult) -> Unit)
142142

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

145145
suspend fun serializeVariables(topLevelVarName: String, descriptorsState: Map<String, SerializedVariablesState>, pathToDescriptor: List<String> = emptyList(),
146146
callback: (SerializationReply) -> Unit)
@@ -570,8 +570,8 @@ class ReplForJupyterImpl(
570570
}
571571

572572
private val serializationQueue = LockQueue<SerializationReply, SerializationArgs>()
573-
override suspend fun serializeVariables(cellId: Int, descriptorsState: Map<String, SerializedVariablesState>, callback: (SerializationReply) -> Unit) {
574-
doWithLock(SerializationArgs(descriptorsState, cellId = cellId, callback = callback), serializationQueue, SerializationReply(cellId, descriptorsState), ::doSerializeVariables)
573+
override suspend fun serializeVariables(cellId: Int, topLevelVarName: String, descriptorsState: Map<String, SerializedVariablesState>, callback: (SerializationReply) -> Unit) {
574+
doWithLock(SerializationArgs(descriptorsState, cellId = cellId, topLevelVarName = topLevelVarName, callback = callback), serializationQueue, SerializationReply(cellId, descriptorsState), ::doSerializeVariables)
575575
}
576576

577577
override suspend fun serializeVariables(topLevelVarName: String, descriptorsState: Map<String, SerializedVariablesState>, pathToDescriptor: List<String>,
@@ -587,7 +587,7 @@ class ReplForJupyterImpl(
587587
finalAns
588588
}
589589
args.descriptorsState.forEach { (name, state) ->
590-
resultMap[name] = variablesSerializer.doIncrementalSerialization(cellId - 1, name, state, args.pathToDescriptor)
590+
resultMap[name] = variablesSerializer.doIncrementalSerialization(cellId - 1, args.topLevelVarName ,name, state, args.pathToDescriptor)
591591
}
592592
log.debug("Serialization cellID: $cellId")
593593
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
@@ -770,7 +770,7 @@ class ReplVarsTest : AbstractSingleReplTest() {
770770
val serializer = repl.variablesSerializer
771771
val descriptor = res.evaluatedVariablesState["l"]!!.fieldDescriptor
772772
val innerList = descriptor["elementData"]!!.fieldDescriptor["data"]
773-
val newData = serializer.doIncrementalSerialization(0, "data", innerList!!)
773+
val newData = serializer.doIncrementalSerialization(0, "l", "data", innerList!!)
774774
assertEquals(2, newData.fieldDescriptor.size)
775775
}
776776

@@ -811,7 +811,7 @@ class ReplVarsTest : AbstractSingleReplTest() {
811811
jupyterId = 2
812812
).metadata.evaluatedVariablesState
813813
val innerList = res["l"]!!.fieldDescriptor["elementData"]!!.fieldDescriptor["data"]
814-
val newData = serializer.doIncrementalSerialization(0, "data", innerList!!)
814+
val newData = serializer.doIncrementalSerialization(0, "l","data", innerList!!)
815815
assertTrue(newData.isContainer)
816816
assertTrue(newData.fieldDescriptor.size > 4)
817817
}
@@ -927,7 +927,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
927927
assertEquals(listOf(1, 2, 3, 4).toString().substring(1, actualContainer.value!!.length + 1), actualContainer.value)
928928

929929
val serializer = repl.variablesSerializer
930-
val newData = serializer.doIncrementalSerialization(0, "data", actualContainer)
930+
val newData = serializer.doIncrementalSerialization(0, "x","data", actualContainer)
931931
}
932932

933933
@Test
@@ -1016,7 +1016,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
10161016

10171017
val serializer = repl.variablesSerializer
10181018

1019-
val newData = serializer.doIncrementalSerialization(0, "i", descriptor["i"]!!)
1019+
val newData = serializer.doIncrementalSerialization(0, "c", "i", descriptor["i"]!!)
10201020
}
10211021

10221022
@Test
@@ -1036,7 +1036,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
10361036
val actualContainer = listData.fieldDescriptor.entries.first().value!!
10371037
val serializer = repl.variablesSerializer
10381038

1039-
val newData = serializer.doIncrementalSerialization(0, listData.fieldDescriptor.entries.first().key, actualContainer)
1039+
val newData = serializer.doIncrementalSerialization(0, "x", listData.fieldDescriptor.entries.first().key, actualContainer)
10401040
val receivedDescriptor = newData.fieldDescriptor
10411041
assertEquals(4, receivedDescriptor.size)
10421042

@@ -1049,7 +1049,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
10491049
}
10501050

10511051
val depthMostNode = actualContainer.fieldDescriptor.entries.first { it.value!!.isContainer }
1052-
val serializationAns = serializer.doIncrementalSerialization(0, depthMostNode.key, depthMostNode.value!!)
1052+
val serializationAns = serializer.doIncrementalSerialization(0, "x", depthMostNode.key, depthMostNode.value!!)
10531053
}
10541054

10551055
@Test
@@ -1067,7 +1067,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
10671067
val serializer = repl.variablesSerializer
10681068
val path = listOf("x", "a")
10691069

1070-
val newData = serializer.doIncrementalSerialization(0, listData.fieldDescriptor.entries.first().key, actualContainer, path)
1070+
val newData = serializer.doIncrementalSerialization(0, "x", listData.fieldDescriptor.entries.first().key, actualContainer, path)
10711071
val receivedDescriptor = newData.fieldDescriptor
10721072
assertEquals(4, receivedDescriptor.size)
10731073

@@ -1108,7 +1108,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
11081108

11091109
val serializer = repl.variablesSerializer
11101110

1111-
var newData = serializer.doIncrementalSerialization(0, "values", valuesDescriptor)
1111+
var newData = serializer.doIncrementalSerialization(0, "x", "values", valuesDescriptor)
11121112
var newDescriptor = newData.fieldDescriptor
11131113
assertEquals("4", newDescriptor["size"]!!.value)
11141114
assertEquals(3, newDescriptor["data"]!!.fieldDescriptor.size)
@@ -1123,7 +1123,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
11231123
val entriesDescriptor = listDescriptors["entries"]!!
11241124
assertEquals("4", valuesDescriptor.fieldDescriptor["size"]!!.value)
11251125
assertTrue(valuesDescriptor.fieldDescriptor["data"]!!.isContainer)
1126-
newData = serializer.doIncrementalSerialization(0, "entries", entriesDescriptor)
1126+
newData = serializer.doIncrementalSerialization(0, "x", "entries", entriesDescriptor)
11271127
newDescriptor = newData.fieldDescriptor
11281128
assertEquals("4", newDescriptor["size"]!!.value)
11291129
assertEquals(4, newDescriptor["data"]!!.fieldDescriptor.size)
@@ -1203,7 +1203,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
12031203
val propertyName = listData.fieldDescriptor.entries.first().key
12041204

12051205
runBlocking {
1206-
repl.serializeVariables(1, mapOf(propertyName to actualContainer)) { result ->
1206+
repl.serializeVariables(1, "x", mapOf(propertyName to actualContainer)) { result ->
12071207
val data = result.descriptorsState
12081208
assertTrue(data.isNotEmpty())
12091209

@@ -1264,7 +1264,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
12641264
val propertyName = listData.fieldDescriptor.entries.first().key
12651265

12661266
runBlocking {
1267-
repl.serializeVariables(1, mapOf(propertyName to actualContainer)) { result ->
1267+
repl.serializeVariables(1, "c", mapOf(propertyName to actualContainer)) { result ->
12681268
val data = result.descriptorsState
12691269
assertTrue(data.isNotEmpty())
12701270

@@ -1279,7 +1279,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
12791279

12801280
val anotherI = originalClass.fieldDescriptor["i"]!!
12811281
runBlocking {
1282-
repl.serializeVariables(1, mapOf(propertyName to anotherI)) { res ->
1282+
repl.serializeVariables(1, "c", mapOf(propertyName to anotherI)) { res ->
12831283
val data = res.descriptorsState
12841284
val innerList = data.entries.last().value
12851285
assertTrue(innerList.isContainer)
@@ -1371,4 +1371,22 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
13711371
state = repl.notebook.unchangedVariables()
13721372
assertTrue(state.isEmpty())
13731373
}
1374+
1375+
@Test
1376+
fun testSerializationClearInfo() {
1377+
var res = eval(
1378+
"""
1379+
val x = listOf(1, 2, 3, 4)
1380+
""".trimIndent(),
1381+
jupyterId = 1
1382+
).metadata.evaluatedVariablesState
1383+
var state = repl.notebook.unchangedVariables()
1384+
res = eval(
1385+
"""
1386+
val x = listOf(1, 2, 3, 4)
1387+
""".trimIndent(),
1388+
jupyterId = 2
1389+
).metadata.evaluatedVariablesState
1390+
val a = 1
1391+
}
13741392
}

0 commit comments

Comments
 (0)