PY-77295 Expand collapsed section on add new cell

(cherry picked from commit c6a8041ba6313e0e6decbed7bdff71fc50e286fc)

GitOrigin-RevId: 4731bd10527423659f965c91a0e01f852556fe0d

notebooks/visualization/src/com/intellij/notebooks/visualization/EditorNotebookExtension.kt

@@ -2,6 +2,6 @@ package com.intellij.notebooks.visualization
 
 import com.intellij.notebooks.visualization.ui.EditorCell
 
-interface CellExtensionFactory {
+interface EditorNotebookExtension {
   fun onCellCreated(cell: EditorCell)
 }

notebooks/visualization/src/com/intellij/notebooks/visualization/EditorNotebookPostprocessor.kt

@@ -0,0 +1,8 @@
+// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.notebooks.visualization
+
+import com.intellij.notebooks.visualization.ui.EditorNotebook
+
+interface EditorNotebookPostprocessor {
+  fun postprocess(editorNotebook: EditorNotebook)
+}

notebooks/visualization/src/com/intellij/notebooks/visualization/NotebookCellInlayManager.kt

@@ -3,7 +3,6 @@ package com.intellij.notebooks.visualization
 import com.intellij.ide.ui.LafManagerListener
 import com.intellij.notebooks.ui.isFoldingEnabledKey
 import com.intellij.notebooks.ui.visualization.NotebookBelowLastCellPanel
-import com.intellij.notebooks.visualization.inlay.JupyterBoundsChangeHandler
 import com.intellij.notebooks.visualization.ui.*
 import com.intellij.notebooks.visualization.ui.EditorCellEventListener.*
 import com.intellij.notebooks.visualization.ui.EditorCellViewEventListener.CellViewCreated
@@ -14,12 +13,10 @@ import com.intellij.openapi.application.runInEdt
 import com.intellij.openapi.diagnostic.thisLogger
 import com.intellij.openapi.editor.CustomFoldRegion
 import com.intellij.openapi.editor.CustomFoldRegionRenderer
-import com.intellij.openapi.editor.Document
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.editor.FoldRegion
 import com.intellij.openapi.editor.colors.EditorColorsListener
 import com.intellij.openapi.editor.colors.EditorColorsManager
-import com.intellij.openapi.editor.event.BulkAwareDocumentListener
 import com.intellij.openapi.editor.event.CaretEvent
 import com.intellij.openapi.editor.event.CaretListener
 import com.intellij.openapi.editor.ex.EditorEx
@@ -40,75 +37,27 @@ import com.intellij.util.concurrency.ThreadingAssertions
 class NotebookCellInlayManager private constructor(
   val editor: EditorImpl,
   private val shouldCheckInlayOffsets: Boolean,
-  private val cellExtensionFactories: List<CellExtensionFactory>,
+  private val notebook: EditorNotebook
 ) : Disposable, NotebookIntervalPointerFactory.ChangeListener {
 
   private val notebookCellLines = NotebookCellLines.get(editor)
 
   private var initialized = false
 
-  private var _cells = mutableListOf<EditorCell>()
-
-  val cells: List<EditorCell> get() = _cells.toList()
+  val cells: List<EditorCell> get() = notebook.cells
 
   val views = mutableMapOf<EditorCell, EditorCellView>()
 
   private var belowLastCellInlay: Inlay<*>? = null
 
-  /**
-   * Listens for inlay changes (called after all inlays are updated). Feel free to convert it to the EP if you need another listener
-   */
-  var changedListener: InlaysChangedListener? = null
-
-  private val cellEventListeners = EventDispatcher.create(EditorCellEventListener::class.java)
-
   private val cellViewEventListeners = EventDispatcher.create(EditorCellViewEventListener::class.java)
 
   private val invalidationListeners = mutableListOf<Runnable>()
 
   private var valid = false
 
-  private var updateCtx: UpdateContext? = null
-
-  /*
-   EditorImpl sets `myDocumentChangeInProgress` attribute to true during document update processing, that prevents correct update
-   of custom folding regions.When this flag is set, folding updates will be postponed until the editor finishes its work.
-   */
-  private var editorIsProcessingDocument = false
-
-  private var postponedUpdates = mutableListOf<UpdateContext>()
-
-  fun <T> update(force: Boolean = false, block: (updateCtx: UpdateContext) -> T): T {
-    val ctx = updateCtx
-    return if (ctx != null) {
-      block(ctx)
-    }
-    else {
-      val newCtx = UpdateContext(force)
-      updateCtx = newCtx
-      try {
-        val jupyterBoundsChangeHandler = JupyterBoundsChangeHandler.get(editor)
-        jupyterBoundsChangeHandler.postponeUpdates()
-        val r = keepScrollingPositionWhile(editor) {
-          val r = block(newCtx)
-          updateCtx = null
-          if (editorIsProcessingDocument) {
-            postponedUpdates.add(newCtx)
-          }
-          else {
-            newCtx.applyUpdates(editor)
-          }
-          r
-        }
-        inlaysChanged()
-        jupyterBoundsChangeHandler.boundsChanged()
-        jupyterBoundsChangeHandler.performPostponed()
-        r
-      }
-      finally {
-        updateCtx = null
-      }
-    }
+  private fun update(force: Boolean = false, block: (UpdateContext) -> Unit) {
+    editor.updateManager.update(force, block)
   }
 
   override fun dispose() {
@@ -116,11 +65,11 @@ class NotebookCellInlayManager private constructor(
   }
 
   fun getCellForInterval(interval: NotebookCellLines.Interval): EditorCell =
-    _cells[interval.ordinal]
+    notebook.cells[interval.ordinal]
 
   fun updateAllOutputs() {
     update {
-      _cells.forEach {
+      notebook.cells.forEach {
         it.updateOutputs()
       }
     }
@@ -161,7 +110,7 @@ class NotebookCellInlayManager private constructor(
 
   private fun addViewportChangeListener() {
     editor.scrollPane.viewport.addChangeListener {
-      _cells.forEach {
+      notebook.cells.forEach {
         it.onViewportChange()
       }
     }
@@ -192,26 +141,12 @@ class NotebookCellInlayManager private constructor(
     setupSelectionUI()
     addBelowLastCellInlay()
 
-    cellEventListeners.addListener(object : EditorCellEventListener {
+    notebook.addCellEventsListener(this, object : EditorCellEventListener {
       override fun onEditorCellEvents(events: List<EditorCellEvent>) {
         updateUI(events)
       }
     })
 
-    editor.document.addDocumentListener(object : BulkAwareDocumentListener.Simple {
-      override fun beforeDocumentChange(document: Document) {
-        editorIsProcessingDocument = true
-      }
-
-      override fun afterDocumentChange(document: Document) {
-        editorIsProcessingDocument = false
-        postponedUpdates.forEach {
-          it.applyUpdates(editor)
-        }
-        postponedUpdates.clear()
-      }
-    }, this)
-
     handleRefreshedDocument()
   }
 
@@ -340,7 +275,7 @@ class NotebookCellInlayManager private constructor(
     }, this)
   }
 
-  private fun editorCells(region: FoldRegion): List<EditorCell> = _cells.filter { cell ->
+  private fun editorCells(region: FoldRegion): List<EditorCell> = notebook.cells.filter { cell ->
     val startOffset = editor.document.getLineStartOffset(cell.intervalPointer.get()!!.lines.first)
     val endOffset = editor.document.getLineEndOffset(cell.intervalPointer.get()!!.lines.last)
     startOffset >= region.startOffset && endOffset <= region.endOffset
@@ -348,28 +283,14 @@ class NotebookCellInlayManager private constructor(
 
   private fun handleRefreshedDocument() {
     ThreadingAssertions.softAssertReadAccess()
-    _cells.forEach {
-      Disposer.dispose(it)
-    }
+    notebook.clear()
     val pointerFactory = NotebookIntervalPointerFactory.get(editor)
 
     update {
-      _cells = notebookCellLines.intervals.map { interval ->
-        createCell(pointerFactory.create(interval))
-      }.toMutableList()
+      notebookCellLines.intervals.forEach { interval ->
+        notebook.addCell(pointerFactory.create(interval))
+      }
     }
-    cellEventListeners.multicaster.onEditorCellEvents(_cells.map { CellCreated(it) })
-  }
-
-  private fun createCell(interval: NotebookIntervalPointer) = EditorCell(editor, this, interval).also {
-    cellExtensionFactories.forEach { factory ->
-      factory.onCellCreated(it)
-    }
-    Disposer.register(this, it)
-  }
-
-  private fun inlaysChanged() {
-    changedListener?.inlaysChanged()
   }
 
   private fun updateCellsFolding(editorCells: List<EditorCell>) = update { updateContext ->
@@ -382,13 +303,16 @@ class NotebookCellInlayManager private constructor(
     fun install(
       editor: EditorImpl,
       shouldCheckInlayOffsets: Boolean,
-      cellExtensionFactories: List<CellExtensionFactory> = listOf(),
+      editorNotebookPostprocessors: List<EditorNotebookPostprocessor> = listOf(),
     ): NotebookCellInlayManager {
       EditorEmbeddedComponentContainer(editor as EditorEx)
+      val updateManager = UpdateManager(editor)
+      Disposer.register(editor.disposable, updateManager)
+      val notebook = createNotebook(editor, editorNotebookPostprocessors)
       val notebookCellInlayManager = NotebookCellInlayManager(
         editor,
         shouldCheckInlayOffsets,
-        cellExtensionFactories
+        notebook
       ).also { Disposer.register(editor.disposable, it) }
       editor.putUserData(isFoldingEnabledKey, Registry.`is`("jupyter.editor.folding.cells"))
       notebookCellInlayManager.initialize()
@@ -396,6 +320,18 @@ class NotebookCellInlayManager private constructor(
       return notebookCellInlayManager
     }
 
+    private fun createNotebook(
+      editor: EditorImpl,
+      editorNotebookPostprocessors: List<EditorNotebookPostprocessor>,
+    ): EditorNotebook {
+      val notebook = EditorNotebook(editor)
+      editorNotebookPostprocessors.forEach {
+        it.postprocess(notebook)
+      }
+      Disposer.register(editor.disposable, notebook)
+      return notebook
+    }
+
     /** NotebookCellInlayManager exist only on Front in RemoteDev. */
     fun get(editor: Editor): NotebookCellInlayManager? {
       return CELL_INLAY_MANAGER_KEY.get(editor)
@@ -406,29 +342,27 @@ class NotebookCellInlayManager private constructor(
   }
 
   override fun onUpdated(event: NotebookIntervalPointersEvent) = update { ctx ->
-    val events = mutableListOf<EditorCellEvent>()
     for (change in event.changes) {
       when (change) {
         is NotebookIntervalPointersEvent.OnEdited -> {
-          val cell = _cells[change.intervalAfter.ordinal]
+          val cell = notebook.cells[change.intervalAfter.ordinal]
           cell.updateInput()
         }
         is NotebookIntervalPointersEvent.OnInserted -> {
           change.subsequentPointers.forEach {
-            val editorCell = createCell(it.pointer)
-            addCell(it.interval.ordinal, editorCell, events)
+            addCell(it.pointer)
           }
         }
         is NotebookIntervalPointersEvent.OnRemoved -> {
           change.subsequentPointers.reversed().forEach {
             val index = it.interval.ordinal
-            removeCell(index, events)
+            removeCell(index)
           }
         }
         is NotebookIntervalPointersEvent.OnSwapped -> {
-          val firstCell = _cells[change.firstOrdinal]
+          val firstCell = notebook.cells[change.firstOrdinal]
           val first = firstCell.intervalPointer
-          val secondCell = _cells[change.secondOrdinal]
+          val secondCell = notebook.cells[change.secondOrdinal]
           firstCell.intervalPointer = secondCell.intervalPointer
           secondCell.intervalPointer = first
           firstCell.update(ctx)
@@ -439,7 +373,7 @@ class NotebookCellInlayManager private constructor(
     event.changes.filterIsInstance<NotebookIntervalPointersEvent.OnInserted>().forEach { change ->
       fixInlaysOffsetsAfterNewCellInsert(change, ctx)
     }
-    cellEventListeners.multicaster.onEditorCellEvents(events)
+
     checkInlayOffsets()
   }
 
@@ -447,13 +381,13 @@ class NotebookCellInlayManager private constructor(
     if (!shouldCheckInlayOffsets) return
 
     val inlaysOffsets = buildSet {
-      for (cell in _cells) {
+      for (cell in notebook.cells) {
         add(editor.document.getLineStartOffset(cell.interval.lines.first))
         add(editor.document.getLineEndOffset(cell.interval.lines.last))
       }
     }
 
-    val wronglyPlacedInlays = _cells.asSequence()
+    val wronglyPlacedInlays = notebook.cells.asSequence()
       .mapNotNull { it.view }
       .flatMap { it.getInlays() }
       .filter { it.offset !in inlaysOffsets }
@@ -471,23 +405,18 @@ class NotebookCellInlayManager private constructor(
     }
   }
 
-  private fun addCell(index: Int, editorCell: EditorCell, events: MutableList<EditorCellEvent>) {
-    _cells.add(index, editorCell)
-    events.add(CellCreated(editorCell))
+  private fun addCell(pointer: NotebookIntervalPointer) {
+    notebook.addCell(pointer)
     invalidateCells()
   }
 
-  private fun removeCell(index: Int, events: MutableList<EditorCellEvent>) {
-    val cell = _cells[index]
-    cell.onBeforeRemove()
-    val removed = _cells.removeAt(index)
-    Disposer.dispose(removed)
-    events.add(CellRemoved(removed))
+  private fun removeCell(index: Int) {
+    notebook.removeCell(index)
     invalidateCells()
   }
 
   fun addCellEventsListener(editorCellEventListener: EditorCellEventListener, disposable: Disposable) {
-    cellEventListeners.addListener(editorCellEventListener, disposable)
+    notebook.addCellEventsListener(disposable, editorCellEventListener)
   }
 
   fun addCellViewEventsListener(editorCellViewEventListener: EditorCellViewEventListener, disposable: Disposable) {
@@ -526,6 +455,7 @@ class NotebookCellInlayManager private constructor(
   }
 }
 
+
 class UpdateContext(val force: Boolean = false) {
 
   private val foldingOperations = mutableListOf<(FoldingModelEx) -> Unit>()

notebooks/visualization/src/com/intellij/notebooks/visualization/observables/DeduplicatedObservableProperty.kt

@@ -0,0 +1,31 @@
+package com.intellij.notebooks.visualization.observables
+
+import com.intellij.openapi.Disposable
+import com.intellij.openapi.observable.dispatcher.SingleEventDispatcher
+import com.intellij.openapi.observable.properties.ObservableMutableProperty
+
+class DeduplicatedObservableProperty<T>(initialValue: T): ObservableMutableProperty<T> {
+
+  private var value: T = initialValue
+
+  private val changeDispatcher = SingleEventDispatcher.create<T>()
+
+  private fun fireChangeEvent(oldValue: T, newValue: T) {
+    if (oldValue != newValue) {
+      changeDispatcher.fireEvent(newValue)
+    }
+  }
+
+  override fun set(value: T) {
+    val oldValue = this.value
+    this.value = value
+    fireChangeEvent(oldValue, value)
+  }
+
+  override fun get(): T {
+    return this.value
+  }
+
+  override fun afterChange(parentDisposable: Disposable?, listener: (T) -> Unit): Unit =
+    changeDispatcher.whenEventHappened(parentDisposable, listener)
+}

notebooks/visualization/src/com/intellij/notebooks/visualization/ui/CustomFoldingEditorCellViewComponent.kt

@@ -50,7 +50,7 @@ class CustomFoldingEditorCellViewComponent(
   }
 
   private fun updateGutterIcons(gutterAction: AnAction?) {
-    cell.manager.update { ctx ->
+    editor.updateManager.update { ctx ->
       gutterActionRenderer = gutterAction?.let { ActionToGutterRendererAdapter(it) }
       ctx.addFoldingOperation { modelEx ->
         foldingRegion?.update()
@@ -65,7 +65,7 @@ class CustomFoldingEditorCellViewComponent(
     updateGutterIcons(cell.gutterAction.get())
   }
 
-  override fun dispose() = cell.manager.update { ctx ->
+  override fun dispose() = editor.updateManager.update { ctx ->
     disposeFolding(ctx)
   }
 

notebooks/visualization/src/com/intellij/notebooks/visualization/ui/EditorCell.kt

@@ -21,9 +21,9 @@ import kotlin.reflect.KClass
 private val CELL_EXTENSION_CONTAINER_KEY = Key<MutableMap<KClass<*>, EditorCellExtension>>("CELL_EXTENSION_CONTAINER_KEY")
 
 class EditorCell(
-  private val editor: EditorEx,
-  val manager: NotebookCellInlayManager,
+  val notebook: EditorNotebook,
   var intervalPointer: NotebookIntervalPointer,
+  private val editor: EditorImpl,
 ) : Disposable, UserDataHolder by UserDataHolderBase() {
 
   val source = AtomicProperty<String>(getSource())
@@ -33,7 +33,7 @@ class EditorCell(
   val interval get() = intervalPointer.get() ?: error("Invalid interval")
 
   val view: EditorCellView?
-    get() = manager.views[this]
+    get() = NotebookCellInlayManager.get(editor)!!.views[this]
 
   var visible = AtomicBooleanProperty(true)
 
@@ -67,7 +67,7 @@ class EditorCell(
   }
 
   fun update() {
-    manager.update { ctx -> update(ctx) }
+    editor.updateManager.update { ctx -> update(ctx) }
   }
 
   fun update(updateCtx: UpdateContext) {
@@ -97,7 +97,7 @@ class EditorCell(
   @PublishedApi
   internal fun createLazyControllers(factory: NotebookCellInlayController.LazyFactory) {
     factory.cellOrdinalsInCreationBlock.add(interval.ordinal)
-    manager.update { ctx ->
+    editor.updateManager.update { ctx ->
       update(ctx)
     }
     factory.cellOrdinalsInCreationBlock.remove(interval.ordinal)
@@ -110,7 +110,7 @@ class EditorCell(
       .firstOrNull { it.getControllerClass() == type.java }
   }
 
-  fun updateOutputs() {
+  fun updateOutputs() = editor.updateManager.update {
     outputs.updateOutputs()
   }
 

notebooks/visualization/src/com/intellij/notebooks/visualization/ui/EditorCellInput.kt

@@ -50,7 +50,7 @@ class EditorCellInput(
     return bounds.y + delimiterPanelSize to bounds.height - delimiterPanelSize
   }
 
-  private fun toggleFolding() = cell.manager.update { ctx ->
+  private fun toggleFolding() = editor.updateManager.update { ctx ->
     folded = !folded
     (component as? InputComponent)?.updateFolding(ctx, folded)
   }
@@ -78,7 +78,7 @@ class EditorCellInput(
     return bounds
   }
 
-  fun updateInput() = cell.manager.update { ctx ->
+  fun updateInput() = editor.updateManager.update { ctx ->
     (component as? InputComponent)?.updateInput(ctx)
   }
 

notebooks/visualization/src/com/intellij/notebooks/visualization/ui/EditorCellView.kt

@@ -148,7 +148,7 @@ class EditorCellView(
     }
   }
 
-  private fun recreateControllers() = cell.manager.update { updateContext ->
+  private fun recreateControllers() = editor.updateManager.update { updateContext ->
     updateContext.addInlayOperation {
       val otherFactories = NotebookCellInlayController.Factory.EP_NAME.extensionList
         .filter { it !is InputFactory }

notebooks/visualization/src/com/intellij/notebooks/visualization/ui/EditorNotebook.kt

@@ -0,0 +1,79 @@
+// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.notebooks.visualization.ui
+
+import com.intellij.notebooks.visualization.EditorNotebookExtension
+import com.intellij.notebooks.visualization.NotebookIntervalPointer
+import com.intellij.notebooks.visualization.ui.EditorCellEventListener.CellCreated
+import com.intellij.notebooks.visualization.ui.EditorCellEventListener.CellRemoved
+import com.intellij.openapi.Disposable
+import com.intellij.openapi.editor.impl.EditorImpl
+import com.intellij.openapi.util.Disposer
+import com.intellij.openapi.util.Disposer.register
+import com.intellij.util.EventDispatcher
+import kotlin.reflect.KClass
+
+class EditorNotebook(private val editor: EditorImpl): Disposable {
+
+  private var _cells = mutableListOf<EditorCell>()
+
+  val cells: List<EditorCell> get() = _cells.toList()
+
+  private val cellEventListeners = EventDispatcher.create(EditorCellEventListener::class.java)
+
+  private val extensions = mutableMapOf<KClass<*>, EditorNotebookExtension>()
+
+  @Suppress("UNCHECKED_CAST")
+  fun <T : EditorNotebookExtension> getExtension(cls: KClass<T>): T? {
+    return extensions[cls] as? T
+  }
+
+  private fun forEachExtension(action: (EditorNotebookExtension) -> Unit) {
+    extensions.values.forEach { action(it) }
+  }
+
+  fun addCellEventsListener(disposable: Disposable, listener: EditorCellEventListener) {
+    cellEventListeners.addListener(listener, disposable)
+  }
+
+  fun addCell(interval: NotebookIntervalPointer) {
+    val editorCell = EditorCell(this, interval, editor).also {
+      forEachExtension { extension ->
+        extension.onCellCreated(it)
+      }
+      register(this, it)
+    }
+    _cells.add(interval.get()!!.ordinal, editorCell)
+    cellEventListeners.multicaster.onEditorCellEvents(listOf(CellCreated(editorCell)))
+  }
+
+  override fun dispose() {
+    extensions.values.forEach {
+      if (it is Disposable) {
+        Disposer.dispose(it)
+      }
+    }
+  }
+
+  fun clear() {
+    _cells.forEach { cell ->
+      Disposer.dispose(cell)
+    }
+    _cells.clear()
+  }
+
+  fun removeCell(index: Int) {
+    val cell = _cells[index]
+    cell.onBeforeRemove()
+    val removed = _cells.removeAt(index)
+    Disposer.dispose(removed)
+    cellEventListeners.multicaster.onEditorCellEvents(listOf(CellRemoved(removed)))
+  }
+
+  fun <T: EditorNotebookExtension> addExtension(type: KClass<T>, extension: T) {
+    extensions[type] = extension
+  }
+}
+
+inline fun <reified T : EditorNotebookExtension> EditorNotebook.getExtension(): T? {
+  return getExtension(T::class)
+}

notebooks/visualization/src/com/intellij/notebooks/visualization/ui/TextEditorCellViewComponent.kt

@@ -83,7 +83,7 @@ class TextEditorCellViewComponent(
     this.highlighters = listOf(highlighter)
   }
 
-  override fun dispose() = cell.manager.update { ctx ->
+  override fun dispose() = editor.updateManager.update { ctx ->
     disposeExistingHighlighter()
     editor.contentComponent.removeMouseListener(mouseListener)
   }

notebooks/visualization/src/com/intellij/notebooks/visualization/ui/UpdateManager.kt

@@ -0,0 +1,99 @@
+// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.notebooks.visualization.ui
+
+import com.intellij.notebooks.visualization.InlaysChangedListener
+import com.intellij.notebooks.visualization.UpdateContext
+import com.intellij.notebooks.visualization.inlay.JupyterBoundsChangeHandler
+import com.intellij.openapi.Disposable
+import com.intellij.openapi.editor.Document
+import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.editor.event.BulkAwareDocumentListener
+import com.intellij.openapi.editor.impl.EditorImpl
+import com.intellij.openapi.util.Key
+
+class UpdateManager(val editor: EditorImpl) : Disposable {
+
+  private var updateCtx: UpdateContext? = null
+
+  /*
+   EditorImpl sets `myDocumentChangeInProgress` attribute to true during document update processing, that prevents correct update
+   of custom folding regions.When this flag is set, folding updates will be postponed until the editor finishes its work.
+   */
+  private var editorIsProcessingDocument = false
+
+  private var postponedUpdates = mutableListOf<UpdateContext>()
+
+  /**
+   * Listens for inlay changes (called after all inlays are updated). Feel free to convert it to the EP if you need another listener
+   */
+  var changedListener: InlaysChangedListener? = null
+
+  init {
+    editor.document.addDocumentListener(object : BulkAwareDocumentListener.Simple {
+      override fun beforeDocumentChange(document: Document) {
+        editorIsProcessingDocument = true
+      }
+
+      override fun afterDocumentChange(document: Document) {
+        editorIsProcessingDocument = false
+        postponedUpdates.forEach {
+          it.applyUpdates(editor)
+        }
+        postponedUpdates.clear()
+        finalizeChanges()
+      }
+    }, this)
+    UPDATE_MANAGER_KEY.set(editor, this)
+  }
+
+  fun <T> update(force: Boolean = false, block: (updateCtx: UpdateContext) -> T): T {
+    val ctx = updateCtx
+    return if (ctx != null) {
+      block(ctx)
+    }
+    else {
+      val newCtx = UpdateContext(force)
+      updateCtx = newCtx
+      try {
+        val jupyterBoundsChangeHandler = JupyterBoundsChangeHandler.get(editor)
+        jupyterBoundsChangeHandler.postponeUpdates()
+        val r = keepScrollingPositionWhile(editor) {
+          val r = block(newCtx)
+          updateCtx = null
+          if (editorIsProcessingDocument) {
+            postponedUpdates.add(newCtx)
+          }
+          else {
+            newCtx.applyUpdates(editor)
+            finalizeChanges()
+          }
+          r
+        }
+        r
+      }
+      finally {
+        updateCtx = null
+      }
+    }
+  }
+
+  private fun finalizeChanges() {
+    inlaysChanged()
+    val jupyterBoundsChangeHandler = JupyterBoundsChangeHandler.get(editor)
+    jupyterBoundsChangeHandler.boundsChanged()
+    jupyterBoundsChangeHandler.performPostponed()
+  }
+
+  private fun inlaysChanged() {
+    changedListener?.inlaysChanged()
+  }
+
+  override fun dispose() {
+
+  }
+}
+
+private val UPDATE_MANAGER_KEY = Key<UpdateManager>("UPDATE_MANAGER_KEY")
+
+val Editor.updateManager
+  get() = UPDATE_MANAGER_KEY.get(this)