fix(ui): preserve viewport position when switching layouts (#185)

Author: benvinegar

src/ui/App.tsx

@@ -99,6 +99,8 @@ export function App({
   const sidebarScrollRef = useRef<ScrollBoxRenderable | null>(null);
   const diffScrollRef = useRef<ScrollBoxRenderable | null>(null);
   const wrapToggleScrollTopRef = useRef<number | null>(null);
+  const layoutToggleScrollTopRef = useRef<number | null>(null);
+  const [layoutToggleRequestId, setLayoutToggleRequestId] = useState(0);
   const [layoutMode, setLayoutMode] = useState<LayoutMode>(bootstrap.initialMode);
   const [themeId, setThemeId] = useState(
     () => resolveTheme(bootstrap.initialTheme, renderer.themeMode).id,
@@ -275,6 +277,13 @@ export function App({
     [maxCodeHorizontalOffset, wrapLines],
   );
 
+  /** Preserve the current review position before changing the active diff layout. */
+  const selectLayoutMode = useCallback((mode: LayoutMode) => {
+    layoutToggleScrollTopRef.current = diffScrollRef.current?.scrollTop ?? 0;
+    setLayoutToggleRequestId((current) => current + 1);
+    setLayoutMode(mode);
+  }, []);
+
   /** Toggle the global agent note layer on or off. */
   const toggleAgentNotes = () => {
     setShowAgentNotes((current) => !current);
@@ -472,7 +481,7 @@ export function App({
         moveToHunk: review.moveToHunk,
         refreshCurrentInput: triggerRefreshCurrentInput,
         requestQuit,
-        selectLayoutMode: setLayoutMode,
+        selectLayoutMode,
         selectThemeId: setThemeId,
         showAgentNotes,
         showHelp,
@@ -497,6 +506,7 @@ export function App({
       moveToAnnotatedHunk,
       requestQuit,
       review.moveToHunk,
+      selectLayoutMode,
       triggerRefreshCurrentInput,
       showAgentNotes,
       showHelp,
@@ -547,7 +557,7 @@ export function App({
     requestQuit,
     scrollCodeHorizontally,
     scrollDiff,
-    selectLayoutMode: setLayoutMode,
+    selectLayoutMode,
     showHelp,
     switchMenu,
     toggleAgentNotes,
@@ -707,6 +717,8 @@ export function App({
           showHunkHeaders={showHunkHeaders}
           wrapLines={wrapLines}
           wrapToggleScrollTop={wrapToggleScrollTopRef.current}
+          layoutToggleScrollTop={layoutToggleScrollTopRef.current}
+          layoutToggleRequestId={layoutToggleRequestId}
           selectedFileTopAlignRequestId={review.selectedFileTopAlignRequestId}
           selectedHunkRevealRequestId={review.selectedHunkRevealRequestId}
           theme={activeTheme}

src/ui/AppHost.interactions.test.tsx

@@ -364,6 +364,10 @@ function firstVisibleAddedLine(frame: string) {
   return frame.match(/line\d{2} = 1\d{2}/)?.[0] ?? null;
 }
 
+function firstVisibleSourceLineNumber(frame: string) {
+  return frame.match(/line(\d{2}) =/)?.[1] ?? null;
+}
+
 function firstVisibleAddedLineNumber(frame: string) {
   return frame.match(/▌\s*(\d+)\s+\+/)?.[1] ?? null;
 }
@@ -1244,6 +1248,67 @@ describe("App interactions", () => {
     }
   });
 
+  test("layout toggles preserve the current viewport anchor across split and stack", async () => {
+    const setup = await testRender(<AppHost bootstrap={createLineScrollBootstrap()} />, {
+      width: 220,
+      height: 12,
+    });
+
+    try {
+      await flush(setup);
+
+      let frame = setup.captureCharFrame();
+      expect(frame).toContain("line01 = 101");
+      expect(frame).not.toContain("line08 = 108");
+
+      for (let index = 0; index < 24; index += 1) {
+        await act(async () => {
+          await setup.mockInput.pressArrow("down");
+        });
+        await flush(setup);
+        frame = setup.captureCharFrame();
+        if (frame.includes("line08 = 108") && !frame.includes("line01 = 101")) {
+          break;
+        }
+      }
+
+      expect(frame).toContain("line08 = 108");
+      expect(frame).not.toContain("line01 = 101");
+      const anchoredLineNumber = firstVisibleSourceLineNumber(frame);
+      expect(anchoredLineNumber).not.toBeNull();
+
+      await act(async () => {
+        await setup.mockInput.typeText("2");
+      });
+      await flush(setup);
+      await act(async () => {
+        await Bun.sleep(80);
+        await setup.renderOnce();
+      });
+
+      frame = setup.captureCharFrame();
+      expect(frame).toContain(`line${anchoredLineNumber} =`);
+      expect(firstVisibleSourceLineNumber(frame)).toBe(anchoredLineNumber);
+
+      await act(async () => {
+        await setup.mockInput.typeText("1");
+      });
+      await flush(setup);
+      await act(async () => {
+        await Bun.sleep(80);
+        await setup.renderOnce();
+      });
+
+      frame = setup.captureCharFrame();
+      expect(frame).toContain(`line${anchoredLineNumber} =`);
+      expect(firstVisibleSourceLineNumber(frame)).toBe(anchoredLineNumber);
+    } finally {
+      await act(async () => {
+        setup.renderer.destroy();
+      });
+    }
+  });
+
   test("Space scrolls down by viewport", async () => {
     // Create a file with many lines so Space has room to scroll
     const before =

src/ui/components/panes/DiffPane.tsx

@@ -15,7 +15,6 @@ import { computeHunkRevealScrollTop } from "../../lib/hunkScroll";
 import {
   measureDiffSectionGeometry,
   type DiffSectionGeometry,
-  type DiffSectionRowBounds,
 } from "../../lib/diffSectionGeometry";
 import {
   buildFileSectionLayouts,
@@ -27,6 +26,11 @@ import {
 } from "../../lib/fileSectionLayout";
 import { diffHunkId, diffSectionId } from "../../lib/ids";
 import { findViewportCenteredHunkTarget } from "../../lib/viewportSelection";
+import {
+  findViewportRowAnchor,
+  resolveViewportRowAnchorTop,
+  type ViewportRowAnchor,
+} from "../../lib/viewportAnchor";
 import type { AppTheme } from "../../themes";
 import { DiffSection } from "./DiffSection";
 import { DiffFileHeaderRow } from "./DiffFileHeaderRow";
@@ -36,99 +40,6 @@ import { prefetchHighlightedDiff } from "../../diff/useHighlightedDiff";
 
 const EMPTY_VISIBLE_AGENT_NOTES: VisibleAgentNote[] = [];
 
-/** Identify the rendered diff row that currently owns the top of the viewport. */
-interface ViewportRowAnchor {
-  fileId: string;
-  rowKey: string;
-  rowOffsetWithin: number;
-}
-
-/** Find the rendered row bounds covering a vertical offset within one file body. */
-function binarySearchRowBounds(sectionRowBounds: DiffSectionRowBounds[], relativeTop: number) {
-  let low = 0;
-  let high = sectionRowBounds.length - 1;
-
-  while (low <= high) {
-    const mid = (low + high) >>> 1;
-    const rowBounds = sectionRowBounds[mid]!;
-
-    if (relativeTop < rowBounds.top) {
-      high = mid - 1;
-    } else if (relativeTop >= rowBounds.top + rowBounds.height) {
-      low = mid + 1;
-    } else {
-      return rowBounds;
-    }
-  }
-
-  return undefined;
-}
-
-/** Capture a stable top-row anchor from the pre-toggle layout so it can be restored later. */
-function findViewportRowAnchor(
-  files: DiffFile[],
-  sectionGeometry: DiffSectionGeometry[],
-  scrollTop: number,
-  headerHeights: number[],
-) {
-  const fileSectionLayouts = buildFileSectionLayouts(
-    files,
-    sectionGeometry.map((metrics) => metrics?.bodyHeight ?? 0),
-    headerHeights,
-  );
-
-  for (let index = 0; index < files.length; index += 1) {
-    const sectionLayout = fileSectionLayouts[index];
-    const bodyTop = sectionLayout?.bodyTop ?? 0;
-    const geometry = sectionGeometry[index];
-    const bodyHeight = geometry?.bodyHeight ?? 0;
-    const relativeTop = scrollTop - bodyTop;
-
-    if (relativeTop >= 0 && relativeTop < bodyHeight && geometry) {
-      const rowBounds = binarySearchRowBounds(geometry.rowBounds, relativeTop);
-      if (rowBounds) {
-        return {
-          fileId: files[index]!.id,
-          rowKey: rowBounds.key,
-          rowOffsetWithin: relativeTop - rowBounds.top,
-        } satisfies ViewportRowAnchor;
-      }
-    }
-  }
-
-  return null;
-}
-
-/** Resolve a captured row anchor into its new scrollTop after wrapping or layout changes. */
-function resolveViewportRowAnchorTop(
-  files: DiffFile[],
-  sectionGeometry: DiffSectionGeometry[],
-  anchor: ViewportRowAnchor,
-  headerHeights: number[],
-) {
-  const fileSectionLayouts = buildFileSectionLayouts(
-    files,
-    sectionGeometry.map((metrics) => metrics?.bodyHeight ?? 0),
-    headerHeights,
-  );
-
-  for (let index = 0; index < files.length; index += 1) {
-    const sectionLayout = fileSectionLayouts[index];
-    const bodyTop = sectionLayout?.bodyTop ?? 0;
-    const file = files[index];
-    const geometry = sectionGeometry[index];
-    if (file?.id === anchor.fileId && geometry) {
-      const rowBounds = geometry.rowBoundsByKey.get(anchor.rowKey);
-      if (rowBounds) {
-        return bodyTop + rowBounds.top + Math.min(anchor.rowOffsetWithin, rowBounds.height - 1);
-      }
-      return bodyTop;
-    }
-  }
-
-  return 0;
-}
-
 /** Keep syntax-highlight warm for the files immediately adjacent to the current selection. */
 function buildAdjacentPrefetchFileIds(files: DiffFile[], selectedFileId?: string) {
   if (!selectedFileId) {
@@ -217,6 +128,8 @@ export function DiffPane({
   showHunkHeaders,
   wrapLines,
   wrapToggleScrollTop,
+  layoutToggleScrollTop = null,
+  layoutToggleRequestId = 0,
   selectedFileTopAlignRequestId = 0,
   selectedHunkRevealRequestId,
   theme,
@@ -243,6 +156,8 @@ export function DiffPane({
   showHunkHeaders: boolean;
   wrapLines: boolean;
   wrapToggleScrollTop: number | null;
+  layoutToggleScrollTop?: number | null;
+  layoutToggleRequestId?: number;
   selectedFileTopAlignRequestId?: number;
   selectedHunkRevealRequestId?: number;
   theme: AppTheme;
@@ -333,8 +248,10 @@ export function DiffPane({
   const prevScrollTopRef = useRef(0);
   const previousSectionGeometryRef = useRef<DiffSectionGeometry[] | null>(null);
   const previousFilesRef = useRef<DiffFile[]>(files);
+  const previousLayoutRef = useRef(layout);
   const previousWrapLinesRef = useRef(wrapLines);
   const previousSelectedFileTopAlignRequestIdRef = useRef(selectedFileTopAlignRequestId);
+  const previousLayoutToggleRequestIdRef = useRef(layoutToggleRequestId);
   const previousSelectedHunkRevealRequestIdRef = useRef(selectedHunkRevealRequestId);
   const pendingFileTopAlignFileIdRef = useRef<string | null>(null);
   const suppressViewportSelectionSyncRef = useRef(false);
@@ -344,6 +261,7 @@ export function DiffPane({
   // Initialized to null so the first render never fires a selection change; a real scroll
   // is required before passive viewport-follow selection can trigger.
   const lastViewportSelectionTopRef = useRef<number | null>(null);
+  const lastViewportRowAnchorRef = useRef<ViewportRowAnchor | null>(null);
 
   /**
    * Ignore viewport-follow selection updates while the pane is scrolling to an explicit selection.
@@ -734,23 +652,29 @@ export function DiffPane({
   );
 
   useLayoutEffect(() => {
+    const layoutChanged = previousLayoutRef.current !== layout;
+    const explicitLayoutToggle = previousLayoutToggleRequestIdRef.current !== layoutToggleRequestId;
     const wrapChanged = previousWrapLinesRef.current !== wrapLines;
     const previousSectionMetrics = previousSectionGeometryRef.current;
     const previousFiles = previousFilesRef.current;
-    const previousSectionHeaderHeights = buildInStreamFileHeaderHeights(previousFiles);
 
-    if (wrapChanged && previousSectionMetrics && previousFiles.length > 0) {
+    if ((layoutChanged || wrapChanged) && previousSectionMetrics && previousFiles.length > 0) {
+      const previousSectionHeaderHeights = buildInStreamFileHeaderHeights(previousFiles);
       const previousScrollTop =
         // Prefer the synchronously captured pre-toggle position so anchor restoration does not
         // race the polling-based viewport snapshot.
-        wrapToggleScrollTop != null
+        wrapChanged && wrapToggleScrollTop != null
           ? wrapToggleScrollTop
-          : Math.max(prevScrollTopRef.current, scrollViewport.top);
+          : layoutChanged && explicitLayoutToggle && layoutToggleScrollTop != null
+            ? layoutToggleScrollTop
+            : (scrollRef.current?.scrollTop ??
+              Math.max(prevScrollTopRef.current, scrollViewport.top));
       const anchor = findViewportRowAnchor(
         previousFiles,
         previousSectionMetrics,
         previousScrollTop,
         previousSectionHeaderHeights,
+        lastViewportRowAnchorRef.current?.stableKey,
       );
       if (anchor) {
         const nextTop = resolveViewportRowAnchorTop(
@@ -763,13 +687,16 @@ export function DiffPane({
           scrollRef.current?.scrollTo(nextTop);
         };
 
+        lastViewportRowAnchorRef.current = anchor;
         suppressViewportSelectionSync();
         restoreViewportAnchor();
         // Retry across a couple of repaint cycles so the restored top-row anchor sticks
         // after wrapped row heights and viewport culling settle.
         const retryDelays = [0, 16, 48];
         const timeouts = retryDelays.map((delay) => setTimeout(restoreViewportAnchor, delay));
 
+        previousLayoutRef.current = layout;
+        previousLayoutToggleRequestIdRef.current = layoutToggleRequestId;
         previousWrapLinesRef.current = wrapLines;
         previousSectionGeometryRef.current = sectionGeometry;
         previousFilesRef.current = files;
@@ -780,11 +707,16 @@ export function DiffPane({
       }
     }
 
+    previousLayoutRef.current = layout;
+    previousLayoutToggleRequestIdRef.current = layoutToggleRequestId;
     previousWrapLinesRef.current = wrapLines;
     previousSectionGeometryRef.current = sectionGeometry;
     previousFilesRef.current = files;
   }, [
     files,
+    layout,
+    layoutToggleRequestId,
+    layoutToggleScrollTop,
     scrollRef,
     scrollViewport.top,
     sectionGeometry,
@@ -794,6 +726,26 @@ export function DiffPane({
     wrapToggleScrollTop,
   ]);
 
+  useLayoutEffect(() => {
+    if (files.length === 0) {
+      lastViewportRowAnchorRef.current = null;
+      return;
+    }
+
+    const currentScrollTop = scrollRef.current?.scrollTop ?? scrollViewport.top;
+    const nextAnchor = findViewportRowAnchor(
+      files,
+      sectionGeometry,
+      currentScrollTop,
+      sectionHeaderHeights,
+      lastViewportRowAnchorRef.current?.stableKey,
+    );
+
+    if (nextAnchor) {
+      lastViewportRowAnchorRef.current = nextAnchor;
+    }
+  }, [files, scrollRef, scrollViewport.top, sectionGeometry, sectionHeaderHeights]);
+
   useLayoutEffect(() => {
     if (previousSelectedFileTopAlignRequestIdRef.current === selectedFileTopAlignRequestId) {
       return;