Skip to content

Commit e34c3cc

Browse files
nikolay-egorovnikolay-egorov
committed
Add possibility to use serialization request with only top-level descriptor name, not cellID
1 parent 1a10cd2 commit e34c3cc

File tree

12 files changed

+152
-17
lines changed

12 files changed

+152
-17
lines changed

‎jupyter-lib/api/src/main/kotlin/org/jetbrains/kotlinx/jupyter/api/VariableState.kt‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,21 +6,21 @@ import kotlin.reflect.jvm.isAccessible
66
import java.lang.reflect.Field
77

88
interface VariableState {
9-
// val property: KProperty<*>
109
val property: Field
1110
val scriptInstance: Any?
1211
val stringValue: String?
1312
val value: Result<Any?>
1413
}
1514

1615
data class VariableStateImpl(
17-
// override val property: KProperty1<Any, *>,
1816
override val property: Field,
1917
override val scriptInstance: Any,
2018
) : VariableState {
2119
private var cachedValue: Result<Any?> = Result.success(null)
2220
private var isRecursive: Boolean = false
2321

22+
// use of Java 9 required
23+
@SuppressWarnings("DEPRECATION")
2424
fun update(): Boolean {
2525
val wasAccessible = property.isAccessible
2626
property.isAccessible = true

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,12 @@ data class SerializedVariablesState(
3030
val fieldDescriptor: MutableMap<String, SerializedVariablesState?> = mutableMapOf()
3131
}
3232

33+
@Serializable
34+
class SerializationReply(
35+
val cellId: Int = 1,
36+
val descriptorsState: Map<String, SerializedVariablesState> = emptyMap()
37+
)
38+
3339
@Serializable
3440
class EvaluatedSnippetMetadata(
3541
val newClasspath: Classpath = emptyList(),

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -123,6 +123,7 @@ class NotebookImpl(
123123
private var currentCellVariables = mapOf<Int, Set<String>>()
124124
private val history = arrayListOf<CodeCellImpl>()
125125
private var mainCellCreated = false
126+
private val unchangedVariables: MutableSet<String> = mutableSetOf()
126127

127128
val displays = DisplayContainerImpl()
128129

@@ -142,12 +143,16 @@ class NotebookImpl(
142143
fun updateVariablesState(evaluator: InternalEvaluator) {
143144
variablesState += evaluator.variablesHolder
144145
currentCellVariables = evaluator.cellVariables
146+
unchangedVariables.clear()
147+
unchangedVariables.addAll(evaluator.getUnchangedVariables())
145148
}
146149

147150
fun updateVariablesState(varsStateUpdate: Map<String, VariableState>) {
148151
variablesState += varsStateUpdate
149152
}
150153

154+
fun unchangedVariables(): Set<String> = unchangedVariables
155+
151156
fun variablesReportAsHTML(): String {
152157
return generateHTMLVarsReport(variablesState)
153158
}

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

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -92,9 +92,6 @@ enum class MessageType(val contentClass: KClass<out MessageContent>) {
9292
SERIALIZATION_REQUEST(SerializationRequest::class),
9393
SERIALIZATION_REPLY(SerializationReply::class);
9494

95-
// TODO: add custom commands
96-
// this custom message should be supported on client-side. either JS or Idea Plugin
97-
9895
val type: String
9996
get() = name.lowercase()
10097
}
@@ -580,12 +577,13 @@ class ListErrorsReply(
580577
@Serializable
581578
class SerializationRequest(
582579
val cellId: Int,
583-
val descriptorsState: Map<String, SerializedVariablesState>
580+
val descriptorsState: Map<String, SerializedVariablesState>,
581+
val topLevelDescriptorName: String = ""
584582
) : MessageContent()
585583

586584
@Serializable
587585
class SerializationReply(
588-
val cellId: Int,
586+
val cellId: Int = 1,
589587
val descriptorsState: Map<String, SerializedVariablesState> = emptyMap()
590588
) : MessageContent()
591589

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

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -305,6 +305,21 @@ fun JupyterConnection.Socket.shellMessagesHandler(msg: Message, repl: ReplForJup
305305
is CommInfoRequest -> {
306306
sendWrapped(msg, makeReplyMessage(msg, MessageType.COMM_INFO_REPLY, content = CommInfoReply(mapOf())))
307307
}
308+
is CommOpen -> {
309+
if (!content.commId.equals(MessageType.SERIALIZATION_REQUEST.name, ignoreCase = true)) {
310+
send(makeReplyMessage(msg, MessageType.NONE))
311+
return
312+
}
313+
log.debug("Message type in CommOpen: $msg, ${msg.type}")
314+
val data = content.data ?: return sendWrapped(msg, makeReplyMessage(msg, MessageType.SERIALIZATION_REPLY))
315+
316+
val messageContent = getVariablesDescriptorsFromJson(data)
317+
GlobalScope.launch(Dispatchers.Default) {
318+
repl.serializeVariables(messageContent.topLevelDescriptorName, messageContent.descriptorsState) { result ->
319+
sendWrapped(msg, makeReplyMessage(msg, MessageType.COMM_OPEN, content = result))
320+
}
321+
}
322+
}
308323
is CompleteRequest -> {
309324
GlobalScope.launch(Dispatchers.Default) {
310325
repl.complete(content.code, content.cursorPos) { result ->
@@ -321,8 +336,14 @@ fun JupyterConnection.Socket.shellMessagesHandler(msg: Message, repl: ReplForJup
321336
}
322337
is SerializationRequest -> {
323338
GlobalScope.launch(Dispatchers.Default) {
324-
repl.serializeVariables(content.cellId, content.descriptorsState) { result ->
325-
sendWrapped(msg, makeReplyMessage(msg, MessageType.SERIALIZATION_REPLY, content = result))
339+
if (content.topLevelDescriptorName.isNotEmpty()) {
340+
repl.serializeVariables(content.topLevelDescriptorName, content.descriptorsState) { result ->
341+
sendWrapped(msg, makeReplyMessage(msg, MessageType.SERIALIZATION_REPLY, content = result))
342+
}
343+
} else {
344+
repl.serializeVariables(content.cellId, content.descriptorsState) { result ->
345+
sendWrapped(msg, makeReplyMessage(msg, MessageType.SERIALIZATION_REPLY, content = result))
346+
}
326347
}
327348
}
328349
}

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

Lines changed: 19 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -122,6 +122,8 @@ interface ReplForJupyter {
122122

123123
suspend fun serializeVariables(cellId: Int, descriptorsState: Map<String, SerializedVariablesState>, callback: (SerializationReply) -> Unit)
124124

125+
suspend fun serializeVariables(topLevelVarName: String, descriptorsState: Map<String, SerializedVariablesState>, callback: (SerializationReply) -> Unit)
126+
125127
val homeDir: File?
126128

127129
val currentClasspath: Collection<String>
@@ -527,15 +529,26 @@ class ReplForJupyterImpl(
527529

528530
private val serializationQueue = LockQueue<SerializationReply, SerializationArgs>()
529531
override suspend fun serializeVariables(cellId: Int, descriptorsState: Map<String, SerializedVariablesState>, callback: (SerializationReply) -> Unit) {
530-
doWithLock(SerializationArgs(cellId, descriptorsState, callback), serializationQueue, SerializationReply(cellId, descriptorsState), ::doSerializeVariables)
532+
doWithLock(SerializationArgs(descriptorsState, cellId = cellId, callback = callback), serializationQueue, SerializationReply(cellId, descriptorsState), ::doSerializeVariables)
533+
}
534+
535+
override suspend fun serializeVariables(topLevelVarName: String, descriptorsState: Map<String, SerializedVariablesState>, callback: (SerializationReply) -> Unit) {
536+
doWithLock(SerializationArgs(descriptorsState, topLevelVarName = topLevelVarName, callback = callback), serializationQueue, SerializationReply(), ::doSerializeVariables)
531537
}
532538

533539
private fun doSerializeVariables(args: SerializationArgs): SerializationReply {
534540
val resultMap = mutableMapOf<String, SerializedVariablesState>()
541+
val cellId = if (args.cellId != -1) args.cellId else {
542+
val watcherInfo = internalEvaluator.findVariableCell(args.topLevelVarName) + 1
543+
val finalAns = if (watcherInfo == - 1) 1 else watcherInfo
544+
finalAns
545+
}
535546
args.descriptorsState.forEach { (name, state) ->
536-
resultMap[name] = variablesSerializer.doIncrementalSerialization(args.cellId - 1, name, state)
547+
resultMap[name] = variablesSerializer.doIncrementalSerialization(cellId - 1, name, state)
537548
}
538-
return SerializationReply(args.cellId, resultMap)
549+
log.debug("Serialization cellID: $cellId")
550+
log.debug("Serialization answer: ${resultMap.entries.first().value.fieldDescriptor}")
551+
return SerializationReply(cellId, resultMap)
539552
}
540553

541554

@@ -572,11 +585,13 @@ class ReplForJupyterImpl(
572585
LockQueueArgs<ListErrorsResult>
573586

574587
private data class SerializationArgs(
575-
val cellId: Int,
576588
val descriptorsState: Map<String, SerializedVariablesState>,
589+
var cellId: Int = -1,
590+
val topLevelVarName: String = "",
577591
override val callback: (SerializationReply) -> Unit
578592
) : LockQueueArgs<SerializationReply>
579593

594+
580595
@JvmInline
581596
private value class LockQueue<T, Args : LockQueueArgs<T>>(
582597
private val args: AtomicReference<Args?> = AtomicReference()

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,4 +30,14 @@ interface InternalEvaluator {
3030
* returns empty data or null
3131
*/
3232
fun popAddedCompiledScripts(): SerializedCompiledScriptsData = SerializedCompiledScriptsData.EMPTY
33+
34+
/**
35+
* Get a cellId where a particular variable is declared
36+
*/
37+
fun findVariableCell(variableName: String): Int
38+
39+
/**
40+
* Returns a set of unaffected variables after execution
41+
*/
42+
fun getUnchangedVariables(): Set<String>
3343
}

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,10 @@ internal class InternalEvaluatorImpl(
4848
return SerializedCompiledScriptsData(scripts)
4949
}
5050

51+
override fun findVariableCell(variableName: String): Int {
52+
return variablesWatcher.findDeclarationAddress(variableName) ?: -1
53+
}
54+
5155
override var writeCompiledClasses: Boolean
5256
get() = classWriter != null
5357
set(value) {

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

Lines changed: 33 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,14 @@
11
package org.jetbrains.kotlinx.jupyter
22

3+
import kotlinx.serialization.Serializable
4+
import kotlinx.serialization.json.Json
5+
import kotlinx.serialization.json.JsonObject
6+
import kotlinx.serialization.json.decodeFromJsonElement
37
import org.jetbrains.kotlinx.jupyter.api.VariableState
48
import org.jetbrains.kotlinx.jupyter.compiler.util.SerializedVariablesState
59
import java.lang.reflect.Field
10+
import kotlin.contracts.ExperimentalContracts
11+
import kotlin.contracts.contract
612
import kotlin.reflect.KClass
713
import kotlin.reflect.KProperty
814
import kotlin.reflect.KProperty1
@@ -23,6 +29,16 @@ enum class PropertiesType {
2329
MIXED
2430
}
2531

32+
@Serializable
33+
data class SerializedCommMessageContent(
34+
val topLevelDescriptorName: String,
35+
val descriptorsState: Map<String, SerializedVariablesState>
36+
)
37+
38+
fun getVariablesDescriptorsFromJson(json: JsonObject): SerializedCommMessageContent {
39+
return Json.decodeFromJsonElement<SerializedCommMessageContent>(json)
40+
}
41+
2642
class ProcessedSerializedVarsState(
2743
val serializedVariablesState: SerializedVariablesState,
2844
val propertiesData: PropertiesData? = null,
@@ -94,9 +110,11 @@ class VariablesSerializer(private val serializationDepth: Int = 2, private val s
94110
} == true
95111
}
96112

97-
val kProperties = if (value != null) value::class.declaredMemberProperties else {
98-
null
99-
}
113+
val kProperties = try {
114+
if (value != null) value::class.declaredMemberProperties else {
115+
null
116+
}
117+
} catch (ex: Exception) {null}
100118
val serializedVersion = SerializedVariablesState(simpleTypeName, getProperString(value), true)
101119
val descriptors = serializedVersion.fieldDescriptor
102120
if (isDescriptorsNeeded) {
@@ -198,6 +216,11 @@ class VariablesSerializer(private val serializationDepth: Int = 2, private val s
198216

199217
private val isSerializationActive: Boolean = System.getProperty(serializationSystemProperty)?.toBooleanStrictOrNull() ?: true
200218

219+
/**
220+
* Cache for not recomputing unchanged variables
221+
*/
222+
val serializedVariablesCache: MutableMap<String, SerializedVariablesState> = mutableMapOf()
223+
201224
fun serializeVariables(cellId: Int, variablesState: Map<String, VariableState>): Map<String, SerializedVariablesState> {
202225
if (!isSerializationActive) return emptyMap()
203226

@@ -375,6 +398,7 @@ class VariablesSerializer(private val serializationDepth: Int = 2, private val s
375398
* Really wanted to use contracts here, but all usages should be provided with this annotation and,
376399
* perhaps, it may be a big overhead
377400
*/
401+
@OptIn(ExperimentalContracts::class)
378402
private fun iterateThrough(
379403
elem: Any,
380404
seenObjectsPerCell: MutableMap<RuntimeObjectWrapper, SerializedVariablesState>?,
@@ -383,6 +407,10 @@ class VariablesSerializer(private val serializationDepth: Int = 2, private val s
383407
instancesPerState: MutableMap<SerializedVariablesState, Any?>,
384408
callInstance: Any
385409
) {
410+
contract {
411+
returns() implies (elem is Field || elem is KProperty1<*, *>)
412+
}
413+
386414
val name = if (elem is Field) elem.name else (elem as KProperty1<Any, *>).name
387415
val value = if (elem is Field) tryGetValueFromProperty(elem, callInstance).toObjectWrapper()
388416
else {
@@ -491,6 +519,8 @@ class VariablesSerializer(private val serializationDepth: Int = 2, private val s
491519
return value
492520
}
493521

522+
// use of Java 9 required
523+
@SuppressWarnings("DEPRECATION")
494524
private fun tryGetValueFromProperty(property: Field, callInstance: Any): Any? {
495525
// some fields may be optimized out like array size. Thus, calling it.isAccessible would return error
496526
val canAccess = try {

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,7 @@ fun ResultsRenderersProcessor.registerDefaultRenderers() {
7373
* Stores info about where a variable Y was declared and info about what are they at the address X.
7474
* K: key, stands for a way of addressing variables, e.g. address.
7575
* V: value, from Variable, choose any suitable type for your variable reference.
76-
* Default: T=Int, V=String
76+
* Default: K=Int, V=String
7777
*/
7878
class VariablesUsagesPerCellWatcher<K : Any, V : Any> {
7979
val cellVariables = mutableMapOf<K, MutableSet<V>>()
@@ -106,5 +106,7 @@ class VariablesUsagesPerCellWatcher<K : Any, V : Any> {
106106
}
107107
}
108108

109+
fun findDeclarationAddress(variableRef: V) = variablesDeclarationInfo[variableRef]
110+
109111
fun ensureStorageCreation(address: K) = cellVariables.putIfAbsent(address, mutableSetOf())
110112
}

0 commit comments

Comments
 (0)