Skip to content

Commit db437cc

Browse files
nikolay-egorovnikolay-egorov
committed
Store all meta-data related to a top-level variable name
1 parent 6c85ff4 commit db437cc

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

+1-3
Original file line numberDiff line numberDiff line change
@@ -345,9 +345,7 @@ fun JupyterConnection.Socket.shellMessagesHandler(msg: Message, repl: ReplForJup
345345
sendWrapped(msg, makeReplyMessage(msg, MessageType.SERIALIZATION_REPLY, content = result))
346346
}
347347
} else {
348-
repl.serializeVariables(content.cellId, content.descriptorsState) { result ->
349-
sendWrapped(msg, makeReplyMessage(msg, MessageType.SERIALIZATION_REPLY, content = result))
350-
}
348+
sendWrapped(msg, makeReplyMessage(msg, MessageType.SERIALIZATION_REPLY, content = null))
351349
}
352350
}
353351
}

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

+4-4
Original file line numberDiff line numberDiff line change
@@ -123,7 +123,7 @@ interface ReplForJupyter {
123123

124124
suspend fun listErrors(code: Code, callback: (ListErrorsResult) -> Unit)
125125

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

128128
suspend fun serializeVariables(topLevelVarName: String, descriptorsState: Map<String, SerializedVariablesState>, pathToDescriptor: List<String> = emptyList(),
129129
callback: (SerializationReply) -> Unit)
@@ -535,8 +535,8 @@ class ReplForJupyterImpl(
535535
}
536536

537537
private val serializationQueue = LockQueue<SerializationReply, SerializationArgs>()
538-
override suspend fun serializeVariables(cellId: Int, descriptorsState: Map<String, SerializedVariablesState>, callback: (SerializationReply) -> Unit) {
539-
doWithLock(SerializationArgs(descriptorsState, cellId = cellId, callback = callback), serializationQueue, SerializationReply(cellId, descriptorsState), ::doSerializeVariables)
538+
override suspend fun serializeVariables(cellId: Int, topLevelVarName: String, descriptorsState: Map<String, SerializedVariablesState>, callback: (SerializationReply) -> Unit) {
539+
doWithLock(SerializationArgs(descriptorsState, cellId = cellId, topLevelVarName = topLevelVarName, callback = callback), serializationQueue, SerializationReply(cellId, descriptorsState), ::doSerializeVariables)
540540
}
541541

542542
override suspend fun serializeVariables(topLevelVarName: String, descriptorsState: Map<String, SerializedVariablesState>, pathToDescriptor: List<String>,
@@ -552,7 +552,7 @@ class ReplForJupyterImpl(
552552
finalAns
553553
}
554554
args.descriptorsState.forEach { (name, state) ->
555-
resultMap[name] = variablesSerializer.doIncrementalSerialization(cellId - 1, name, state, args.pathToDescriptor)
555+
resultMap[name] = variablesSerializer.doIncrementalSerialization(cellId - 1, args.topLevelVarName ,name, state, args.pathToDescriptor)
556556
}
557557
log.debug("Serialization cellID: $cellId")
558558
log.debug("Serialization answer: ${resultMap.entries.first().value.fieldDescriptor}")

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

+29-34
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

+30-12
Original file line numberDiff line numberDiff line change
@@ -767,7 +767,7 @@ class ReplVarsTest : AbstractSingleReplTest() {
767767
val serializer = repl.variablesSerializer
768768
val descriptor = res.evaluatedVariablesState["l"]!!.fieldDescriptor
769769
val innerList = descriptor["elementData"]!!.fieldDescriptor["data"]
770-
val newData = serializer.doIncrementalSerialization(0, "data", innerList!!)
770+
val newData = serializer.doIncrementalSerialization(0, "l", "data", innerList!!)
771771
assertEquals(2, newData.fieldDescriptor.size)
772772
}
773773

@@ -808,7 +808,7 @@ class ReplVarsTest : AbstractSingleReplTest() {
808808
jupyterId = 2
809809
).metadata.evaluatedVariablesState
810810
val innerList = res["l"]!!.fieldDescriptor["elementData"]!!.fieldDescriptor["data"]
811-
val newData = serializer.doIncrementalSerialization(0, "data", innerList!!)
811+
val newData = serializer.doIncrementalSerialization(0, "l","data", innerList!!)
812812
assertTrue(newData.isContainer)
813813
assertTrue(newData.fieldDescriptor.size > 4)
814814
}
@@ -924,7 +924,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
924924
assertEquals(listOf(1, 2, 3, 4).toString().substring(1, actualContainer.value!!.length + 1), actualContainer.value)
925925

926926
val serializer = repl.variablesSerializer
927-
val newData = serializer.doIncrementalSerialization(0, "data", actualContainer)
927+
val newData = serializer.doIncrementalSerialization(0, "x","data", actualContainer)
928928
}
929929

930930
@Test
@@ -1013,7 +1013,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
10131013

10141014
val serializer = repl.variablesSerializer
10151015

1016-
val newData = serializer.doIncrementalSerialization(0, "i", descriptor["i"]!!)
1016+
val newData = serializer.doIncrementalSerialization(0, "c", "i", descriptor["i"]!!)
10171017
}
10181018

10191019
@Test
@@ -1033,7 +1033,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
10331033
val actualContainer = listData.fieldDescriptor.entries.first().value!!
10341034
val serializer = repl.variablesSerializer
10351035

1036-
val newData = serializer.doIncrementalSerialization(0, listData.fieldDescriptor.entries.first().key, actualContainer)
1036+
val newData = serializer.doIncrementalSerialization(0, "x", listData.fieldDescriptor.entries.first().key, actualContainer)
10371037
val receivedDescriptor = newData.fieldDescriptor
10381038
assertEquals(4, receivedDescriptor.size)
10391039

@@ -1046,7 +1046,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
10461046
}
10471047

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

10521052
@Test
@@ -1064,7 +1064,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
10641064
val serializer = repl.variablesSerializer
10651065
val path = listOf("x", "a")
10661066

1067-
val newData = serializer.doIncrementalSerialization(0, listData.fieldDescriptor.entries.first().key, actualContainer, path)
1067+
val newData = serializer.doIncrementalSerialization(0, "x", listData.fieldDescriptor.entries.first().key, actualContainer, path)
10681068
val receivedDescriptor = newData.fieldDescriptor
10691069
assertEquals(4, receivedDescriptor.size)
10701070

@@ -1105,7 +1105,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
11051105

11061106
val serializer = repl.variablesSerializer
11071107

1108-
var newData = serializer.doIncrementalSerialization(0, "values", valuesDescriptor)
1108+
var newData = serializer.doIncrementalSerialization(0, "x", "values", valuesDescriptor)
11091109
var newDescriptor = newData.fieldDescriptor
11101110
assertEquals("4", newDescriptor["size"]!!.value)
11111111
assertEquals(3, newDescriptor["data"]!!.fieldDescriptor.size)
@@ -1120,7 +1120,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
11201120
val entriesDescriptor = listDescriptors["entries"]!!
11211121
assertEquals("4", valuesDescriptor.fieldDescriptor["size"]!!.value)
11221122
assertTrue(valuesDescriptor.fieldDescriptor["data"]!!.isContainer)
1123-
newData = serializer.doIncrementalSerialization(0, "entries", entriesDescriptor)
1123+
newData = serializer.doIncrementalSerialization(0, "x", "entries", entriesDescriptor)
11241124
newDescriptor = newData.fieldDescriptor
11251125
assertEquals("4", newDescriptor["size"]!!.value)
11261126
assertEquals(4, newDescriptor["data"]!!.fieldDescriptor.size)
@@ -1200,7 +1200,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
12001200
val propertyName = listData.fieldDescriptor.entries.first().key
12011201

12021202
runBlocking {
1203-
repl.serializeVariables(1, mapOf(propertyName to actualContainer)) { result ->
1203+
repl.serializeVariables(1, "x", mapOf(propertyName to actualContainer)) { result ->
12041204
val data = result.descriptorsState
12051205
assertTrue(data.isNotEmpty())
12061206

@@ -1261,7 +1261,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
12611261
val propertyName = listData.fieldDescriptor.entries.first().key
12621262

12631263
runBlocking {
1264-
repl.serializeVariables(1, mapOf(propertyName to actualContainer)) { result ->
1264+
repl.serializeVariables(1, "c", mapOf(propertyName to actualContainer)) { result ->
12651265
val data = result.descriptorsState
12661266
assertTrue(data.isNotEmpty())
12671267

@@ -1276,7 +1276,7 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
12761276

12771277
val anotherI = originalClass.fieldDescriptor["i"]!!
12781278
runBlocking {
1279-
repl.serializeVariables(1, mapOf(propertyName to anotherI)) { res ->
1279+
repl.serializeVariables(1, "c", mapOf(propertyName to anotherI)) { res ->
12801280
val data = res.descriptorsState
12811281
val innerList = data.entries.last().value
12821282
assertTrue(innerList.isContainer)
@@ -1368,4 +1368,22 @@ class ReplVarsSerializationTest : AbstractSingleReplTest() {
13681368
state = repl.notebook.unchangedVariables()
13691369
assertTrue(state.isEmpty())
13701370
}
1371+
1372+
@Test
1373+
fun testSerializationClearInfo() {
1374+
var res = eval(
1375+
"""
1376+
val x = listOf(1, 2, 3, 4)
1377+
""".trimIndent(),
1378+
jupyterId = 1
1379+
).metadata.evaluatedVariablesState
1380+
var state = repl.notebook.unchangedVariables()
1381+
res = eval(
1382+
"""
1383+
val x = listOf(1, 2, 3, 4)
1384+
""".trimIndent(),
1385+
jupyterId = 2
1386+
).metadata.evaluatedVariablesState
1387+
val a = 1
1388+
}
13711389
}

0 commit comments

Comments
 (0)