PY-79420 [Jupyter] Change editor position keeping logic to make it similar to JupyterLab

(cherry picked from commit abf51469ce24ea3eb8361d931b4ab115dbd233b9)

GitOrigin-RevId: a4bcf3bfbfb48c8fda5e57ad9aec488adb47fcd2

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

@@ -8,23 +8,16 @@ import com.intellij.notebooks.visualization.ui.EditorCellViewEventListener.CellV
 import com.intellij.notebooks.visualization.ui.EditorCellViewEventListener.EditorCellViewEvent
 import com.intellij.notebooks.visualization.ui.EditorLayerController.Companion.EDITOR_LAYER_CONTROLLER_KEY
 import com.intellij.openapi.application.ApplicationManager
-import com.intellij.openapi.application.WriteIntentReadAction
 import com.intellij.openapi.client.ClientSystemInfo
 import com.intellij.openapi.editor.Caret
 import com.intellij.openapi.editor.Editor
-import com.intellij.openapi.editor.event.*
-import com.intellij.openapi.editor.ex.util.EditorScrollingPositionKeeper
+import com.intellij.openapi.editor.event.CaretEvent
+import com.intellij.openapi.editor.event.CaretListener
+import com.intellij.openapi.editor.event.EditorMouseEventArea
 import com.intellij.openapi.editor.impl.EditorImpl
 import com.intellij.openapi.util.Disposer
-import com.intellij.openapi.util.use
 import com.intellij.ui.ComponentUtil
-import java.awt.BorderLayout
-import java.awt.Color
-import java.awt.Component
-import java.awt.Graphics
-import java.awt.Graphics2D
-import java.awt.GraphicsEnvironment
-import java.awt.Point
+import java.awt.*
 import java.awt.event.InputEvent
 import java.awt.event.MouseEvent
 import java.awt.event.MouseEvent.MOUSE_PRESSED
@@ -45,6 +38,8 @@ class DecoratedEditor private constructor(
   override var mouseOverCell: EditorCellView? = null
     private set
 
+  override val editorPositionKeeper: NotebookPositionKeeper = NotebookPositionKeeper(editorImpl)
+
   private val selectionModel = EditorCellSelectionModel(manager)
 
   private var selectionUpdateScheduled = AtomicBoolean(false)
@@ -170,7 +165,7 @@ class DecoratedEditor private constructor(
     }
 
     override fun validateTree() {
-      keepScrollingPositionWhile(editor) {
+      editor.notebookEditor.editorPositionKeeper.keepScrollingPositionWhile {
         JupyterBoundsChangeHandler.get(editor).postponeUpdates()
         super.validateTree()
         JupyterBoundsChangeHandler.get(editor).performPostponed()
@@ -317,20 +312,6 @@ class DecoratedEditor private constructor(
   }
 }
 
-internal fun <T> keepScrollingPositionWhile(editor: Editor, task: () -> T): T {
-  return WriteIntentReadAction.compute<T, Nothing> {
-    EditorScrollingPositionKeeper(editor).use { keeper ->
-      if (editor.isDisposed) {
-        return@compute task()
-      }
-      keeper.savePosition()
-      val r = task()
-      keeper.restorePosition(false)
-      r
-    }
-  }
-}
-
 /** lists assumed to be ordered and non-empty  */
 private fun hasIntersection(cells: List<NotebookCellLines.Interval>, others: List<NotebookCellLines.Interval>): Boolean =
   !(cells.last().ordinal < others.first().ordinal || cells.first().ordinal > others.last().ordinal)

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

@@ -56,7 +56,7 @@ internal class EditorEmbeddedComponentLayoutManager(private val editor: EditorEx
 
   override fun layoutContainer(parent: Container) {
     synchronized(parent.treeLock) {
-      keepScrollingPositionWhile(editor) {
+      editor.notebookEditor.editorPositionKeeper.keepScrollingPositionWhile {
         val visibleWidth = maxOf(myEditorScrollPane.getViewport().getWidth() - myEditorScrollPane.getVerticalScrollBar().getWidth(), 0)
         for (entry in constraints) {
           val component: JComponent = entry.first

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

@@ -7,6 +7,7 @@ import com.intellij.notebooks.visualization.NotebookCellLines
 interface NotebookEditor {
   val mouseOverCell: EditorCellView?
   fun inlayClicked(clickedCell: NotebookCellLines.Interval, ctrlPressed: Boolean, shiftPressed: Boolean, mouseButton: Int)
+  val editorPositionKeeper: NotebookPositionKeeper
 }
 
 internal val notebookEditorKey = Key.create<NotebookEditor>(NotebookEditor::class.java.name)

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

@@ -0,0 +1,79 @@
+// 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.ui
+
+import com.intellij.openapi.application.WriteIntentReadAction
+import com.intellij.openapi.editor.impl.EditorImpl
+
+class NotebookPositionKeeper(val editor: EditorImpl) {
+
+  fun getPosition(useCaretPositon: Boolean = true): Position {
+    val visibleArea = editor.scrollingModel.getVisibleAreaOnScrollingFinished()
+
+    val topLeftCornerOffset = if (useCaretPositon) {
+      val caretY = editor.visualLineToY(editor.caretModel.visualPosition.line)
+      if (visibleArea.height > 0 && (caretY + editor.lineHeight <= visibleArea.y || caretY >= (visibleArea.y + visibleArea.height))) {
+        editor.logicalPositionToOffset(editor.xyToLogicalPosition(visibleArea.location))
+      }
+      else {
+        editor.caretModel.offset
+      }
+    }
+    else {
+      editor.logicalPositionToOffset(editor.xyToLogicalPosition(visibleArea.location))
+    }
+    val viewportShift = editor.offsetToXY(topLeftCornerOffset).y - visibleArea.y
+    val position = Position(topLeftCornerOffset, viewportShift)
+    return position
+  }
+
+  fun restorePosition(position: Position) {
+    val (topLeftCornerOffset, viewportShift) = position
+    val scrollingModel = editor.scrollingModel
+    scrollingModel.disableAnimation()
+    val newY = editor.offsetToXY(topLeftCornerOffset).y - viewportShift
+    scrollingModel.scrollVertically(newY)
+    scrollingModel.enableAnimation()
+  }
+
+  fun <T> keepScrollingPositionWhile(task: () -> T): T {
+    return WriteIntentReadAction.compute<T, Nothing> {
+      if (editor.isDisposed) {
+        return@compute task()
+      }
+      val position = getPosition(false)
+      val (r, newOffset) = getOffsetProvider(position).use { offsetProvider ->
+        task() to offsetProvider.getOffset()
+      }
+      restorePosition(Position(newOffset, position.viewportShift))
+      r
+    }
+  }
+
+  private fun getOffsetProvider(position: Position): OffsetProvider {
+    return if (editor.caretModel.offset == position.topLeftCornerOffset) {
+      object : OffsetProvider {
+        override fun getOffset(): Int = editor.caretModel.offset
+
+        override fun close() {}
+      }
+    }
+    else {
+      object : OffsetProvider {
+        val myTopLeftCornerMarker = editor.document.createRangeMarker(position.topLeftCornerOffset, position.topLeftCornerOffset)
+
+        override fun getOffset(): Int = myTopLeftCornerMarker.startOffset
+
+        override fun close() {
+          myTopLeftCornerMarker.dispose()
+        }
+      }
+    }
+  }
+
+  private interface OffsetProvider : AutoCloseable {
+    fun getOffset(): Int
+  }
+
+  data class Position(val topLeftCornerOffset: Int, val viewportShift: Int)
+
+}

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

@@ -56,7 +56,7 @@ class UpdateManager(val editor: EditorImpl) : Disposable {
       updateCtx = newCtx
       try {
         if (keepScrollingPositon) {
-          keepScrollingPositionWhile(editor) {
+          editor.notebookEditor.editorPositionKeeper.keepScrollingPositionWhile {
             updateImpl(newCtx, block)
           }
         }