DS-1911 Notebook scrolls to a strange position when collapsing cells …

…from command mode

GitOrigin-RevId: 9607c4dfeb6543ff8a0479d959ca8d4e1e5a98db

visualisation/src/org/jetbrains/plugins/notebooks/editor/NotebookCellInlayManager.kt

@@ -26,7 +26,6 @@ import com.intellij.util.ui.update.MergingUpdateQueue
 import com.intellij.util.ui.update.Update
 import org.jetbrains.annotations.TestOnly
 import org.jetbrains.plugins.notebooks.editor.outputs.impl.DOCUMENT_BEING_UPDATED
-import org.jetbrains.plugins.notebooks.editor.outputs.impl.markScrollingPosition
 import java.awt.Graphics
 import javax.swing.JComponent
 import kotlin.math.max
@@ -189,7 +188,6 @@ class NotebookCellInlayManager private constructor(val editor: EditorImpl) {
   }
 
   private fun updateConsequentInlays(interestingRange: IntRange) {
-    markScrollingPosition(editor)
     val matchingIntervals = notebookCellLines.getMatchingCells(interestingRange)
     val fullInterestingRange =
       if (matchingIntervals.isNotEmpty()) matchingIntervals.first().lines.first..matchingIntervals.last().lines.last

visualisation/src/org/jetbrains/plugins/notebooks/editor/outputs/impl/NotebookOutputCollapsingActions.kt

@@ -4,11 +4,13 @@ import com.intellij.openapi.actionSystem.AnActionEvent
 import com.intellij.openapi.actionSystem.PlatformDataKeys
 import com.intellij.openapi.actionSystem.ToggleAction
 import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.editor.LogicalPosition
 import com.intellij.openapi.editor.impl.EditorImpl
 import com.intellij.openapi.project.DumbAware
 import com.intellij.util.castSafelyTo
 import org.jetbrains.plugins.notebooks.editor.NotebookCellInlayManager
 import org.jetbrains.plugins.notebooks.editor.NotebookCellLines
+import org.jetbrains.plugins.notebooks.editor.getCell
 import org.jetbrains.plugins.notebooks.editor.notebookCellLinesInterval
 import org.jetbrains.plugins.notebooks.editor.outputs.NotebookOutputInlayController
 
@@ -22,6 +24,8 @@ internal class NotebookOutputCollapseAllAction private constructor() : ToggleAct
     !allCollapsingComponents(e).any { it.isSeen }
 
   override fun setSelected(e: AnActionEvent, state: Boolean) {
+    markScrollingPositionBeforeOutputCollapseToggle(e)
+
     for (component in allCollapsingComponents(e)) {
       component.isSeen = !state
     }
@@ -48,6 +52,8 @@ internal class NotebookOutputCollapseAllInCellAction private constructor() : Tog
   }
 
   override fun setSelected(e: AnActionEvent, state: Boolean) {
+    markScrollingPositionBeforeOutputCollapseToggle(e)
+
     getCollapsingComponents(e)?.forEach {
       it.isSeen = !state
     }
@@ -64,10 +70,7 @@ internal class NotebookOutputCollapseSingleInCellAction private constructor() :
     getExpectedComponent(e)?.isSeen?.let { !it } ?: false
 
   override fun setSelected(e: AnActionEvent, state: Boolean) {
-    val inlayManager = e.notebookCellInlayManager
-    if (inlayManager != null) {
-      markScrollingPosition(inlayManager.editor)
-    }
+    markScrollingPositionBeforeOutputCollapseToggle(e)
     getExpectedComponent(e)?.isSeen = !state
   }
 
@@ -95,4 +98,36 @@ private val AnActionEvent.notebookCellInlayManager: NotebookCellInlayManager?
   get() = getData(PlatformDataKeys.EDITOR)?.let(NotebookCellInlayManager.Companion::get)
 
 private val AnActionEvent.notebookEditor: EditorImpl?
-  get() = notebookCellInlayManager?.editor
+  get() = notebookCellInlayManager?.editor
+
+/**
+ * Chooses the output's cell or the next one depending on scrolling position
+ */
+private fun findTargetCellForCollapseAction(editor: Editor, outputCell: NotebookCellLines.Interval): NotebookCellLines.Interval {
+  val outputCellBottom = editor.logicalPositionToXY(LogicalPosition(outputCell.lines.last, 0)).y + editor.lineHeight
+  val visibleArea = editor.scrollingModel.visibleAreaOnScrollingFinished
+  if (outputCellBottom < visibleArea.y) {
+    if (outputCell.lines.last < editor.document.lineCount - 1) {
+      val nextCell = editor.getCell(outputCell.lines.last + 1)
+      val nextCellTop = editor.logicalPositionToXY(LogicalPosition(nextCell.lines.first, 0)).y
+      if (nextCellTop < visibleArea.y + visibleArea.height) {
+        return nextCell
+      }
+    }
+  }
+  return outputCell
+}
+
+private fun markScrollingPositionBeforeOutputCollapseToggle(e: AnActionEvent) {
+  val inlayManager = e.notebookCellInlayManager
+  if (inlayManager != null) {
+    val cell = e.dataContext.notebookCellLinesInterval
+    if (cell != null) {
+      val targetCell = findTargetCellForCollapseAction(inlayManager.editor, cell)
+      markScrollingPosition(inlayManager.editor, targetCell)
+    }
+    else {
+      markScrollingPosition(inlayManager.editor)
+    }
+  }
+}

visualisation/src/org/jetbrains/plugins/notebooks/editor/outputs/impl/NotebookOutputInlayControllerImplUtil.kt

@@ -10,6 +10,7 @@ import com.intellij.openapi.fileEditor.TextEditor
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.vfs.VirtualFile
 import org.jetbrains.annotations.TestOnly
+import org.jetbrains.plugins.notebooks.editor.NotebookCellLines
 import org.jetbrains.plugins.notebooks.editor.SwingClientProperty
 import org.jetbrains.plugins.notebooks.editor.cellSelectionModel
 import org.jetbrains.plugins.notebooks.editor.outputs.NotebookOutputInlayController
@@ -72,18 +73,33 @@ fun markScrollingPosition(virtualFile: VirtualFile, project: Project) {
 }
 
 fun markScrollingPosition(editor: Editor) {
+  val visibleArea: Rectangle = editor.scrollingModel.visibleAreaOnScrollingFinished
+  val selectedCell = editor.cellSelectionModel?.selectedCells?.firstOrNull()
+  var targetCell: NotebookCellLines.Interval? = null
+  if (selectedCell != null) {
+    val cellYTop = editor.logicalPositionToXY(LogicalPosition(selectedCell.lines.first, 0)).y
+    if (cellYTop >= visibleArea.y &&
+        cellYTop <= visibleArea.y + visibleArea.height) {
+      targetCell = selectedCell
+    }
+  }
+  markScrollingPosition(editor, targetCell)
+}
+
+fun markScrollingPosition(editor: Editor, targetCell: NotebookCellLines.Interval?) {
   val visibleArea: Rectangle = editor.scrollingModel.visibleAreaOnScrollingFinished
   if (visibleArea.height <= 0) {
     return
   }
   var topLine: Int = editor.xyToVisualPosition(visibleArea.location).line
 
   if (editor is EditorImpl) {
-    val selectedCell = editor.cellSelectionModel?.selectedCells?.firstOrNull()
-    if (selectedCell != null && editor.logicalPositionToXY(LogicalPosition(selectedCell.lines.first, 0)).y >= visibleArea.y && editor.logicalPositionToXY(LogicalPosition(selectedCell.lines.first, 0)).y <= visibleArea.y + visibleArea.height) {
-      topLine = selectedCell.lines.first
+    if (targetCell != null) {
+      topLine = targetCell.lines.first
     } else {
-      while (topLine < editor.document.lineCount - 1 && editor.logicalPositionToXY(LogicalPosition(topLine, 0)).y < visibleArea.y) {
+      // Looking for the first cell that starts on the screen
+      while (topLine < editor.document.lineCount - 1 && editor.logicalPositionToXY(LogicalPosition(topLine, 0)).y < visibleArea.y &&
+             editor.logicalPositionToXY(LogicalPosition(topLine + 1, 0)).y <= visibleArea.y + visibleArea.height) {
         topLine++
       }
     }