Add suspended cache validation; Add possibility to remove old variables from cache; some improvements

Author: nikolay-egorov

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

@@ -578,7 +578,8 @@ class ListErrorsReply(
 class SerializationRequest(
     val cellId: Int,
     val descriptorsState: Map<String, SerializedVariablesState>,
-    val topLevelDescriptorName: String = ""
+    val topLevelDescriptorName: String = "",
+    val pathToDescriptor: List<String> = emptyList()
 ) : MessageContent()
 
 @Serializable

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

@@ -315,7 +315,11 @@ fun JupyterConnection.Socket.shellMessagesHandler(msg: Message, repl: ReplForJup
 
             val messageContent = getVariablesDescriptorsFromJson(data)
             GlobalScope.launch(Dispatchers.Default) {
-                repl.serializeVariables(messageContent.topLevelDescriptorName, messageContent.descriptorsState) { result ->
+                repl.serializeVariables(
+                    messageContent.topLevelDescriptorName,
+                    messageContent.descriptorsState,
+                    messageContent.pathToDescriptor
+                ) { result ->
                     sendWrapped(msg, makeReplyMessage(msg, MessageType.COMM_OPEN, content = result))
                 }
             }
@@ -337,7 +341,7 @@ fun JupyterConnection.Socket.shellMessagesHandler(msg: Message, repl: ReplForJup
         is SerializationRequest -> {
             GlobalScope.launch(Dispatchers.Default) {
                 if (content.topLevelDescriptorName.isNotEmpty()) {
-                    repl.serializeVariables(content.topLevelDescriptorName, content.descriptorsState) { result ->
+                    repl.serializeVariables(content.topLevelDescriptorName, content.descriptorsState, content.pathToDescriptor) { result ->
                         sendWrapped(msg, makeReplyMessage(msg, MessageType.SERIALIZATION_REPLY, content = result))
                     }
                 } else {

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

@@ -5,6 +5,10 @@ import jupyter.kotlin.DependsOn
 import jupyter.kotlin.KotlinContext
 import jupyter.kotlin.KotlinKernelHostProvider
 import jupyter.kotlin.Repository
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.GlobalScope
+import kotlinx.coroutines.launch
+import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.config.KotlinCompilerVersion
 import org.jetbrains.kotlinx.jupyter.api.Code
 import org.jetbrains.kotlinx.jupyter.api.ExecutionCallback
@@ -48,6 +52,7 @@ import org.jetbrains.kotlinx.jupyter.repl.CellExecutor
 import org.jetbrains.kotlinx.jupyter.repl.CompletionResult
 import org.jetbrains.kotlinx.jupyter.repl.ContextUpdater
 import org.jetbrains.kotlinx.jupyter.repl.EvalResult
+import org.jetbrains.kotlinx.jupyter.repl.EvalResultEx
 import org.jetbrains.kotlinx.jupyter.repl.InternalEvaluator
 import org.jetbrains.kotlinx.jupyter.repl.KotlinCompleter
 import org.jetbrains.kotlinx.jupyter.repl.ListErrorsResult
@@ -120,7 +125,8 @@ interface ReplForJupyter {
 
     suspend fun serializeVariables(cellId: Int, descriptorsState: Map<String, SerializedVariablesState>, callback: (SerializationReply) -> Unit)
 
-    suspend fun serializeVariables(topLevelVarName: String, descriptorsState: Map<String, SerializedVariablesState>, callback: (SerializationReply) -> Unit)
+    suspend fun serializeVariables(topLevelVarName: String, descriptorsState: Map<String, SerializedVariablesState>, pathToDescriptor: List<String> = emptyList(),
+                                   callback: (SerializationReply) -> Unit)
 
     val homeDir: File?
 
@@ -191,7 +197,7 @@ class ReplForJupyterImpl(
 
     override val variablesSerializer = VariablesSerializer()
 
-    private val librariesScanner = LibrariesScanner(notebook)
+    val librariesScanner = LibrariesScanner(notebook)
     private val resourcesProcessor = LibraryResourcesProcessorImpl()
 
     override var outputConfig
@@ -347,7 +353,7 @@ class ReplForJupyterImpl(
     )
 
     private var evalContextEnabled = false
-    private fun withEvalContext(action: () -> EvalResult): EvalResult {
+    private fun <T> withEvalContext(action: () -> T): T {
         return synchronized(this) {
             evalContextEnabled = true
             try {
@@ -365,14 +371,14 @@ class ReplForJupyterImpl(
         else context.compilationConfiguration.asSuccess()
     }
 
-    /**
-     * Used for debug purposes.
-     * @see ReplCommand
-     */
+    @TestOnly
+    @Suppress("unused")
     private fun printVariables(isHtmlFormat: Boolean = false) = log.debug(
         if (isHtmlFormat) notebook.variablesReportAsHTML() else notebook.variablesReport()
     )
 
+    @TestOnly
+    @Suppress("unused")
     private fun printUsagesInfo(cellId: Int, usedVariables: Set<String>?) {
         log.debug(buildString {
             if (usedVariables == null || usedVariables.isEmpty()) {
@@ -386,7 +392,7 @@ class ReplForJupyterImpl(
         })
     }
 
-    fun evalEx(code: Code, displayHandler: DisplayHandler?, jupyterId: Int): EvalResult {
+    fun evalEx(code: Code, displayHandler: DisplayHandler?, jupyterId: Int): EvalResultEx {
         return withEvalContext {
             rethrowAsLibraryException(LibraryProblemPart.BEFORE_CELL_CALLBACKS) {
                 beforeCellExecution.forEach { executor.execute(it) }
@@ -426,8 +432,15 @@ class ReplForJupyterImpl(
             // printUsagesInfo(jupyterId, cellVariables[jupyterId - 1])
             val serializedData = variablesSerializer.serializeVariables(jupyterId - 1, notebook.variablesState, notebook.unchangedVariables())
 
-            EvalResult(
+            GlobalScope.launch(Dispatchers.Default) {
+                variablesSerializer.tryValidateCache(jupyterId - 1, notebook.cellVariables)
+            }
+
+            EvalResultEx(
                 result.result.value,
+                rendered,
+                result.scriptInstance,
+                result.result.name,
                 EvaluatedSnippetMetadata(newClasspath, compiledData, newImports, serializedData),
             )
         }
@@ -525,8 +538,9 @@ class ReplForJupyterImpl(
         doWithLock(SerializationArgs(descriptorsState, cellId = cellId, callback = callback), serializationQueue, SerializationReply(cellId, descriptorsState), ::doSerializeVariables)
     }
 
-    override suspend fun serializeVariables(topLevelVarName: String, descriptorsState: Map<String, SerializedVariablesState>, callback: (SerializationReply) -> Unit) {
-        doWithLock(SerializationArgs(descriptorsState, topLevelVarName = topLevelVarName, callback = callback), serializationQueue, SerializationReply(), ::doSerializeVariables)
+    override suspend fun serializeVariables(topLevelVarName: String, descriptorsState: Map<String, SerializedVariablesState>, pathToDescriptor: List<String>,
+                                            callback: (SerializationReply) -> Unit) {
+        doWithLock(SerializationArgs(descriptorsState, topLevelVarName = topLevelVarName, callback = callback, pathToDescriptor = pathToDescriptor), serializationQueue, SerializationReply(), ::doSerializeVariables)
     }
 
     private fun doSerializeVariables(args: SerializationArgs): SerializationReply {
@@ -537,7 +551,7 @@ class ReplForJupyterImpl(
             finalAns
         }
         args.descriptorsState.forEach { (name, state) ->
-            resultMap[name] = variablesSerializer.doIncrementalSerialization(cellId - 1, name, state)
+            resultMap[name] = variablesSerializer.doIncrementalSerialization(cellId - 1, name, state, args.pathToDescriptor)
         }
         log.debug("Serialization cellID: $cellId")
         log.debug("Serialization answer: ${resultMap.entries.first().value.fieldDescriptor}")
@@ -581,6 +595,7 @@ class ReplForJupyterImpl(
         val descriptorsState: Map<String, SerializedVariablesState>,
         var cellId: Int = -1,
         val topLevelVarName: String = "",
+        val pathToDescriptor: List<String> = emptyList(),
         override val callback: (SerializationReply) -> Unit
     ) : LockQueueArgs<SerializationReply>
 

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

@@ -151,7 +151,6 @@ internal class InternalEvaluatorImpl(
 
     private fun updateVariablesState(cellId: Int) {
         variablesWatcher.removeOldUsages(cellId)
-
         variablesHolder.forEach {
             val state = it.value as VariableStateImpl
             val oldValue = state.stringValue
@@ -173,11 +172,23 @@ internal class InternalEvaluatorImpl(
             it.name
         }.toHashSet()
         val ans = mutableMapOf<String, VariableStateImpl>()
+        // maybe remove known declarations
+        val addedDeclarations = mutableSetOf<String>()
+
         fields.forEach { property ->
             if (!memberKPropertiesNames.contains(property.name)) return@forEach
 
             val state = VariableStateImpl(property, cellClassInstance)
             variablesWatcher.addDeclaration(cellId, property.name)
+            addedDeclarations.add(property.name)
+
+            // try check values
+            if (variablesHolder.containsKey(property.name)) {
+                val seenState = variablesHolder[property.name]
+                if (seenState?.value?.equals(state.value) == true) {
+                    addedDeclarations.remove(property.name)
+                }
+            }
 
             // it was val, now it's var
             if (isValField(property)) {
@@ -189,6 +200,9 @@ internal class InternalEvaluatorImpl(
 
             ans[property.name] = state
         }
+        // remove old
+        variablesWatcher.removeOldDeclarations(cellId, addedDeclarations)
+
         return ans
     }
 
@@ -199,7 +213,7 @@ internal class InternalEvaluatorImpl(
     private fun updateDataAfterExecution(lastExecutionCellId: Int, resultValue: ResultValue) {
         variablesWatcher.ensureStorageCreation(lastExecutionCellId)
         variablesHolder += getVisibleVariables(resultValue, lastExecutionCellId)
-
+        // remove unreached variables
         updateVariablesState(lastExecutionCellId)
     }
 }