Skip to content

Commit f20216b

Browse files
nikolay-egorovnikolay-egorov
committed
Add incremental serialization support prototype; support cyclic references
1 parent d4f46fc commit f20216b

File tree

3 files changed

+276
-4
lines changed

3 files changed

+276
-4
lines changed

‎jupyter-lib/shared-compiler/src/main/kotlin/org/jetbrains/kotlinx/jupyter/compiler/util/serializedCompiledScript.kt‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,12 +20,23 @@ data class SerializedCompiledScriptsData(
2020
}
2121
}
2222

23+
@Serializable
24+
data class SerializedVariablesState(
25+
val name: String = "",
26+
val type: String = "",
27+
val value: String? = null,
28+
val isContainer: Boolean = false
29+
) {
30+
val fieldDescriptor: MutableMap<String, SerializedVariablesState?> = mutableMapOf()
31+
}
32+
33+
2334
@Serializable
2435
class EvaluatedSnippetMetadata(
2536
val newClasspath: Classpath = emptyList(),
2637
val compiledData: SerializedCompiledScriptsData = SerializedCompiledScriptsData.EMPTY,
2738
val newImports: List<String> = emptyList(),
28-
val evaluatedVariablesState: Map<String, String?> = mutableMapOf()
39+
val evaluatedVariablesState: SerializedVariablesState = SerializedVariablesState()
2940
) {
3041
companion object {
3142
val EMPTY = EvaluatedSnippetMetadata()

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -420,10 +420,10 @@ class ReplForJupyterImpl(
420420
notebook.updateVariablesState(internalEvaluator)
421421
// printVars()
422422
// printUsagesInfo(jupyterId, cellVariables[jupyterId - 1])
423+
val entry = notebook.variablesState.entries.last()
424+
val serializedVarsState = serializeVariableState(entry.key, entry.value)
423425

424-
425-
val variablesStateUpdate = notebook.variablesState.mapValues { it.value.stringValue }
426-
EvalResult(rendered, EvaluatedSnippetMetadata(newClasspath, compiledData, newImports, variablesStateUpdate))
426+
EvalResult(rendered, EvaluatedSnippetMetadata(newClasspath, compiledData, newImports, serializedVarsState))
427427
}
428428
}
429429

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

Lines changed: 261 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,261 @@
1+
package org.jetbrains.kotlinx.jupyter
2+
3+
import org.jetbrains.kotlinx.jupyter.api.VariableState
4+
import org.jetbrains.kotlinx.jupyter.compiler.util.SerializedVariablesState
5+
import kotlin.reflect.KClass
6+
import kotlin.reflect.KProperty
7+
import kotlin.reflect.KProperty1
8+
import kotlin.reflect.full.declaredMemberProperties
9+
import kotlin.reflect.full.isSubclassOf
10+
import kotlin.reflect.jvm.isAccessible
11+
import kotlin.reflect.jvm.jvmErasure
12+
13+
typealias FieldDescriptor = Map<String, SerializedVariablesState?>
14+
typealias MutableFieldDescriptor = MutableMap<String, SerializedVariablesState?>
15+
typealias PropertiesData = Collection<KProperty1<out Any, *>>
16+
17+
data class ProcessedSerializedVarsState(
18+
val serializedVariablesState: SerializedVariablesState,
19+
val propertiesData: PropertiesData?
20+
)
21+
22+
data class ProcessedDescriptorsState(
23+
// perhaps, better tp make SerializedVariablesState -> PropertiesData?
24+
val processedSerializedVarsState: MutableMap<SerializedVariablesState, PropertiesData?> = mutableMapOf(),
25+
val instancesPerState: MutableMap<SerializedVariablesState, Any?> = mutableMapOf()
26+
)
27+
28+
class VariablesSerializer(private val serializationStep: Int = 2, private val serializationLimit: Int = 10000) {
29+
30+
private val seenObjectsPerCell: MutableMap<Int, MutableMap<Any, SerializedVariablesState>> = mutableMapOf()
31+
32+
var currentSerializeCount: Int = 0
33+
34+
/**
35+
* Stores info computed descriptors in a cell
36+
*/
37+
private val computedDescriptorsPerCell: MutableMap<Int, ProcessedDescriptorsState> = mutableMapOf()
38+
39+
40+
fun serializeVariables(cellId: Int, variablesState: Map<String, VariableState>): Map<String, SerializedVariablesState> {
41+
return variablesState.mapValues { serializeVariableState(cellId, it.key, it.value) }
42+
}
43+
44+
fun doIncrementalSerialization(cellId: Int, propertyName: String, serializedVariablesState: SerializedVariablesState): SerializedVariablesState {
45+
val cellDescriptors = computedDescriptorsPerCell[cellId] ?: return serializedVariablesState
46+
return updateVariableState(propertyName, cellDescriptors, serializedVariablesState)
47+
}
48+
49+
/**
50+
* @param evaluatedDescriptorsState - origin variable state to get value from
51+
* @param serializedVariablesState - current state of recursive state to go further
52+
*/
53+
private fun updateVariableState(propertyName: String, evaluatedDescriptorsState: ProcessedDescriptorsState,
54+
serializedVariablesState: SerializedVariablesState) : SerializedVariablesState {
55+
val value = evaluatedDescriptorsState.instancesPerState[serializedVariablesState]
56+
val propertiesData = evaluatedDescriptorsState.processedSerializedVarsState[serializedVariablesState] ?: return serializedVariablesState
57+
val property = propertiesData.firstOrNull {
58+
it.name == propertyName
59+
} ?: return serializedVariablesState
60+
61+
return serializeVariableState(propertyName, property, value)
62+
}
63+
64+
65+
fun serializeVariableState(cellId: Int, name: String, variableState: VariableState): SerializedVariablesState {
66+
return serializeVariableState(cellId, name, variableState.property, variableState.value)
67+
}
68+
69+
fun serializeVariableState(cellId: Int, name: String, property: KProperty<*>, value: Any?): SerializedVariablesState {
70+
val processedData = createSerializeVariableState(name, property, value)
71+
val serializedVersion = processedData.serializedVariablesState
72+
73+
if (seenObjectsPerCell.containsKey(cellId)) {
74+
seenObjectsPerCell[cellId]!!.clear()
75+
}
76+
seenObjectsPerCell.putIfAbsent(cellId, mutableMapOf())
77+
// always override?
78+
computedDescriptorsPerCell[cellId] = ProcessedDescriptorsState()
79+
val currentCellDescriptors = computedDescriptorsPerCell[cellId]
80+
currentCellDescriptors!!.processedSerializedVarsState[serializedVersion] = processedData.propertiesData
81+
82+
if (value != null) {
83+
seenObjectsPerCell[cellId]!![value] = serializedVersion
84+
}
85+
if (serializedVersion.isContainer) {
86+
iterateThroughContainerMembers(cellId, value, serializedVersion.fieldDescriptor, processedData.propertiesData)
87+
}
88+
return processedData.serializedVariablesState
89+
}
90+
91+
92+
private fun iterateThroughContainerMembers(cellId: Int, callInstance: Any?, descriptor: MutableFieldDescriptor, properties: PropertiesData?, currentDepth: Int = 0): Unit {
93+
if (properties == null || callInstance == null || currentDepth > serializationStep) return
94+
95+
val serializedIteration = mutableMapOf<String, ProcessedSerializedVarsState>()
96+
val callInstances = mutableMapOf<String, Any?>()
97+
98+
seenObjectsPerCell.putIfAbsent(cellId, mutableMapOf())
99+
val seenObjectsPerCell = seenObjectsPerCell[cellId]
100+
val currentCellDescriptors = computedDescriptorsPerCell[cellId]!!
101+
val instancesPerState = currentCellDescriptors.instancesPerState
102+
103+
for (it in properties) {
104+
if (currentSerializeCount > serializationLimit) {
105+
break
106+
}
107+
it as KProperty1<Any, *>
108+
val name = it.name
109+
val wasAccessible = it.isAccessible
110+
it.isAccessible = true
111+
val value = it.get(callInstance)
112+
113+
if (!seenObjectsPerCell!!.containsKey(value)) {
114+
serializedIteration[name] = createSerializeVariableState(name, it, value)
115+
descriptor[name] = serializedIteration[name]!!.serializedVariablesState
116+
}
117+
instancesPerState[descriptor[name]!!] = value
118+
119+
if (value != null && !seenObjectsPerCell.containsKey(value)) {
120+
if (descriptor[name] != null) {
121+
seenObjectsPerCell[value] = descriptor[name]!!
122+
}
123+
}
124+
it.isAccessible = wasAccessible
125+
currentSerializeCount++
126+
}
127+
128+
serializedIteration.forEach {
129+
val serializedVariablesState = it.value.serializedVariablesState
130+
val name = it.key
131+
if (serializedVariablesState.isContainer) {
132+
iterateThroughContainerMembers(cellId, callInstances[name], serializedVariablesState.fieldDescriptor,
133+
it.value.propertiesData, currentDepth + 1)
134+
}
135+
}
136+
}
137+
138+
139+
}
140+
// TODO: place code bellow to the VariablesSerializer once it's good
141+
/**
142+
* Map of seen objects.
143+
* Key: hash code of actual value
144+
* Value: this value
145+
*/
146+
val seenObjects: MutableMap<Any, SerializedVariablesState> = mutableMapOf()
147+
var currentSerializeCount: Int = 0
148+
val computedDescriptorsPerCell: Map<Int, ProcessedDescriptorsState> = mutableMapOf()
149+
150+
151+
fun serializeVariableState(name: String, property: KProperty<*>, value: Any?): SerializedVariablesState {
152+
val processedData = createSerializeVariableState(name, property, value)
153+
val serializedVersion = processedData.serializedVariablesState
154+
if (value != null) {
155+
seenObjects[value] = serializedVersion
156+
}
157+
if (serializedVersion.isContainer) {
158+
iterateThroughContainerMembers(value, serializedVersion.fieldDescriptor, processedData.propertiesData)
159+
}
160+
return processedData.serializedVariablesState
161+
}
162+
163+
fun serializeVariableState(name: String, variableState: VariableState): SerializedVariablesState {
164+
return serializeVariableState(name, variableState.property, variableState.value)
165+
}
166+
167+
// maybe let it be global
168+
fun createSerializeVariableState(name: String, property: KProperty<*>, value: Any?): ProcessedSerializedVarsState {
169+
val returnType = property.returnType
170+
val classifier = returnType.classifier as KClass<*>
171+
val membersProperties = if (value != null) value::class.declaredMemberProperties else null
172+
val isContainer = if (membersProperties != null) membersProperties.size > 1 else false
173+
val serializedVariablesState = SerializedVariablesState(name, classifier.simpleName.toString(),
174+
getProperString(value), isContainer)
175+
176+
return ProcessedSerializedVarsState(serializedVariablesState, membersProperties)
177+
}
178+
179+
fun createSerializeVariableState(name: String, variableState: VariableState): ProcessedSerializedVarsState {
180+
val returnType = variableState.property.returnType
181+
val classifier = returnType.classifier as KClass<*>
182+
val property = variableState.property
183+
val javaField = property.returnType.jvmErasure
184+
185+
val membersProperties = if (variableState.value != null) variableState.value!!::class.declaredMemberProperties else null
186+
val isContainer = if (membersProperties != null) membersProperties.size > 1 else false
187+
val serializedVariablesState = SerializedVariablesState(name, classifier.simpleName.toString(), variableState.stringValue, isContainer)
188+
189+
return ProcessedSerializedVarsState(serializedVariablesState, membersProperties)
190+
}
191+
192+
fun iterateThroughContainerMembers(callInstance: Any?, descriptor: MutableFieldDescriptor, properties: PropertiesData?, currentDepth: Int = 0): Unit {
193+
if (properties == null || callInstance == null || currentDepth > 2) return
194+
195+
val serializedIteration = mutableMapOf<String, ProcessedSerializedVarsState>()
196+
val callInstances = mutableMapOf<String, Any?>()
197+
198+
for (it in properties) {
199+
if (currentSerializeCount > 1000) {
200+
break
201+
}
202+
it as KProperty1<Any, *>
203+
val name = it.name
204+
val wasAccessible = it.isAccessible
205+
it.isAccessible = true
206+
val value = it.get(callInstance)
207+
208+
if (!seenObjects.containsKey(value)) {
209+
serializedIteration[name] = createSerializeVariableState(name, it, value)
210+
descriptor[name] = serializedIteration[name]?.serializedVariablesState
211+
}
212+
213+
if (value != null && !seenObjects.containsKey(value)) {
214+
if (descriptor[name] != null) {
215+
seenObjects[value] = descriptor[name]!!
216+
}
217+
}
218+
it.isAccessible = wasAccessible
219+
currentSerializeCount++
220+
}
221+
222+
serializedIteration.forEach {
223+
val serializedVariablesState = it.value.serializedVariablesState
224+
val name = it.key
225+
if (serializedVariablesState.isContainer) {
226+
iterateThroughContainerMembers(callInstances[name], serializedVariablesState.fieldDescriptor,
227+
it.value.propertiesData, currentDepth + 1)
228+
}
229+
}
230+
}
231+
232+
233+
fun getProperString(value: Any?) : String {
234+
value ?: return "null"
235+
236+
val kClass = value::class
237+
val isFromJavaArray = kClass.java.isArray
238+
if (isFromJavaArray || kClass.isSubclassOf(Array::class)) {
239+
value as Array<*>
240+
return value.toString()
241+
}
242+
val isCollection = kClass.isSubclassOf(Collection::class)
243+
if (isCollection) {
244+
value as Collection<*>
245+
return buildString {
246+
value.forEach {
247+
append(it.toString(), ", ")
248+
}
249+
}
250+
}
251+
val isMap = kClass.isSubclassOf(Map::class)
252+
if (isMap) {
253+
value as Map<Any, Any?>
254+
return buildString {
255+
value.forEach {
256+
append(it.key, '=', it.value, "\n")
257+
}
258+
}
259+
}
260+
return value.toString()
261+
}

0 commit comments

Comments
 (0)