PY-80533 Jupyter: Refactor NotebookEditorCellView

Signed-off-by: Nikita.Ashihmin <nikita.ashihmin@jetbrains.com>

GitOrigin-RevId: b38b72071822c833020a2f02b6e82264b2906714

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

@@ -2,19 +2,25 @@ package com.intellij.notebooks.visualization
 
 import com.intellij.notebooks.visualization.ui.EditorCell
 import com.intellij.notebooks.visualization.ui.EditorCellViewComponent
-import com.intellij.openapi.editor.impl.EditorImpl
+import com.intellij.notebooks.visualization.ui.TextEditorCellInputFactory
 import com.intellij.openapi.extensions.ExtensionPointName
 
 /**
  * Marker interface for factories producing custom editors for cells
  */
 interface EditorCellInputFactory {
-  fun createComponent(editor: EditorImpl, cell: EditorCell): EditorCellViewComponent
+  fun createComponent(cell: EditorCell): EditorCellViewComponent
 
-  fun supports(editor: EditorImpl, cell: EditorCell): Boolean
+  fun supports(cell: EditorCell): Boolean
 
   companion object {
     @JvmField
     val EP_NAME: ExtensionPointName<EditorCellInputFactory> = ExtensionPointName.create<EditorCellInputFactory>("org.jetbrains.plugins.notebooks.inputFactory")
+
+
+    fun create(cell: EditorCell): EditorCellViewComponent {
+      val inputFactory = EP_NAME.extensionsIfPointIsRegistered.firstOrNull { it.supports(cell) } ?: TextEditorCellInputFactory()
+      return inputFactory.createComponent(cell)
+    }
   }
 }

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

@@ -179,7 +179,7 @@ class NotebookCellInlayManager private constructor(
     cell: EditorCell,
     ctx: UpdateContext,
   ) {
-    val view = EditorCellView(editor, cell, this)
+    val view = EditorCellView(cell)
     Disposer.register(cell, view)
     view.updateCellFolding(ctx)
     views[cell] = view
@@ -348,6 +348,11 @@ class NotebookCellInlayManager private constructor(
             change.subsequentPointers.forEach {
               addCell(it.pointer)
             }
+            //After insert we need fix ranges of previous cell
+            change.subsequentPointers.forEach {
+              val prevCell = getCellOrNull(it.interval.ordinal - 1)
+              prevCell?.checkAndRebuildInlays()
+            }
           }
           is NotebookIntervalPointersEvent.OnRemoved -> {
             change.subsequentPointers.reversed().forEach {
@@ -363,6 +368,10 @@ class NotebookCellInlayManager private constructor(
             secondCell.intervalPointer = first
             firstCell.update(ctx)
             secondCell.update(ctx)
+            firstCell.checkAndRebuildInlays()
+            secondCell.checkAndRebuildInlays()
+            getCellOrNull(firstCell.interval.ordinal - 1)?.checkAndRebuildInlays()
+            getCellOrNull(secondCell.interval.ordinal - 1)?.checkAndRebuildInlays()
           }
         }
       }
@@ -424,8 +433,4 @@ class NotebookCellInlayManager private constructor(
   fun getCell(pointer: NotebookIntervalPointer): EditorCell {
     return getCell(pointer.get()!!)
   }
-
-  internal fun getInputFactories(): Sequence<EditorCellInputFactory> {
-    return EditorCellInputFactory.EP_NAME.extensionList.asSequence()
-  }
 }

notebooks/visualization/src/com/intellij/notebooks/visualization/controllers/NotebookCellController.kt

@@ -1,4 +1,16 @@
 // Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package com.intellij.notebooks.visualization.controllers
 
-interface NotebookCellController
+import com.intellij.openapi.Disposable
+
+interface NotebookCellController : Disposable.Default {
+  /**
+   * As there are so many possible document editing operations that can destroy cell inlays by removing document range they attached to,
+   * the only option we have to preserve consistency is to check inlays validity
+   * and recreate them if needed.
+   * This logic is supposed to be as simple as check `isValid` and `offset` attributes of inlays
+   * so it should not introduce significant performance degradation.
+   */
+  fun checkAndRebuildInlays() {
+  }
+}

notebooks/visualization/src/com/intellij/notebooks/visualization/controllers/selfUpdate/SelfManagedCellController.kt

@@ -2,16 +2,8 @@
 package com.intellij.notebooks.visualization.controllers.selfUpdate
 
 import com.intellij.notebooks.visualization.controllers.NotebookCellController
-import com.intellij.openapi.Disposable
 
 /**
  * Controller which does not rely on external recreate and allow safe updater UI
- * All controllers in the future should implement this interface
  */
-interface SelfManagedCellController : NotebookCellController, Disposable.Default {
-
-  /**
-   * Update internal state of controller
-   */
-  fun selfUpdate()
-}
+interface SelfManagedCellController : NotebookCellController

notebooks/visualization/src/com/intellij/notebooks/visualization/controllers/selfUpdate/common/NotebookCellSelfHighlighterController.kt

@@ -31,7 +31,7 @@ abstract class NotebookCellSelfHighlighterController(
   open fun getTextAttribute(): TextAttributes? = null
   open fun customizeHighlighter(cellHighlighter: RangeHighlighterEx) {}
 
-  override fun selfUpdate() {
+  override fun checkAndRebuildInlays() {
     if (highlighter?.isValid == true &&
         highlighter?.startOffset == editorCell.interval.getCellStartOffset(editor) &&
         highlighter?.endOffset == editorCell.interval.getCellEndOffset(editor)

notebooks/visualization/src/com/intellij/notebooks/visualization/controllers/selfUpdate/common/NotebookCellSelfInlayController.kt

@@ -42,7 +42,7 @@ abstract class NotebookCellSelfInlayController(
 
   abstract fun createLineMarkerRender(createdHighlighter: RangeHighlighterEx): NotebookLineMarkerRenderer?
 
-  override fun selfUpdate() {
+  override fun checkAndRebuildInlays() {
     editor.updateManager.update { updater ->
       updater.addInlayOperation {
         editorCell.intervalOrNull ?: return@addInlayOperation
@@ -60,7 +60,7 @@ abstract class NotebookCellSelfInlayController(
   }
 
   open fun updateHighlight() {
-    highlighterController.selfUpdate()
+    highlighterController.checkAndRebuildInlays()
   }
 
   private fun createInlay(): Inlay<*> {

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

@@ -19,11 +19,9 @@ import javax.swing.BoxLayout
 import javax.swing.JComponent
 import javax.swing.JPanel
 
-class CustomFoldingEditorCellViewComponent(
-  internal val component: JComponent,
-  private val editor: EditorEx,
-  private val cell: EditorCell,
-) : EditorCellViewComponent() {
+class CustomFoldingEditorCellViewComponent(private val cell: EditorCell, internal val component: JComponent)
+  : EditorCellViewComponent() {
+  private val editor: EditorEx = cell.editor
 
   private var foldingRegion: CustomFoldRegion? = null
 
@@ -108,7 +106,6 @@ class CustomFoldingEditorCellViewComponent(
 
   override fun addInlayBelow(presentation: InlayPresentation) {
     val inlayComponent = object : JComponent() {
-
       init {
         enableEvents(MOUSE_EVENT_MASK or MOUSE_MOTION_EVENT_MASK)
       }

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

@@ -77,6 +77,10 @@ class EditorCell(
     source.set(interval.getContentText(editor))
   }
 
+  fun checkAndRebuildInlays() {
+    view?.checkAndRebuildInlays()
+  }
+
   fun onViewportChange() {
     view?.onViewportChanges()
   }

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

@@ -2,23 +2,16 @@ package com.intellij.notebooks.visualization.ui
 
 import com.intellij.notebooks.ui.visualization.NotebookUtil.notebookAppearance
 import com.intellij.notebooks.visualization.EditorCellInputFactory
-import com.intellij.notebooks.visualization.NotebookCellLines
 import com.intellij.notebooks.visualization.ui.cellsDnD.EditorCellDragAssistant
 import com.intellij.openapi.editor.Inlay
 import com.intellij.openapi.util.Disposer
 import com.intellij.openapi.util.registry.Registry
 import java.awt.Rectangle
 
-class EditorCellInput(
-  componentFactory: EditorCellInputFactory,
-  val cell: EditorCell,
-) : EditorCellViewComponent() {
+class EditorCellInput(val cell: EditorCell) : EditorCellViewComponent() {
   private val editor = cell.editor
-  val interval: NotebookCellLines.Interval
-    get() = cell.intervalPointer.get() ?: error("Invalid interval")
 
-
-  val component: EditorCellViewComponent = componentFactory.createComponent(editor, cell).also { add(it) }
+  val component: EditorCellViewComponent = EditorCellInputFactory.create(cell).also { add(it) }
 
   private val dragAssistant = when (Registry.`is`("jupyter.editor.dnd.cells")) {
     true -> EditorCellDragAssistant(editor, this, ::fold, ::unfold).also { Disposer.register(this, it) }
@@ -39,7 +32,7 @@ class EditorCellInput(
       return Pair(0, 0)
     }
 
-    val delimiterPanelSize = if (interval.ordinal == 0) {
+    val delimiterPanelSize = if (cell.interval.ordinal == 0) {
       editor.notebookAppearance.aboveFirstCellDelimiterHeight
     }
     else {
@@ -66,7 +59,7 @@ class EditorCellInput(
   }
 
   fun getBlockElementsInRange(): List<Inlay<*>> {
-    val linesRange = interval.lines
+    val linesRange = cell.interval.lines
     val startOffset = editor.document.getLineStartOffset(linesRange.first)
     val endOffset = editor.document.getLineEndOffset(linesRange.last)
     return editor.inlayModel.getBlockElementsInRange(startOffset, endOffset)

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

@@ -1,28 +1,23 @@
 package com.intellij.notebooks.visualization.ui
 
 import com.intellij.notebooks.ui.bind
-import com.intellij.notebooks.visualization.EditorCellInputFactory
-import com.intellij.notebooks.visualization.NotebookCellInlayManager
 import com.intellij.notebooks.visualization.NotebookCellLines
 import com.intellij.notebooks.visualization.UpdateContext
 import com.intellij.notebooks.visualization.controllers.selfUpdate.SelfManagedCellController
 import com.intellij.notebooks.visualization.controllers.selfUpdate.SelfManagedControllerFactory
 import com.intellij.notebooks.visualization.ui.cellsDnD.DropHighlightable
-import com.intellij.openapi.Disposable
 import com.intellij.openapi.application.runInEdt
 import com.intellij.openapi.editor.EditorKind
-import com.intellij.openapi.editor.impl.EditorImpl
 import com.intellij.openapi.util.registry.Registry
 import java.awt.Rectangle
 
 
-class EditorCellView(
-  val editor: EditorImpl,
-  val cell: EditorCell,
-  private val cellInlayManager: NotebookCellInlayManager,
-) : EditorCellViewComponent(), Disposable {
-  val input: EditorCellInput = createEditorCellInput()
+class EditorCellView(val cell: EditorCell) : EditorCellViewComponent() {
+  private val editor = cell.editor
 
+  val input: EditorCellInput = EditorCellInput(cell).also {
+    add(it)
+  }
   var outputs: EditorCellOutputsView? = null
     private set
 
@@ -40,63 +35,36 @@ class EditorCellView(
     SelfManagedControllerFactory.createControllers(this)
   }
 
-
-  // We are storing last lines range for highlighters to prevent highlighters unnecessary recreation on the same lines.
-  private var lastHighLightersLines: IntRange? = null
-
-
   init {
     cell.source.bind(this) {
       updateInput()
     }
     cell.isSelected.bind(this) { selected ->
-      updateSelected()
+      updateFolding()
     }
     cell.isHovered.bind(this) {
       updateHovered()
     }
 
-    updateSelfManaged()
     updateOutputs()
-  }
-
-  private fun updateSelected() {
-    updateFolding()
-    updateCellHighlight()
-  }
-
-  override fun dispose() {
-    super.dispose()
-  }
-
-  private fun createEditorCellInput(): EditorCellInput {
-    val inputFactory = getInputFactories().firstOrNull { it.supports(editor, cell) } ?: TextEditorCellInputFactory()
-    return EditorCellInput(inputFactory, cell).also {
-      add(it)
-    }
+    checkAndRebuildInlays()
   }
 
   fun update(updateContext: UpdateContext) {
     input.updateInput()
-    updateSelfManaged()
     updateOutputs()
     updateCellFolding(updateContext)
   }
 
-  private fun updateSelfManaged() {
-    controllers.forEach {
-      it.selfUpdate()
-    }
-  }
-
-  private fun updateInput() = runInEdt {
-    updateCellHighlight()
+  private fun updateInput() {
     input.updateInput()
-    checkAndRebuildInlays()
   }
 
   override fun doCheckAndRebuildInlays() {
-    updateSelfManaged()
+    controllers.forEach {
+      it.checkAndRebuildInlays()
+    }
+    cell.cellFrameManager?.updateCellFrameShow()
   }
 
   private fun updateOutputs() = runInEdt {
@@ -105,7 +73,6 @@ class EditorCellView(
         outputs = EditorCellOutputsView(editor, cell).also {
           add(it)
         }
-        updateCellHighlight()
         updateFolding()
       }
       else {
@@ -123,8 +90,6 @@ class EditorCellView(
   private fun hasOutputs() = cell.interval.type == NotebookCellLines.CellType.CODE
                              && (editor.editorKind != EditorKind.DIFF || Registry.`is`("jupyter.diff.viewer.output"))
 
-  private fun getInputFactories(): Sequence<EditorCellInputFactory> = cellInlayManager.getInputFactories()
-
   fun onViewportChanges() {
     input.onViewportChange()
     outputs?.onViewportChange()
@@ -134,18 +99,6 @@ class EditorCellView(
     updateFolding()
   }
 
-  private fun updateCellHighlight(force: Boolean = false) {
-    val interval = cell.interval
-
-    if (!force && interval.lines == lastHighLightersLines) {
-      return
-    }
-    lastHighLightersLines = IntRange(interval.lines.first, interval.lines.last)
-    updateSelfManaged()
-
-  }
-
-
   private fun updateFolding() {
     input.folding.visible = isHovered || isSelected
     input.folding.selected = isSelected

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

@@ -2,12 +2,12 @@ package com.intellij.notebooks.visualization.ui
 
 import com.intellij.codeInsight.hints.presentation.InlayPresentation
 import com.intellij.notebooks.visualization.UpdateContext
-import com.intellij.openapi.Disposable
+import com.intellij.notebooks.visualization.controllers.NotebookCellController
 import com.intellij.openapi.util.Disposer
 import java.awt.Rectangle
 import java.util.*
 
-abstract class EditorCellViewComponent : Disposable.Default {
+abstract class EditorCellViewComponent : NotebookCellController {
   protected var parent: EditorCellViewComponent? = null
 
   private val _children = mutableListOf<EditorCellViewComponent>()
@@ -52,14 +52,8 @@ abstract class EditorCellViewComponent : Disposable.Default {
     throw UnsupportedOperationException("Operation is not supported")
   }
 
-  /**
-   * As there are so many possible document editing operations that can destroy cell inlays by removing document range they attached to,
-   * the only option we have to preserve consistency is to check inlays validity
-   * and recreate them if needed.
-   * This logic is supposed to be as simple as check `isValid` and `offset` attributes of inlays
-   * so it should not introduce significant performance degradation.
-   */
-  fun checkAndRebuildInlays() {
+
+  final override fun checkAndRebuildInlays() {
     _children.forEach { it.checkAndRebuildInlays() }
     doCheckAndRebuildInlays()
   }

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

@@ -2,9 +2,8 @@
 package com.intellij.notebooks.visualization.ui
 
 import com.intellij.notebooks.visualization.EditorCellInputFactory
-import com.intellij.openapi.editor.impl.EditorImpl
 
 class TextEditorCellInputFactory : EditorCellInputFactory {
-  override fun createComponent(editor: EditorImpl, cell: EditorCell): TextEditorCellViewComponent = TextEditorCellViewComponent(cell)
-  override fun supports(editor: EditorImpl, cell: EditorCell): Boolean = true
+  override fun createComponent(cell: EditorCell): TextEditorCellViewComponent = TextEditorCellViewComponent(cell)
+  override fun supports(cell: EditorCell): Boolean = true
 }

notebooks/visualization/src/com/intellij/notebooks/visualization/ui/cell/frame/EditorCellFrameManager.kt

@@ -100,7 +100,7 @@ class EditorCellFrameManager(private val editorCell: EditorCell) : Disposable {
     return line2DDouble
   }
 
-  private fun updateCellFrameShow() {
+  fun updateCellFrameShow() {
     if (cellType == CellType.MARKDOWN) {
       updateCellFrameShowMarkdown()
     }

notebooks/visualization/src/com/intellij/notebooks/visualization/ui/cell/runCell/EditorCellRunGutterController.kt

@@ -49,7 +49,7 @@ class EditorCellRunGutterController(
     cell.gutterAction.set(null)
   }
 
-  override fun selfUpdate() {}
+  override fun checkAndRebuildInlays() {}
 
   private fun updateGutterAction() {
     //For markdown, it will set up in markdown component

notebooks/visualization/src/com/intellij/notebooks/visualization/ui/cell/toolbar/EditorCellActionsToolbarController.kt

@@ -65,7 +65,7 @@ internal class EditorCellActionsToolbarController(
     updateToolbarVisibility()
   }
 
-  override fun selfUpdate() {
+  override fun checkAndRebuildInlays() {
     val component = targetComponent ?: return
     updateToolbarPosition(component)
   }