PY-65441 Expand/Collapse jupyter code/markdown cells

Merge-request: IJ-MR-126097
Merged-by: Anton Efimchuk <Anton.Efimchuk@jetbrains.com>

GitOrigin-RevId: 04b7b1b9745fbe7da922cda2060416e3d457be1c

notebooks/notebook-ui/src/org/jetbrains/plugins/notebooks/ui/NotebookEditorUiUtil.kt

@@ -1,5 +1,6 @@
 package org.jetbrains.plugins.notebooks.ui
 
+import com.intellij.openapi.util.Key
 import javax.swing.JPanel
 import javax.swing.plaf.PanelUI
 
@@ -21,3 +22,5 @@ open class SteadyUIPanel(private val steadyUi: PanelUI) : JPanel() {
     setUI(steadyUi)
   }
 }
+
+val isFoldingEnabledKey = Key.create<Boolean>("jupyter.editor.folding.cells")

notebooks/notebook-ui/src/org/jetbrains/plugins/notebooks/ui/editor/DefaultNotebookEditorAppearance.kt

@@ -6,12 +6,15 @@ import com.intellij.openapi.editor.colors.EditorColors
 import com.intellij.openapi.editor.colors.EditorColorsScheme
 import com.intellij.openapi.editor.colors.TextAttributesKey
 import com.intellij.openapi.editor.impl.EditorImpl
+import com.intellij.ui.JBColor
 import com.intellij.ui.NewUiValue
 import com.intellij.util.ui.JBUI
 import org.jetbrains.plugins.notebooks.ui.editor.actions.command.mode.NotebookEditorMode
 import org.jetbrains.plugins.notebooks.ui.editor.actions.command.mode.currentMode
 import org.jetbrains.plugins.notebooks.ui.visualization.DefaultNotebookEditorAppearanceSizes
 import org.jetbrains.plugins.notebooks.ui.visualization.NotebookEditorAppearance
+import org.jetbrains.plugins.notebooks.ui.visualization.NotebookEditorAppearance.Companion.CELL_STRIPE_COLOR
+import org.jetbrains.plugins.notebooks.ui.visualization.NotebookEditorAppearance.Companion.CELL_UNDER_CURSOR_STRIPE_HOVER_COLOR
 import org.jetbrains.plugins.notebooks.ui.visualization.NotebookEditorAppearanceSizes
 import java.awt.Color
 
@@ -56,7 +59,6 @@ object DefaultNotebookEditorAppearance : NotebookEditorAppearance,
 
   val CELL_UNDER_CARET_COMMAND_MODE_STRIPE_COLOR = ColorKey.createColorKey("JUPYTER.CELL_UNDER_CARET_COMMAND_MODE_STRIPE_COLOR")
   val CELL_UNDER_CARET_EDITOR_MODE_STRIPE_COLOR = ColorKey.createColorKey("JUPYTER.CELL_UNDER_CARET_EDITOR_MODE_STRIPE_COLOR")
-  private val CELL_UNDER_CURSOR_STRIPE_HOVER_COLOR = ColorKey.createColorKey("JUPYTER.CELL_UNDER_CURSOR_STRIPE_HOVER_COLOR")
 
   // see org.jetbrains.plugins.notebooks.visualization.CaretBasedCellSelectionModelKt.getSelectionLines
   private fun isCellSelected(editor: Editor, lines: IntRange): Boolean {
@@ -93,6 +95,14 @@ object DefaultNotebookEditorAppearance : NotebookEditorAppearance,
     return null
   }
 
+  override fun getCellStripeHoverColor(editor: Editor): Color {
+    return editor.colorsScheme.getColor(CELL_UNDER_CURSOR_STRIPE_HOVER_COLOR) ?: JBColor.BLUE
+  }
+
+  override fun getCellStripeColor(editor: Editor): Color {
+    return editor.colorsScheme.getColor(CELL_STRIPE_COLOR) ?: JBColor.GRAY
+  }
+
   override fun getCellLeftLineWidth(): Int =
     when (currentMode()) {
       NotebookEditorMode.EDIT -> EDIT_MODE_CELL_LEFT_LINE_WIDTH

notebooks/notebook-ui/src/org/jetbrains/plugins/notebooks/ui/editor/NewUINotebookDiffEditorAppearance.kt

@@ -1,7 +1,9 @@
 package org.jetbrains.plugins.notebooks.ui.editor
 
+import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.editor.colors.ColorKey
 import com.intellij.openapi.editor.colors.EditorColorsScheme
+import com.intellij.ui.JBColor
 import org.jetbrains.plugins.notebooks.ui.visualization.NotebookEditorAppearance
 import org.jetbrains.plugins.notebooks.ui.visualization.NotebookEditorAppearanceSizes
 import java.awt.Color
@@ -18,6 +20,12 @@ object NewUINotebookDiffEditorAppearance: NotebookEditorAppearance,
   override fun shouldShowCellLineNumbers(): Boolean = false
   override fun shouldShowExecutionCounts(): Boolean = false  // not needed for DIFF -> execution does not reach it
   override fun shouldShowOutExecutionCounts(): Boolean = false
+  override fun getCellStripeHoverColor(editor: Editor): Color {
+    return editor.colorsScheme.getColor(NotebookEditorAppearance.CELL_UNDER_CURSOR_STRIPE_HOVER_COLOR) ?: JBColor.BLUE
+  }
+  override fun getCellStripeColor(editor: Editor): Color {
+    return editor.colorsScheme.getColor(NotebookEditorAppearance.CELL_STRIPE_COLOR) ?: JBColor.GRAY
+  }
 }
 
 

notebooks/notebook-ui/src/org/jetbrains/plugins/notebooks/ui/visualization/NotebookEditorAppearance.kt

@@ -1,6 +1,7 @@
 // Copyright 2000-2022 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package org.jetbrains.plugins.notebooks.ui.visualization
 
+import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.editor.colors.ColorKey
 import com.intellij.openapi.editor.colors.EditorColorsScheme
 import com.intellij.openapi.editor.impl.EditorImpl
@@ -18,6 +19,8 @@ interface NotebookEditorAppearance: NotebookEditorAppearanceColors, NotebookEdit
     val NOTEBOOK_APPEARANCE_KEY = Key.create<NotebookEditorAppearance>(NotebookEditorAppearance::class.java.name)
     val CODE_CELL_BACKGROUND = ColorKey.createColorKey("JUPYTER.CODE_CELL_BACKGROUND")
     internal val CODE_CELL_BACKGROUND_NEW_UI = ColorKey.createColorKey("JUPYTER.CODE_CELL_BACKGROUND_NEW_UI")
+    internal val CELL_UNDER_CURSOR_STRIPE_HOVER_COLOR = ColorKey.createColorKey("JUPYTER.CELL_UNDER_CURSOR_STRIPE_HOVER_COLOR")
+    internal val CELL_STRIPE_COLOR = ColorKey.createColorKey("JUPYTER.CELL_STRIPE_COLOR")
   }
 
   fun getCaretRowColor(scheme: EditorColorsScheme): Color?
@@ -85,6 +88,9 @@ interface NotebookEditorAppearanceColors {
    */
   fun getCellStripeColor(editor: EditorImpl, lines: IntRange): Color? = null
   fun getCellStripeHoverColor(editor: EditorImpl, lines: IntRange): Color? = null
+
+  fun getCellStripeColor(editor: Editor): Color
+  fun getCellStripeHoverColor(editor: Editor): Color
 }
 
 interface NotebookEditorAppearanceFlags {

notebooks/notebook-ui/src/org/jetbrains/plugins/notebooks/ui/visualization/NotebookLineMarkerRenderer.kt

@@ -12,6 +12,7 @@ import com.intellij.openapi.editor.impl.EditorImpl
 import com.intellij.openapi.editor.impl.view.FontLayoutService
 import com.intellij.openapi.editor.markup.LineMarkerRendererEx
 import com.intellij.openapi.editor.markup.RangeHighlighter
+import org.jetbrains.plugins.notebooks.ui.isFoldingEnabledKey
 import java.awt.Graphics
 import java.awt.Point
 import java.awt.Rectangle
@@ -23,7 +24,7 @@ abstract class NotebookLineMarkerRenderer(private val inlayId: Long? = null) : L
 
   override fun getPosition(): LineMarkerRendererEx.Position = LineMarkerRendererEx.Position.CUSTOM
 
-  protected fun getInlayBounds(editor: EditorEx, linesRange: IntRange) : Rectangle? {
+  protected fun getInlayBounds(editor: EditorEx, linesRange: IntRange): Rectangle? {
     val startOffset = editor.document.getLineStartOffset(linesRange.first)
     val endOffset = editor.document.getLineEndOffset(linesRange.last)
     val inlays = editor.inlayModel.getBlockElementsInRange(startOffset, endOffset)
@@ -56,10 +57,12 @@ class NotebookBelowCellCellGutterLineMarkerRenderer(private val highlighter: Ran
 
 class MarkdownCellGutterLineMarkerRenderer(private val highlighter: RangeHighlighter, inlayId: Long) : NotebookLineMarkerRenderer(inlayId) {
   override fun paint(editor: Editor, g: Graphics, r: Rectangle) {
-    editor as EditorImpl
-    val lines = IntRange(editor.document.getLineNumber(highlighter.startOffset), editor.document.getLineNumber(highlighter.endOffset))
-    val inlayBounds = getInlayBounds(editor, lines) ?: return
-    paintCellGutter(inlayBounds, lines, editor, g, r)
+    if (editor.getUserData(isFoldingEnabledKey) != true) {
+      editor as EditorImpl
+      val lines = IntRange(editor.document.getLineNumber(highlighter.startOffset), editor.document.getLineNumber(highlighter.endOffset))
+      val inlayBounds = getInlayBounds(editor, lines) ?: return
+      paintCellGutter(inlayBounds, lines, editor, g, r)
+    }
   }
 }
 
@@ -127,9 +130,11 @@ class NotebookTextCellBackgroundLineMarkerRenderer(private val highlighter: Rang
     val height = editor.offsetToXY(editor.document.getLineEndOffset(lines.last)).y + editor.lineHeight - top
 
     paintCaretRow(editor, g, lines)
-    val appearance = editor.notebookAppearance
-    appearance.getCellStripeColor(editor, lines)?.let {
-      paintCellStripe(appearance, g, r, it, top, height)
+    if (editor.getUserData(isFoldingEnabledKey) != true) {
+      val appearance = editor.notebookAppearance
+      appearance.getCellStripeColor(editor, lines)?.let {
+        paintCellStripe(appearance, g, r, it, top, height)
+      }
     }
   }
 }

notebooks/notebook-ui/src/org/jetbrains/plugins/notebooks/ui/visualization/NotebookUtil.kt

@@ -7,6 +7,7 @@ import com.intellij.openapi.editor.LineNumberConverter
 import com.intellij.openapi.editor.colors.EditorColors
 import com.intellij.openapi.editor.ex.EditorEx
 import com.intellij.openapi.editor.impl.EditorImpl
+import org.jetbrains.plugins.notebooks.ui.isFoldingEnabledKey
 import org.jetbrains.plugins.notebooks.ui.visualization.NotebookEditorAppearance.Companion.NOTEBOOK_APPEARANCE_KEY
 import java.awt.Color
 import java.awt.Graphics
@@ -42,13 +43,15 @@ inline fun paintNotebookCellBackgroundGutter(
   }
 
   actionBetweenBackgroundAndStripe()
-  if (editor.editorKind == EditorKind.DIFF) return
-  if (stripe != null) {
-    paintCellStripe(appearance, g, r, stripe, top, height)
-  }
-  if (stripeHover != null) {
-    g.color = stripeHover
-    g.fillRect(r.width - appearance.getLeftBorderWidth(), top, appearance.getCellLeftLineHoverWidth(), height)
+  if (editor.getUserData(isFoldingEnabledKey) != true) {
+    if (editor.editorKind == EditorKind.DIFF) return
+    if (stripe != null) {
+      paintCellStripe(appearance, g, r, stripe, top, height)
+    }
+    if (stripeHover != null) {
+      g.color = stripeHover
+      g.fillRect(r.width - appearance.getLeftBorderWidth(), top, appearance.getCellLeftLineHoverWidth(), height)
+    }
   }
 }
 

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

@@ -3,7 +3,8 @@ package org.jetbrains.plugins.notebooks.visualization
 import com.intellij.openapi.diagnostic.thisLogger
 import com.intellij.openapi.editor.Document
 import com.intellij.openapi.editor.event.DocumentEvent
-import com.intellij.openapi.editor.event.DocumentListener
+import com.intellij.openapi.editor.ex.PrioritizedDocumentListener
+import com.intellij.openapi.editor.impl.EditorDocumentPriorities
 import com.intellij.openapi.util.TextRange
 import com.intellij.util.EventDispatcher
 import com.intellij.util.concurrency.ThreadingAssertions
@@ -35,11 +36,11 @@ class NonIncrementalCellLines private constructor(private val document: Document
     return intervals.listIterator(ordinal)
   }
 
-  private fun notifyChanged(oldCells: List<NotebookCellLines.Interval>,
-                            oldAffectedCells: List<NotebookCellLines.Interval>,
-                            newCells: List<NotebookCellLines.Interval>,
-                            newAffectedCells: List<NotebookCellLines.Interval>,
-                            documentEvent: DocumentEvent) {
+  private fun createEvent(oldCells: List<NotebookCellLines.Interval>,
+                                     newCells: List<NotebookCellLines.Interval>,
+                                     oldAffectedCells: List<NotebookCellLines.Interval>,
+                                     newAffectedCells: List<NotebookCellLines.Interval>,
+                                     documentEvent: DocumentEvent): NotebookCellLinesEvent {
     val (trimmedOldCells, trimmedNewCells) =
       if (oldCells == newCells) {
         Pair(emptyList(), emptyList())
@@ -70,15 +71,22 @@ class NonIncrementalCellLines private constructor(private val document: Document
       newAffectedIntervals = newAffectedCells,
       modificationStamp = modificationStamp,
     )
+    return event
+  }
 
+  private fun notify(event: NotebookCellLinesEvent) {
     catchThrowableAndLog {
       intervalListeners.multicaster.documentChanged(event)
     }
   }
 
-  private fun createDocumentListener() = object : DocumentListener {
+  private fun createDocumentListener() = object : PrioritizedDocumentListener {
     private var oldAffectedCells: List<NotebookCellLines.Interval> = emptyList()
 
+    private val postponedEvents = mutableListOf<NotebookCellLinesEvent>()
+
+    override fun getPriority(): Int = EditorDocumentPriorities.INLAY_MODEL + 1
+
     override fun beforeDocumentChange(event: DocumentEvent) {
       oldAffectedCells = getAffectedCells(intervals, document, TextRange(event.offset, event.offset + event.oldLength))
 
@@ -99,7 +107,21 @@ class NonIncrementalCellLines private constructor(private val document: Document
       intervals = intervalsGenerator.makeIntervals(document)
 
       val newAffectedCells = getAffectedCells(intervals, document, TextRange(event.offset, event.offset + event.newLength))
-      notifyChanged(oldIntervals, oldAffectedCells, intervals, newAffectedCells, event)
+      val newEvent = createEvent(oldIntervals, intervals, oldAffectedCells, newAffectedCells, event)
+      if (event.document.isInBulkUpdate) {
+        postponedEvents.add(newEvent)
+      } else {
+        notify(newEvent)
+      }
+    }
+
+    override fun bulkUpdateFinished(document: Document) {
+      if (postponedEvents.isNotEmpty()) {
+        postponedEvents.forEach {
+          notify(it)
+        }
+      }
+      postponedEvents.clear()
     }
   }
 

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

@@ -34,6 +34,11 @@ interface NotebookCellInlayController {
     }
   }
 
+  /**
+   * Marker interface for factories producing custom editors for cells
+   */
+  interface InputFactory
+
   val inlay: Inlay<*>
 
   val factory: Factory

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

@@ -1,27 +1,19 @@
 package org.jetbrains.plugins.notebooks.visualization
 
-import com.intellij.ide.DataManager
 import com.intellij.ide.ui.LafManagerListener
-import com.intellij.openapi.Disposable
-import com.intellij.openapi.actionSystem.DataProvider
-import com.intellij.openapi.actionSystem.PlatformCoreDataKeys
-import com.intellij.openapi.actionSystem.PlatformDataKeys
 import com.intellij.openapi.application.ApplicationManager
-import com.intellij.openapi.diagnostic.logger
-import com.intellij.openapi.diagnostic.thisLogger
-import com.intellij.openapi.editor.Document
+import com.intellij.openapi.application.runInEdt
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.editor.EditorKind
 import com.intellij.openapi.editor.Inlay
 import com.intellij.openapi.editor.colors.EditorColorsListener
 import com.intellij.openapi.editor.colors.EditorColorsManager
-import com.intellij.openapi.editor.event.DocumentEvent
-import com.intellij.openapi.editor.event.DocumentListener
+import com.intellij.openapi.editor.ex.FoldingListener
 import com.intellij.openapi.editor.ex.RangeHighlighterEx
 import com.intellij.openapi.editor.impl.EditorImpl
 import com.intellij.openapi.editor.markup.*
-import com.intellij.openapi.util.Disposer
 import com.intellij.openapi.util.Key
+import com.intellij.openapi.util.registry.Registry
 import com.intellij.util.Processor
 import com.intellij.util.SmartList
 import com.intellij.util.asSafely
@@ -31,14 +23,14 @@ import com.intellij.util.ui.update.MergingUpdateQueue
 import com.intellij.util.ui.update.Update
 import org.jetbrains.annotations.TestOnly
 import org.jetbrains.plugins.notebooks.ui.visualization.notebookAppearance
-import org.jetbrains.plugins.notebooks.visualization.UpdateInlaysTask.Companion.CELL_MARKER
-import org.jetbrains.plugins.notebooks.visualization.outputs.NotebookOutputInlayController
+import org.jetbrains.plugins.notebooks.visualization.ui.EditorCell
+import org.jetbrains.plugins.notebooks.ui.isFoldingEnabledKey
 import java.awt.Graphics
 import javax.swing.JComponent
 import kotlin.math.max
 import kotlin.math.min
 
-class NotebookCellInlayManager private constructor(val editor: EditorImpl) {
+class NotebookCellInlayManager private constructor(val editor: EditorImpl) : NotebookIntervalPointerFactory.ChangeListener {
   private val inlays: MutableMap<Inlay<*>, NotebookCellInlayController> = HashMap()
   private val notebookCellLines = NotebookCellLines.get(editor)
   private val viewportQueue = MergingUpdateQueue("NotebookCellInlayManager Viewport Update", 100, true, null, editor.disposable, null, true)
@@ -47,13 +39,17 @@ class NotebookCellInlayManager private constructor(val editor: EditorImpl) {
   private val updateQueue = MergingUpdateQueue("NotebookCellInlayManager Interval Update", 20, true, null, editor.disposable, null, true)
   private var initialized = false
 
+  private var _cells = mutableListOf<EditorCell>()
+
+  val cells: List<EditorCell> get() = _cells.toList()
+
   /**
    * 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
 
   fun inlaysForInterval(interval: NotebookCellLines.Interval): Iterable<NotebookCellInlayController> =
-    getMatchingInlaysForLines(interval.lines)
+    _cells[interval.ordinal].controllers
 
   /** It's public, but think twice before using it. Called many times in a row, it can freeze UI. Consider using [update] instead. */
   fun updateImmediately(lines: IntRange) {
@@ -84,6 +80,7 @@ class NotebookCellInlayManager private constructor(val editor: EditorImpl) {
 
   private fun addViewportChangeListener() {
     editor.scrollPane.viewport.addChangeListener {
+      scheduleUpdatePositions()
       viewportQueue.queue(object : Update("Viewport change") {
         override fun run() {
           if (editor.isDisposed) return
@@ -106,8 +103,6 @@ class NotebookCellInlayManager private constructor(val editor: EditorImpl) {
 
     handleRefreshedDocument()
 
-    addDocumentListener()
-
     val connection = ApplicationManager.getApplication().messageBus.connect(editor.disposable)
     connection.subscribe(EditorColorsManager.TOPIC, EditorColorsListener {
       updateAll()
@@ -120,9 +115,21 @@ class NotebookCellInlayManager private constructor(val editor: EditorImpl) {
 
     addViewportChangeListener()
 
+    editor.foldingModel.addListener(object : FoldingListener {
+      override fun onFoldProcessingEnd() {
+        scheduleUpdatePositions()
+      }
+    }, editor.disposable)
+
     initialized = true
   }
 
+  private fun scheduleUpdatePositions() {
+    runInEdt {
+      _cells.forEach { cell -> cell.updatePositions() }
+    }
+  }
+
   private fun refreshHighlightersLookAndFeel() {
     for (highlighter in editor.markupModel.allHighlighters) {
       if (highlighter.customRenderer === NotebookCellHighlighterRenderer) {
@@ -133,78 +140,22 @@ class NotebookCellInlayManager private constructor(val editor: EditorImpl) {
 
   private fun handleRefreshedDocument() {
     ThreadingAssertions.softAssertReadAccess()
-    val factories = NotebookCellInlayController.Factory.EP_NAME.extensionList
-    for (interval in notebookCellLines.intervals) {
-      for (factory in factories) {
-        val controller = failSafeCompute(factory, editor, emptyList(), notebookCellLines.intervals.listIterator(interval.ordinal))
-        if (controller != null) {
-          rememberController(controller, interval)
-        }
-      }
+    _cells.forEach {
+      it.dispose()
     }
+    val pointerFactory = NotebookIntervalPointerFactory.get(editor)
+    _cells = notebookCellLines.intervals.map { interval ->
+      createCell(pointerFactory.create(interval))
+    }.toMutableList()
     addHighlighters(notebookCellLines.intervals)
     inlaysChanged()
   }
 
-  private fun addDocumentListener() {
-    val documentListener = object : DocumentListener {
-      private var matchingCellsBeforeChange: List<NotebookCellLines.Interval> = emptyList()
-      private var isBulkModeEnabled = false
-
-      private fun interestingLogicalLines(document: Document, startOffset: Int, length: Int): IntRange {
-        // Adding one additional line is needed to handle deletions at the end of the document.
-        val end =
-          if (startOffset + length <= document.textLength) document.getLineNumber(startOffset + length)
-          else document.lineCount + 1
-        return document.getLineNumber(startOffset)..end
-      }
-
-      override fun bulkUpdateStarting(document: Document) {
-        isBulkModeEnabled = true
-        matchingCellsBeforeChange = notebookCellLines.getMatchingCells(0 until document.lineCount)
-      }
-
-      override fun beforeDocumentChange(event: DocumentEvent) {
-        if (isBulkModeEnabled) return
-        val document = event.document
-        val logicalLines = interestingLogicalLines(document, event.offset, event.oldLength)
-
-        matchingCellsBeforeChange = notebookCellLines.getMatchingCells(logicalLines)
-      }
-
-      override fun documentChanged(event: DocumentEvent) {
-        if (isBulkModeEnabled) return
-        if (event.oldLength == 0 && event.newFragment.contains(CELL_MARKER)) {
-          refreshInlays()
-        }
-
-        val logicalLines = interestingLogicalLines(event.document, event.offset, event.newLength)
-        ensureInlaysAndHighlightersExist(matchingCellsBeforeChange, logicalLines)
-      }
-
-      override fun bulkUpdateFinished(document: Document) {
-        isBulkModeEnabled = false
-
-        // bulk mode is over, now we could access inlays, let's update them all
-        refreshInlays()
-        ensureInlaysAndHighlightersExist(matchingCellsBeforeChange, 0 until document.lineCount)
-      }
-    }
-
-    editor.document.addDocumentListener(documentListener, editor.disposable)
-  }
-
-  /**
-   * Hack. When we are adding cell in notebook, previous cell changes their range and we need to update it.
-   */
-  private fun refreshInlays() {
-    val outputInlays = inlays.values.filterIsInstance<NotebookOutputInlayController>()
-    for (outputInlay in outputInlays) {
-      val oldInlay = outputInlay.checkAndUpdateInlayPosition() ?: continue
-      inlays.remove(oldInlay)
-      inlays[outputInlay.inlay] = outputInlay
-    }
-  }
+  private fun createCell(interval: NotebookIntervalPointer) = EditorCell(
+    editor,
+    notebookCellLines,
+    interval
+  )
 
   private fun ensureInlaysAndHighlightersExist(matchingCellsBeforeChange: List<NotebookCellLines.Interval>, logicalLines: IntRange) {
     val interestingRange =
@@ -238,86 +189,15 @@ class NotebookCellInlayManager private constructor(val editor: EditorImpl) {
     }
     addHighlighters(intervalsToAddHighlightersFor.values)
 
-    val allMatchingInlays: MutableList<Pair<Int, NotebookCellInlayController>> = getMatchingInlaysForLines(fullInterestingRange)
-      .mapTo(mutableListOf()) {
-        editor.document.getLineNumber(it.inlay.offset) to it
-      }
-    val allFactories = NotebookCellInlayController.Factory.EP_NAME.extensionList
-
     for (interval in matchingIntervals) {
-      val seenControllersByFactory: Map<NotebookCellInlayController.Factory, MutableList<NotebookCellInlayController>> =
-        allFactories.associateWith { SmartList() }
-      allMatchingInlays.removeIf { (inlayLine, controller) ->
-        if (inlayLine in interval.lines) {
-          seenControllersByFactory[controller.factory]?.add(controller)
-          true
-        }
-        else false
-      }
-      for ((factory, controllers) in seenControllersByFactory) {
-        val actualController = if (!editor.isDisposed) {
-          failSafeCompute(factory, editor, controllers, notebookCellLines.intervals.listIterator(interval.ordinal))
-        }
-        else {
-          null
-        }
-        if (actualController != null) {
-          rememberController(actualController, interval)
-        }
-        for (oldController in controllers) {
-          if (oldController != actualController) {
-            Disposer.dispose(oldController.inlay, false)
-          }
-        }
-      }
+      _cells[interval.ordinal].update()
     }
 
     NotebookGutterLineMarkerManager().putHighlighters(editor)
 
-    for ((_, controller) in allMatchingInlays) {
-      Disposer.dispose(controller.inlay, false)
-    }
     inlaysChanged()
   }
 
-  private data class NotebookCellDataProvider(
-    val editor: EditorImpl,
-    val component: JComponent,
-    val interval: NotebookCellLines.Interval,
-  ) : DataProvider {
-    override fun getData(key: String): Any? =
-      when (key) {
-        NOTEBOOK_CELL_LINES_INTERVAL_DATA_KEY.name -> interval
-        PlatformCoreDataKeys.CONTEXT_COMPONENT.name -> component
-        PlatformDataKeys.EDITOR.name -> editor
-        else -> null
-      }
-  }
-
-  private fun rememberController(controller: NotebookCellInlayController, interval: NotebookCellLines.Interval) {
-    val inlay = controller.inlay
-    inlay.renderer.asSafely<JComponent>()?.let { component ->
-      val oldProvider = DataManager.getDataProvider(component)
-      if (oldProvider != null && oldProvider !is NotebookCellDataProvider) {
-        LOG.error("Overwriting an existing CLIENT_PROPERTY_DATA_PROVIDER. Old provider: $oldProvider")
-      }
-      DataManager.removeDataProvider(component)
-      DataManager.registerDataProvider(component, NotebookCellDataProvider(editor, component, interval))
-    }
-    if (inlays.put(inlay, controller) !== controller) {
-      val disposable = Disposable {
-        inlay.renderer.asSafely<JComponent>()?.let { DataManager.removeDataProvider(it) }
-        inlays.remove(inlay)
-      }
-      if (Disposer.isDisposed(inlay)) {
-        disposable.dispose()
-      }
-      else {
-        Disposer.register(inlay, disposable)
-      }
-    }
-  }
-
   private fun getMatchingHighlightersForLines(lines: IntRange): List<RangeHighlighterEx> =
     mutableListOf<RangeHighlighterEx>()
       .also { list ->
@@ -331,19 +211,9 @@ class NotebookCellInlayManager private constructor(val editor: EditorImpl) {
         })
       }
 
-  private fun getMatchingInlaysForLines(lines: IntRange): List<NotebookCellInlayController> =
-    getMatchingInlaysForOffsets(
-      editor.document.getLineStartOffset(saturateLine(lines.first)),
-      editor.document.getLineEndOffset(saturateLine(lines.last)))
-
   private fun saturateLine(line: Int): Int =
     line.coerceAtMost(editor.document.lineCount - 1).coerceAtLeast(0)
 
-  private fun getMatchingInlaysForOffsets(startOffset: Int, endOffset: Int): List<NotebookCellInlayController> =
-    editor.inlayModel
-      .getBlockElementsInRange(startOffset, endOffset)
-      .mapNotNull(inlays::get)
-
   private val NotebookCellLines.Interval.shouldHaveHighlighter: Boolean
     get() = type == NotebookCellLines.CellType.CODE
 
@@ -379,19 +249,6 @@ class NotebookCellInlayManager private constructor(val editor: EditorImpl) {
       }
     }
 
-  private fun failSafeCompute(factory: NotebookCellInlayController.Factory,
-                              editor: EditorImpl,
-                              controllers: Collection<NotebookCellInlayController>,
-                              intervalIterator: ListIterator<NotebookCellLines.Interval>): NotebookCellInlayController? {
-    try {
-      return factory.compute(editor, controllers, intervalIterator)
-    }
-    catch (t: Throwable) {
-      thisLogger().error("${factory.javaClass.name} shouldn't throw exceptions at NotebookCellInlayController.Factory.compute(...)", t)
-      return null
-    }
-  }
-
   @TestOnly
   fun getInlays(): MutableMap<Inlay<*>, NotebookCellInlayController> = inlays
 
@@ -401,11 +258,12 @@ class NotebookCellInlayManager private constructor(val editor: EditorImpl) {
   }
 
   companion object {
-    private val LOG = logger<NotebookCellInlayManager>()
-
     @JvmStatic
     fun install(editor: EditorImpl) {
-      NotebookCellInlayManager(editor).initialize()
+      val notebookCellInlayManager = NotebookCellInlayManager(editor)
+      editor.putUserData(isFoldingEnabledKey, Registry.`is`("jupyter.editor.folding.cells"))
+      NotebookIntervalPointerFactory.get(editor).changeListeners.addListener(notebookCellInlayManager, editor.disposable)
+      notebookCellInlayManager.initialize()
     }
 
     @JvmStatic
@@ -413,6 +271,44 @@ class NotebookCellInlayManager private constructor(val editor: EditorImpl) {
 
     private val key = Key.create<NotebookCellInlayManager>(NotebookCellInlayManager::class.java.name)
   }
+
+  override fun onUpdated(event: NotebookIntervalPointersEvent) {
+    var start = Int.MAX_VALUE
+    var end = Int.MIN_VALUE
+    for (change in event.changes) {
+      when (change) {
+        is NotebookIntervalPointersEvent.OnEdited -> {
+          start = minOf(start, change.intervalBefore.lines.first, change.intervalAfter.lines.first)
+          end = maxOf(end, change.intervalBefore.lines.last, change.intervalAfter.lines.last)
+        }
+        is NotebookIntervalPointersEvent.OnInserted -> {
+          change.subsequentPointers.forEach {
+            _cells.add(it.interval.ordinal, createCell(it.pointer))
+          }
+          start = minOf(start, change.subsequentPointers.first().interval.lines.first)
+          end = maxOf(end, change.subsequentPointers.last().interval.lines.last)
+          scheduleUpdatePositions()
+        }
+        is NotebookIntervalPointersEvent.OnRemoved -> {
+          change.subsequentPointers.reversed().forEach {
+            val removed = _cells.removeAt(it.interval.ordinal)
+            removed.dispose()
+          }
+          start = minOf(start, change.subsequentPointers.first().interval.lines.first)
+          end = maxOf(end, change.subsequentPointers.last().interval.lines.last)
+          scheduleUpdatePositions()
+        }
+        is NotebookIntervalPointersEvent.OnSwapped -> {
+          val first = _cells[change.firstOrdinal].intervalPointer
+          _cells[change.firstOrdinal].intervalPointer = _cells[change.secondOrdinal].intervalPointer
+          _cells[change.secondOrdinal].intervalPointer = first
+          start = minOf(start, change.first.interval.lines.first)
+          end = maxOf(end, change.second.interval.lines.last)
+        }
+      }
+    }
+    updateConsequentInlays(start..end)
+  }
 }
 
 /**
@@ -473,8 +369,4 @@ private class UpdateInlaysTask(private val manager: NotebookCellInlayManager,
     pointerSet.addAll(update.pointerSet)
     return true
   }
-
-  companion object {
-    const val CELL_MARKER = "#%%"
-  }
-}
+                               }

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

@@ -3,8 +3,6 @@ package org.jetbrains.plugins.notebooks.visualization
 import com.intellij.openapi.editor.*
 import com.intellij.openapi.editor.event.CaretEvent
 import com.intellij.openapi.editor.event.CaretListener
-import com.intellij.openapi.editor.event.DocumentEvent
-import com.intellij.openapi.editor.event.DocumentListener
 import com.intellij.openapi.editor.ex.EditorEx
 import com.intellij.openapi.editor.ex.RangeHighlighterEx
 import com.intellij.openapi.editor.impl.EditorImpl
@@ -21,10 +19,11 @@ import java.awt.Rectangle
 class NotebookGutterLineMarkerManager {
 
   fun attachHighlighters(editor: EditorEx) {
-    editor.addEditorDocumentListener(object : DocumentListener {
-      override fun documentChanged(event: DocumentEvent) = putHighlighters(editor)
-      override fun bulkUpdateFinished(document: Document) = putHighlighters(editor)
-    })
+    NotebookIntervalPointerFactory.get(editor).changeListeners.addListener(object: NotebookIntervalPointerFactory.ChangeListener  {
+      override fun onUpdated(event: NotebookIntervalPointersEvent) {
+        putHighlighters(editor)
+      }
+    }, (editor as EditorImpl).disposable)
 
     editor.caretModel.addCaretListener(object : CaretListener {
       override fun caretPositionChanged(event: CaretEvent) {

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

@@ -42,13 +42,6 @@ private class NotebookIntervalPointerImpl(@Volatile var interval: NotebookCellLi
 
 private typealias NotebookIntervalPointersEventChanges = ArrayList<Change>
 
-
-private sealed interface ChangesContext
-
-private data class DocumentChangedContext(var redoContext: RedoContext? = null) : ChangesContext
-private data class UndoContext(val changes: List<Change>) : ChangesContext
-private data class RedoContext(val changes: List<Change>) : ChangesContext
-
 /**
  * One unique NotebookIntervalPointer exists for each current interval. You can use NotebookIntervalPointer as map key.
  * [NotebookIntervalPointerFactoryImpl] automatically supports undo/redo for [documentChanged] and [modifyPointers] calls.
@@ -62,7 +55,7 @@ class NotebookIntervalPointerFactoryImpl(private val notebookCellLines: Notebook
                                          undoManager: UndoManager?,
                                          private val project: Project) : NotebookIntervalPointerFactory, NotebookCellLines.IntervalListener {
   private val pointers = ArrayList<NotebookIntervalPointerImpl>()
-  private var changesContext: ChangesContext? = null
+  private var postponedEvent: NotebookIntervalPointersEvent? = null
   override val changeListeners: EventDispatcher<NotebookIntervalPointerFactory.ChangeListener> =
     EventDispatcher.create(NotebookIntervalPointerFactory.ChangeListener::class.java)
 
@@ -109,77 +102,44 @@ class NotebookIntervalPointerFactoryImpl(private val notebookCellLines: Notebook
   override fun documentChanged(event: NotebookCellLinesEvent) {
     ThreadingAssertions.assertWriteAccess()
     try {
-      val pointersEvent = when (val context = changesContext) {
-        is DocumentChangedContext -> documentChangedByAction(event, context)
-        is UndoContext -> documentChangedByUndo(event, context)
-        is RedoContext -> documentChangedByRedo(event, context)
-        null -> documentChangedByAction(event, null) // changesContext is null if undo manager is unavailable
+      if (validUndoManager?.isUndoOrRedoInProgress != true) {
+        documentChangedByAction(event)
+      } else {
+        val e = postponedEvent
+        if (e != null) {
+          onUpdated(e)
+        }
       }
-      onUpdated(pointersEvent)
     }
     catch (ex: Exception) {
       thisLogger().error(ex)
       // DS-3893 consume exception and log it, actions changing document should work as usual
-    }
-    finally {
-      changesContext = null
+    } finally {
+      postponedEvent = null
     }
   }
 
-  override fun beforeDocumentChange(event: NotebookCellLinesEventBeforeChange) {
-    ThreadingAssertions.assertWriteAccess()
-    val undoManager = validUndoManager
-    if (undoManager == null || undoManager.isUndoOrRedoInProgress) return
-    val context = DocumentChangedContext()
-    try {
-      undoManager.undoableActionPerformed(object : BasicUndoableAction() {
-        override fun undo() {}
-
-        override fun redo() {
-          changesContext = context.redoContext
-        }
-      })
-      changesContext = context
-    }
-    catch (ex: Exception) {
-      thisLogger().error(ex)
-      // DS-3893 consume exception, don't prevent document updating
-    }
-  }
-
-  private fun documentChangedByAction(event: NotebookCellLinesEvent,
-                                      documentChangedContext: DocumentChangedContext?): NotebookIntervalPointersEvent {
-    val eventChanges = NotebookIntervalPointersEventChanges()
-
-    updateChangedIntervals(event, eventChanges)
-    updateShiftedIntervals(event)
+  private fun documentChangedByAction(event: NotebookCellLinesEvent) {
+    val eventChanges = updateChangedIntervals(event)
+    val shiftChanges = updateShiftedIntervals(event)
 
     validUndoManager?.undoableActionPerformed(object : BasicUndoableAction(documentReference) {
       override fun undo() {
-        changesContext = UndoContext(eventChanges)
+        ThreadingAssertions.assertWriteAccess()
+        updatePointersByChanges(invertChanges(shiftChanges))
+        val invertChanges = invertChanges(eventChanges)
+        updatePointersByChanges(invertChanges)
+        onUpdated(NotebookIntervalPointersEvent(invertChanges, event, EventSource.UNDO_ACTION))
       }
 
-      override fun redo() {}
+      override fun redo() {
+        ThreadingAssertions.assertWriteAccess()
+        updatePointersByChanges(eventChanges)
+        updatePointersByChanges(shiftChanges)
+        postponedEvent = NotebookIntervalPointersEvent(eventChanges, event, EventSource.REDO_ACTION)
+      }
     })
-
-    documentChangedContext?.let {
-      it.redoContext = RedoContext(eventChanges)
-    }
-
-    return NotebookIntervalPointersEvent(eventChanges, event, EventSource.ACTION)
-  }
-
-  private fun documentChangedByUndo(event: NotebookCellLinesEvent, context: UndoContext): NotebookIntervalPointersEvent {
-    val invertedChanges = invertChanges(context.changes)
-    updatePointersByChanges(invertedChanges)
-    updateShiftedIntervals(event)
-    return NotebookIntervalPointersEvent(invertedChanges, event, EventSource.UNDO_ACTION)
-  }
-
-  private fun documentChangedByRedo(event: NotebookCellLinesEvent, context: RedoContext): NotebookIntervalPointersEvent {
-    updatePointersByChanges(context.changes)
-    updateShiftedIntervals(event)
-    return NotebookIntervalPointersEvent(context.changes, event, EventSource.REDO_ACTION)
+    onUpdated(NotebookIntervalPointersEvent(eventChanges, event, EventSource.ACTION))
   }
 
   private fun updatePointersByChanges(changes: List<Change>) {
@@ -215,7 +175,8 @@ class NotebookIntervalPointerFactoryImpl(private val notebookCellLines: Notebook
     return old.type == new.type && old.language == new.language
   }
 
-  private fun updateChangedIntervals(e: NotebookCellLinesEvent, eventChanges: NotebookIntervalPointersEventChanges) {
+  private fun updateChangedIntervals(e: NotebookCellLinesEvent): NotebookIntervalPointersEventChanges {
+    val eventChanges = NotebookIntervalPointersEventChanges()
     when {
       !e.isIntervalsChanged() -> {
         // content edited without affecting intervals values
@@ -257,17 +218,25 @@ class NotebookIntervalPointerFactoryImpl(private val notebookCellLines: Notebook
         }
       }
     }
+    return eventChanges
   }
 
-  private fun updateShiftedIntervals(event: NotebookCellLinesEvent) {
+  private fun updateShiftedIntervals(event: NotebookCellLinesEvent): NotebookIntervalPointersEventChanges {
     val invalidPointersStart =
       event.newIntervals.firstOrNull()?.let { it.ordinal + event.newIntervals.size }
       ?: event.oldIntervals.firstOrNull()?.ordinal
       ?: pointers.size
 
+    val eventChanges = NotebookIntervalPointersEventChanges()
+    val intervals = notebookCellLines.intervals
     for (i in invalidPointersStart until pointers.size) {
-      pointers[i].interval = notebookCellLines.intervals[i]
+      val ptr = pointers[i]
+      val intervalBefore = ptr.interval!!
+      val intervalAfter = intervals[i]
+      ptr.interval = intervals[i]
+      eventChanges.add(OnEdited(ptr, intervalBefore, intervalAfter))
     }
+    return eventChanges
   }
 
   private fun applyChanges(changes: Iterable<NotebookIntervalPointerFactory.Change>, eventChanges: NotebookIntervalPointersEventChanges) {

notebooks/visualization/src/org/jetbrains/plugins/notebooks/visualization/outputs/NotebookOutputInlayController.kt

@@ -12,6 +12,7 @@ import com.intellij.openapi.util.Disposer
 import com.intellij.openapi.util.registry.Registry
 import com.intellij.util.asSafely
 import com.intellij.util.messages.Topic
+import org.jetbrains.plugins.notebooks.ui.isFoldingEnabledKey
 import org.jetbrains.plugins.notebooks.ui.visualization.notebookAppearance
 import org.jetbrains.plugins.notebooks.visualization.*
 import org.jetbrains.plugins.notebooks.visualization.outputs.NotebookOutputComponentFactory.Companion.gutterPainter
@@ -80,22 +81,6 @@ class NotebookOutputInlayController private constructor(
     }
   }
 
-  fun checkAndUpdateInlayPosition(): Inlay<*>? {
-    val lines = intervalPointer.get()?.lines ?: return null
-    val lineEnd = computeInlayOffset(editor.document, lines)
-    if (inlay.offset != lineEnd) {
-      val oldInlay = inlay
-      isInReplaceInlay = true
-      Disposer.dispose(inlay)
-      inlay = createInlay()
-      isInReplaceInlay = false
-      return oldInlay
-    }
-    else {
-      return null
-    }
-  }
-
   private fun createInlay() = editor.addComponentInlay(
     outerComponent,
     isRelatedToPrecedingText = true,
@@ -124,7 +109,9 @@ class NotebookOutputInlayController private constructor(
     for (collapsingComponent in innerComponent.components) {
       val mainComponent = (collapsingComponent as CollapsingComponent).mainComponent
 
-      collapsingComponent.paintGutter(editor, yOffset, g)
+      if (editor.getUserData(isFoldingEnabledKey) != true) {
+        collapsingComponent.paintGutter(editor, yOffset, g)
+      }
 
       mainComponent.gutterPainter?.let { painter ->
         mainComponent.yOffsetFromEditor(editor)?.let { yOffset ->
@@ -278,6 +265,14 @@ class NotebookOutputInlayController private constructor(
     return result
   }
 
+  fun toggle() {
+    collapsingComponents.forEach { collapsingComponent ->
+      if (collapsingComponent.isWorthCollapsing) {
+        collapsingComponent.isSeen = !collapsingComponent.isSeen
+      }
+    }
+  }
+
   class Factory : NotebookCellInlayController.Factory {
     override fun compute(
       editor: EditorImpl,

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

@@ -0,0 +1,135 @@
+package org.jetbrains.plugins.notebooks.visualization.ui
+
+import com.intellij.ide.DataManager
+import com.intellij.openapi.actionSystem.DataProvider
+import com.intellij.openapi.actionSystem.PlatformCoreDataKeys
+import com.intellij.openapi.actionSystem.PlatformDataKeys
+import com.intellij.openapi.diagnostic.logger
+import com.intellij.openapi.diagnostic.thisLogger
+import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.editor.ex.EditorEx
+import com.intellij.openapi.editor.impl.EditorImpl
+import com.intellij.openapi.util.Disposer
+import com.intellij.util.asSafely
+import org.jetbrains.plugins.notebooks.visualization.NOTEBOOK_CELL_LINES_INTERVAL_DATA_KEY
+import org.jetbrains.plugins.notebooks.visualization.NotebookCellInlayController
+import org.jetbrains.plugins.notebooks.visualization.NotebookCellLines
+import org.jetbrains.plugins.notebooks.visualization.NotebookIntervalPointer
+import org.jetbrains.plugins.notebooks.visualization.outputs.NotebookOutputInlayController
+import javax.swing.JComponent
+
+class EditorCell(
+  private val editor: Editor,
+  private val intervals: NotebookCellLines,
+  internal var intervalPointer: NotebookIntervalPointer
+) {
+
+  private var _controllers: List<NotebookCellInlayController> = emptyList()
+  val controllers: List<NotebookCellInlayController>
+    get() = _controllers + (input.inputController?.let { listOf(it) } ?: emptyList())
+
+  private val interval get() = intervalPointer.get() ?: error("Invalid interval")
+
+  private var input: EditorCellInput = EditorCellInput(
+    editor as EditorEx,
+    { currentController: NotebookCellInlayController? ->
+      getInputFactories().firstNotNullOfOrNull {
+        failSafeCompute(it, editor, currentController?.let { listOf(it) }
+                                    ?: emptyList(), intervals.intervals.listIterator(interval.ordinal))
+      }
+    }, intervalPointer)
+
+  private var output: EditorCellOutput? = null
+
+  init {
+    update()
+  }
+
+  fun dispose() {
+    controllers.forEach { controller ->
+      disposeController(controller)
+    }
+    input.dispose()
+    output?.dispose()
+  }
+
+  private fun disposeController(controller: NotebookCellInlayController) {
+    val inlay = controller.inlay
+    inlay.renderer.asSafely<JComponent>()?.let { DataManager.removeDataProvider(it) }
+    Disposer.dispose(inlay)
+  }
+
+  fun update() {
+    val otherFactories = NotebookCellInlayController.Factory.EP_NAME.extensionList
+      .filter { it !is NotebookCellInlayController.InputFactory }
+
+    val controllersToDispose = controllers.toMutableSet()
+    _controllers = if (!editor.isDisposed) {
+      otherFactories.mapNotNull { factory -> failSafeCompute(factory, editor, controllers, intervals.intervals.listIterator(interval.ordinal)) }
+    }
+    else {
+      emptyList()
+    }
+    controllersToDispose.removeAll(controllers.toSet())
+    controllersToDispose.forEach { disposeController(it) }
+    for (controller in controllers) {
+      val inlay = controller.inlay
+      inlay.renderer.asSafely<JComponent>()?.let { component ->
+        val oldProvider = DataManager.getDataProvider(component)
+        if (oldProvider != null && oldProvider !is NotebookCellDataProvider) {
+          LOG.error("Overwriting an existing CLIENT_PROPERTY_DATA_PROVIDER. Old provider: $oldProvider")
+        }
+        DataManager.removeDataProvider(component)
+        DataManager.registerDataProvider(component, NotebookCellDataProvider(editor, component) { interval })
+      }
+    }
+    input.update()
+    output?.dispose()
+    val outputController = controllers.filterIsInstance<NotebookOutputInlayController>().firstOrNull()
+    if (outputController != null) {
+      output = EditorCellOutput(editor as EditorEx, outputController)
+    }
+  }
+
+  private fun getInputFactories(): Sequence<NotebookCellInlayController.Factory> {
+    return NotebookCellInlayController.Factory.EP_NAME.extensionList.asSequence()
+      .filter { it is NotebookCellInlayController.InputFactory }
+  }
+
+  private fun failSafeCompute(factory: NotebookCellInlayController.Factory,
+                              editor: Editor,
+                              controllers: Collection<NotebookCellInlayController>,
+                              intervalIterator: ListIterator<NotebookCellLines.Interval>): NotebookCellInlayController? {
+    try {
+      return factory.compute(editor as EditorImpl, controllers, intervalIterator)
+    }
+    catch (t: Throwable) {
+      thisLogger().error("${factory.javaClass.name} shouldn't throw exceptions at NotebookCellInlayController.Factory.compute(...)", t)
+      return null
+    }
+  }
+
+  fun updatePositions() {
+    input.updatePositions()
+    output?.updatePositions()
+  }
+
+  companion object {
+    private val LOG = logger<EditorCell>()
+  }
+
+  private data class NotebookCellDataProvider(
+    val editor: Editor,
+    val component: JComponent,
+    val intervalProvider: () -> NotebookCellLines.Interval,
+  ) : DataProvider {
+    override fun getData(key: String): Any? =
+      when (key) {
+        NOTEBOOK_CELL_LINES_INTERVAL_DATA_KEY.name -> intervalProvider()
+        PlatformCoreDataKeys.CONTEXT_COMPONENT.name -> component
+        PlatformDataKeys.EDITOR.name -> editor
+        else -> null
+      }
+  }
+
+}

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

@@ -0,0 +1,43 @@
+package org.jetbrains.plugins.notebooks.visualization.ui
+
+import com.intellij.openapi.editor.ex.EditorEx
+import org.jetbrains.plugins.notebooks.ui.isFoldingEnabledKey
+import org.jetbrains.plugins.notebooks.visualization.NotebookCellInlayController
+import org.jetbrains.plugins.notebooks.visualization.NotebookIntervalPointer
+
+class EditorCellFolding(editor: EditorEx, toggleListener: () -> Unit) {
+
+  private val foldingBar: EditorCellFoldingBar?
+
+  private val binder: EditorCellFoldingBarLocationBinder?
+
+  init {
+    val isFoldingEnabled = editor.getUserData(isFoldingEnabledKey) ?: false
+    if (isFoldingEnabled) {
+      foldingBar = EditorCellFoldingBar(editor, toggleListener)
+      binder = EditorCellFoldingBarLocationBinder(editor, foldingBar)
+    }
+    else {
+      foldingBar = null
+      binder = null
+    }
+  }
+
+  fun bindTo(controller: NotebookCellInlayController?) {
+    binder?.bindTo(controller)
+  }
+
+  fun dispose() {
+    binder?.dispose()
+    foldingBar?.dispose()
+  }
+
+  fun bindTo(interval: NotebookIntervalPointer) {
+    binder?.bindTo(interval)
+  }
+
+  fun updatePosition() {
+    binder?.updatePositions()
+  }
+
+}

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

@@ -0,0 +1,50 @@
+package org.jetbrains.plugins.notebooks.visualization.ui
+
+import com.intellij.openapi.editor.ex.EditorEx
+import org.jetbrains.plugins.notebooks.ui.visualization.notebookAppearance
+import java.awt.Cursor
+import java.awt.Dimension
+import java.awt.Point
+import java.awt.event.MouseAdapter
+import java.awt.event.MouseEvent
+import javax.swing.JPanel
+
+internal class EditorCellFoldingBar(
+  private val editor: EditorEx,
+  private val toggleListener: () -> Unit
+) {
+
+  val panel = JPanel().also {
+
+    val appearance = editor.notebookAppearance
+
+    it.background = appearance.getCellStripeColor(editor)
+    it.cursor = Cursor.getPredefinedCursor(Cursor.HAND_CURSOR)
+    it.addMouseListener(object : MouseAdapter() {
+      override fun mouseEntered(e: MouseEvent) {
+        it.background = appearance.getCellStripeHoverColor(editor)
+      }
+
+      override fun mouseExited(e: MouseEvent) {
+        it.background = appearance.getCellStripeColor(editor)
+      }
+
+      override fun mouseClicked(e: MouseEvent) {
+        toggleListener.invoke()
+      }
+    })
+  }
+
+  init {
+    editor.gutterComponentEx.add(panel)
+  }
+
+  fun dispose() {
+    editor.gutterComponentEx.remove(panel)
+  }
+
+  fun setLocation(y: Int, height: Int) {
+    panel.location = Point(editor.gutterComponentEx.extraLineMarkerFreePaintersAreaOffset, y)
+    panel.size = Dimension(8, height)
+  }
+}

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

@@ -0,0 +1,64 @@
+package org.jetbrains.plugins.notebooks.visualization.ui
+
+import com.intellij.openapi.editor.ex.EditorEx
+import com.intellij.util.asSafely
+import org.jetbrains.plugins.notebooks.visualization.NotebookCellInlayController
+import org.jetbrains.plugins.notebooks.visualization.NotebookIntervalPointer
+import java.awt.event.ComponentAdapter
+import java.awt.event.ComponentEvent
+import javax.swing.JComponent
+
+internal class EditorCellFoldingBarLocationBinder(private val editor: EditorEx, private val foldingBar: EditorCellFoldingBar) {
+
+  private val listener = object : ComponentAdapter() {
+    override fun componentResized(e: ComponentEvent) {
+      foldingBar.setLocation(e.component.location.y, e.component.size.height)
+    }
+  }
+
+  private var currentController: NotebookCellInlayController? = null
+  private var currentInterval: NotebookIntervalPointer? = null
+
+  fun dispose() {
+    unbind()
+  }
+
+  private fun unbind() {
+    currentController?.let {
+      it.inlay.renderer.asSafely<JComponent>()?.removeComponentListener(listener)
+    }
+    currentController = null
+  }
+
+  fun bindTo(controller: NotebookCellInlayController?) {
+    if (controller != currentController) {
+      unbind()
+      if (controller != null) {
+        controller.inlay.renderer.asSafely<JComponent>()?.addComponentListener(listener)
+        currentController = controller
+      }
+    }
+  }
+
+  fun bindTo(interval: NotebookIntervalPointer) {
+    currentInterval = interval
+  }
+
+  fun updatePositions() {
+    val controller = currentController
+    if (controller != null) {
+      val component = controller.inlay.renderer as JComponent
+      foldingBar.setLocation(component.location.y, component.size.height)
+    }
+    else {
+      val interval = currentInterval?.get()
+      if (interval != null) {
+        val startOffset = editor.document.getLineStartOffset(interval.lines.first)
+        val endOffset = editor.document.getLineEndOffset(interval.lines.last)
+        val top = editor.offsetToXY(startOffset).y
+        val height = editor.offsetToXY(endOffset).y + editor.lineHeight - top
+        foldingBar.setLocation(top, height)
+      }
+    }
+  }
+}

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

@@ -0,0 +1,95 @@
+package org.jetbrains.plugins.notebooks.visualization.ui
+
+import com.intellij.openapi.editor.ex.EditorEx
+import com.intellij.openapi.util.Disposer
+import com.intellij.openapi.util.TextRange
+import org.jetbrains.plugins.notebooks.visualization.NotebookCellInlayController
+import org.jetbrains.plugins.notebooks.visualization.NotebookIntervalPointer
+
+internal class EditorCellInput(
+  private val editor: EditorEx,
+  private val inputControllerFactory: ((NotebookCellInlayController?) -> NotebookCellInlayController?)?,
+  private val intervalPointer: NotebookIntervalPointer
+) {
+
+  private val folding: EditorCellFolding = EditorCellFolding(editor) {
+    if (inputControllerFactory == null) {
+      toggleTextFolding()
+    }
+    else {
+      toggleFolding(inputControllerFactory)
+    }
+  }
+
+  private fun toggleFolding(inputControllerFactory: (NotebookCellInlayController?) -> NotebookCellInlayController?) {
+    val controller = inputController
+    inputController = if (controller != null) {
+      Disposer.dispose(controller.inlay)
+      toggleTextFolding()
+      null
+    }
+    else {
+      toggleTextFolding()
+      inputControllerFactory.invoke(inputController)
+    }
+    folding.bindTo(inputController)
+  }
+
+  private fun toggleTextFolding() {
+    val interval = intervalPointer.get() ?: error("Invalid interval")
+    val startOffset = editor.document.getLineStartOffset(interval.lines.first)
+    val endOffset = editor.document.getLineEndOffset(interval.lines.last)
+    val foldingModel = editor.foldingModel
+    val foldRegion = foldingModel.getFoldRegion(startOffset, endOffset)
+    if (foldRegion == null) {
+      foldingModel.runBatchFoldingOperation {
+        val text = editor.document.getText(TextRange(startOffset, endOffset))
+        val placeholder = text.lines().drop(1).firstOrNull { it.trim().isNotEmpty() }?.ellipsis(30) ?: "..."
+        foldingModel.createFoldRegion(startOffset, endOffset, placeholder, null, true)
+      }
+    }
+    else {
+      foldingModel.runBatchFoldingOperation {
+        foldingModel.removeFoldRegion(foldRegion)
+      }
+    }
+  }
+
+  internal var inputController: NotebookCellInlayController? = createOrUpdateController()
+
+  init {
+    folding.bindTo(inputController)
+    folding.bindTo(intervalPointer)
+  }
+
+  fun dispose() {
+    folding.dispose()
+    inputController?.let { controller -> Disposer.dispose(controller.inlay) }
+  }
+
+  fun update() {
+    inputController = createOrUpdateController()
+  }
+
+  private fun createOrUpdateController(): NotebookCellInlayController? {
+    val actualController = inputControllerFactory?.invoke(inputController)
+    if (actualController != inputController) {
+      inputController?.let { controller -> Disposer.dispose(controller.inlay) }
+      folding.bindTo(actualController)
+    }
+    return actualController
+  }
+
+  fun updatePositions() {
+    folding.updatePosition()
+  }
+}
+
+private fun String.ellipsis(length: Int): String {
+  return if (this.length > length) {
+    substring(0, length - 3) + "..."
+  }
+  else {
+    this
+  }
+}

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

@@ -0,0 +1,20 @@
+package org.jetbrains.plugins.notebooks.visualization.ui
+
+import com.intellij.openapi.editor.ex.EditorEx
+import org.jetbrains.plugins.notebooks.visualization.outputs.NotebookOutputInlayController
+
+internal class EditorCellOutput(editor: EditorEx, private val outputController: NotebookOutputInlayController) {
+
+  private val folding = EditorCellFolding(editor) { outputController.toggle() }.also {
+    it.bindTo(outputController)
+  }
+
+  fun updatePositions() {
+    folding.updatePosition()
+  }
+
+  fun dispose() {
+    folding.dispose()
+  }
+
+}