DS-6503 Fix invalid interval error on cell removal

GitOrigin-RevId: 6c3c08d556a116d884a89f8ef452be0340579204

notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/NotebookCellInlayManager.kt

@@ -20,10 +20,10 @@ import com.intellij.util.EventDispatcher
 import com.intellij.util.SmartList
 import com.intellij.util.concurrency.ThreadingAssertions
 import org.jetbrains.plugins.notebooks.ui.isFoldingEnabledKey
-import org.jetbrains.plugins.notebooks.visualization.ui.EditorCell
-import org.jetbrains.plugins.notebooks.visualization.ui.EditorCellEventListener
+import org.jetbrains.plugins.notebooks.visualization.ui.*
 import org.jetbrains.plugins.notebooks.visualization.ui.EditorCellEventListener.*
-import org.jetbrains.plugins.notebooks.visualization.ui.EditorCellView
+import org.jetbrains.plugins.notebooks.visualization.ui.EditorCellViewEventListener.CellViewCreated
+import org.jetbrains.plugins.notebooks.visualization.ui.EditorCellViewEventListener.CellViewRemoved
 import org.jetbrains.plugins.notebooks.visualization.ui.keepScrollingPositionWhile
 import java.util.*
 
@@ -46,6 +46,8 @@ class NotebookCellInlayManager private constructor(
 
   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
@@ -230,7 +232,7 @@ class NotebookCellInlayManager private constructor(
     inlaysChanged()
   }
 
-  private fun createCell(interval: NotebookIntervalPointer) = EditorCell(editor, interval) { cell ->
+  private fun createCell(interval: NotebookIntervalPointer) = EditorCell(editor, this, interval) { cell ->
     EditorCellView(editor, notebookCellLines, cell, this).also { Disposer.register(cell, it) }
   }.also { Disposer.register(this, it) }
 
@@ -349,6 +351,18 @@ class NotebookCellInlayManager private constructor(
     cellEventListeners.addListener(editorCellEventListener, disposable)
   }
 
+  fun addCellViewEventsListener(editorCellViewEventListener: EditorCellViewEventListener, disposable: Disposable) {
+    cellViewEventListeners.addListener(editorCellViewEventListener, disposable)
+  }
+
+  internal fun fireCellViewCreated(cellView: EditorCellView) {
+    cellViewEventListeners.multicaster.onEditorCellViewEvents(listOf(CellViewCreated(cellView)))
+  }
+
+  internal fun fireCellViewRemoved(cellView: EditorCellView) {
+    cellViewEventListeners.multicaster.onEditorCellViewEvents(listOf(CellViewRemoved(cellView)))
+  }
+
   fun getCell(index: Int): EditorCell {
     return cells[index]
   }

notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/ui/DecoratedEditor.kt

@@ -14,6 +14,8 @@ import com.intellij.openapi.util.Disposer
 import com.intellij.openapi.util.use
 import com.intellij.openapi.vfs.VirtualFile
 import org.jetbrains.plugins.notebooks.visualization.NotebookCellInlayManager
+import org.jetbrains.plugins.notebooks.visualization.ui.EditorCellViewEventListener.CellViewRemoved
+import org.jetbrains.plugins.notebooks.visualization.ui.EditorCellViewEventListener.EditorCellViewEvent
 import java.awt.AWTEvent
 import java.awt.BorderLayout
 import java.awt.GraphicsEnvironment
@@ -41,6 +43,15 @@ private class DecoratedEditor(private val original: TextEditor, private val mana
     manager.onInvalidate {
       component.revalidate()
     }
+    manager.addCellViewEventsListener(object : EditorCellViewEventListener {
+      override fun onEditorCellViewEvents(events: List<EditorCellViewEvent>) {
+        events.asSequence().filterIsInstance<CellViewRemoved>().forEach {
+          if (it.view == mouseOverCell) {
+            mouseOverCell = null
+          }
+        }
+      }
+    }, this)
   }
 
   private fun setupScrollPaneListener() {

notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/ui/EditorCell.kt

@@ -8,6 +8,7 @@ import com.intellij.openapi.util.TextRange
 import com.intellij.openapi.util.UserDataHolder
 import com.intellij.openapi.util.UserDataHolderBase
 import org.jetbrains.plugins.notebooks.visualization.NotebookCellInlayController
+import org.jetbrains.plugins.notebooks.visualization.NotebookCellInlayManager
 import org.jetbrains.plugins.notebooks.visualization.NotebookCellLines
 import org.jetbrains.plugins.notebooks.visualization.NotebookIntervalPointer
 import org.jetbrains.plugins.notebooks.visualization.execution.ExecutionEvent
@@ -17,6 +18,7 @@ import kotlin.reflect.KClass
 
 class EditorCell(
   private val editor: EditorEx,
+  private val manager: NotebookCellInlayManager,
   internal var intervalPointer: NotebookIntervalPointer,
   private val viewFactory: (EditorCell) -> EditorCellView,
 ) : Disposable, UserDataHolder by UserDataHolderBase() {
@@ -34,7 +36,7 @@ class EditorCell(
 
   val interval get() = intervalPointer.get() ?: error("Invalid interval")
 
-  var view: EditorCellView? = viewFactory(this)
+  var view: EditorCellView? = createView()
 
   var visible: Boolean = true
     set(value) {
@@ -43,13 +45,12 @@ class EditorCell(
 
       if (value) {
         view?.let {
-          Disposer.dispose(it)
-          view = null
+          disposeView(it)
         }
       }
       else {
         if (view == null) {
-          view = viewFactory(this).also { Disposer.register(this, it) }
+          view = createView()
           view?.selected = selected
           gutterAction?.let { view?.setGutterAction(it) }
           view?.updateExecutionStatus(executionCount, progressStatus, executionStartTime, executionEndTime)
@@ -57,6 +58,18 @@ class EditorCell(
       }
     }
 
+  private fun createView(): EditorCellView {
+    val view = viewFactory(this).also { Disposer.register(this, it) }
+    manager.fireCellViewCreated(view)
+    return view
+  }
+
+  private fun disposeView(it: EditorCellView) {
+    Disposer.dispose(it)
+    view = null
+    manager.fireCellViewRemoved(it)
+  }
+
   var selected: Boolean = false
     set(value) {
       if (field != value) {
@@ -76,7 +89,7 @@ class EditorCell(
   private var executionEndTime: ZonedDateTime? = null
 
   override fun dispose() {
-    view?.let { Disposer.dispose(it) }
+    view?.let { disposeView(it) }
   }
 
   fun update(force: Boolean = false) {

notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/ui/EditorCellInput.kt

@@ -92,7 +92,6 @@ class EditorCellInput(
 
   override fun doDispose() {
     folding.dispose()
-    runCellButton?.dispose()
     component.dispose()
   }
 

notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/ui/EditorCellRunGutterButton.kt

@@ -50,8 +50,6 @@ class EditorCellRunGutterButton(private val editor: EditorEx, private val cell:
 
   inner class DummyEmptyAction: AnAction(AllIcons.Empty) { override fun actionPerformed(e: AnActionEvent) { } }
 
-  fun dispose() = hideRunButton()
-
   companion object {
     private val runCellAction = ActionManager.getInstance().getAction("NotebookRunCellAction")
     private val interruptKernelAction = ActionManager.getInstance().getAction("JupyterInterruptKernelAction")

notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/ui/EditorCellViewEventListener.kt

@@ -0,0 +1,11 @@
+package org.jetbrains.plugins.notebooks.visualization.ui
+
+import java.util.*
+
+interface EditorCellViewEventListener : EventListener {
+  fun onEditorCellViewEvents(events: List<EditorCellViewEvent>)
+
+  sealed interface EditorCellViewEvent
+  data class CellViewCreated(val cell: EditorCellView) : EditorCellViewEvent
+  data class CellViewRemoved(val view: EditorCellView) : EditorCellViewEvent
+}