perf(webapp): keep trace tree scrolling smooth on large runs

Scrolling the run trace tree no longer slows down as the trace grows.

Mono-RevId: 97d00a66801f05663abb1b034e9fde52c9dd767e

Author: Mariefay

apps/webapp/app/components/primitives/TreeView/TreeView.tsx

@@ -2,7 +2,7 @@ import type { VirtualItem, Virtualizer } from "@tanstack/react-virtual";
 import { useVirtualizer } from "@tanstack/react-virtual";
 import { motion } from "framer-motion";
 import type { MutableRefObject, RefObject } from "react";
-import { useEffect, useMemo, useReducer, useRef } from "react";
+import { useCallback, useEffect, useMemo, useReducer, useRef } from "react";
 import { cn } from "~/utils/cn";
 import type { NodeState, NodesState } from "./reducer";
 import { reducer } from "./reducer";
@@ -213,9 +213,11 @@ export function useTree<TData, TFilterValue>({
 
   const [state, dispatch] = useReducer(
     reducer,
-    concreteStateFromInput({ tree, selectedId, collapsedIds, filter })
+    { tree, selectedId, collapsedIds, filter },
+    concreteStateFromInput
   );
-  const currentSelectedId = selectedIdFromState(state.nodes);
+  // Scans every node, so only redo it when node state changes, not on scroll renders.
+  const currentSelectedId = useMemo(() => selectedIdFromState(state.nodes), [state.nodes]);
 
   // id -> index lookup so getNodeProps resolves in O(1) instead of scanning
   // the whole tree per rendered row.
@@ -278,13 +280,22 @@ export function useTree<TData, TFilterValue>({
     dispatch({ type: "UPDATE_FILTER", payload: { filter: latestFilterRef.current } });
   }, [serializedFilterValue]);
 
+  // A new getItemKey makes the virtualizer re-measure every row, so only
+  // change it when the visible rows change.
+  const getItemKey = useCallback(
+    (index: number) => state.visibleNodeIds[index],
+    [state.visibleNodeIds]
+  );
+
   // oxlint-disable-next-line react/incompatible-library -- TanStack Virtual is not compatible with compiler memoization.
   const virtualizer = useVirtualizer({
     count: state.visibleNodeIds.length,
-    getItemKey: (index) => state.visibleNodeIds[index],
+    getItemKey,
     getScrollElement: () => parentRef.current,
     estimateSize: (index: number) => {
-      const treeItem = tree[index];
+      const id = state.visibleNodeIds[index];
+      const treeIndex = id === undefined ? undefined : treeIndexById.get(id);
+      const treeItem = treeIndex === undefined ? undefined : tree[treeIndex];
       if (!treeItem) return 0;
       return estimatedRowHeight({
         node: treeItem,

apps/webapp/test/treeViewVirtualizer.dom.test.ts

@@ -0,0 +1,121 @@
+// @vitest-environment jsdom
+import { createElement } from "react";
+import { createRoot, type Root } from "react-dom/client";
+import { act } from "react-dom/test-utils";
+import { afterEach, describe, expect, it } from "vitest";
+import {
+  flattenTree,
+  useTree,
+  type FlatTreeItem,
+  type UseTreeStateOutput,
+} from "~/components/primitives/TreeView/TreeView";
+
+type Data = { title: string };
+
+const tree = flattenTree<Data>({
+  id: "root",
+  data: { title: "root" },
+  children: [
+    {
+      id: "a",
+      data: { title: "a" },
+      children: [
+        { id: "a1", data: { title: "a1" } },
+        { id: "a2", data: { title: "a2" } },
+      ],
+    },
+    { id: "b", data: { title: "b" } },
+  ],
+});
+
+let container: HTMLDivElement | undefined;
+let root: Root | undefined;
+
+afterEach(() => {
+  if (root) act(() => root!.unmount());
+  container?.remove();
+  container = undefined;
+  root = undefined;
+});
+
+function render() {
+  let output: UseTreeStateOutput | undefined;
+  const estimates = new Map<number, string>();
+  const parentRef = { current: null };
+  const estimatedRowHeight = ({ node, index }: { node: FlatTreeItem<Data>; index: number }) => {
+    estimates.set(index, node.id);
+    return 32;
+  };
+  let filterCalls = 0;
+  const filter = {
+    value: "match-all",
+    fn: () => {
+      filterCalls++;
+      return true;
+    },
+  };
+
+  function Harness({ tick: _tick }: { tick: number }) {
+    output = useTree({ tree, parentRef, estimatedRowHeight, filter });
+    return null;
+  }
+
+  let tick = 0;
+  container = document.createElement("div");
+  document.body.appendChild(container);
+  root = createRoot(container);
+  act(() => root!.render(createElement(Harness, { tick })));
+
+  return {
+    output: () => output!,
+    estimates,
+    filterCalls: () => filterCalls,
+    rerender: () => {
+      tick++;
+      act(() => root!.render(createElement(Harness, { tick })));
+    },
+  };
+}
+
+describe("useTree virtualizer options", () => {
+  it("keeps getItemKey identity across re-renders with unchanged rows", () => {
+    const harness = render();
+    const before = harness.output().virtualizer.options.getItemKey;
+
+    harness.rerender();
+
+    expect(harness.output().virtualizer.options.getItemKey).toBe(before);
+  });
+
+  it("changes getItemKey when the visible rows change", () => {
+    const harness = render();
+    const before = harness.output().virtualizer.options.getItemKey;
+
+    act(() => harness.output().collapseNode("a"));
+
+    expect(harness.output().virtualizer.options.getItemKey).not.toBe(before);
+  });
+
+  it("doesn't rebuild the initial tree state on re-render", () => {
+    const harness = render();
+    const before = harness.filterCalls();
+
+    harness.rerender();
+
+    expect(harness.filterCalls()).toBe(before);
+  });
+
+  it("estimates each row from the node shown in that row", () => {
+    const harness = render();
+    act(() => harness.output().collapseNode("a"));
+    harness.estimates.clear();
+
+    harness.output().virtualizer.getTotalSize();
+
+    expect([...harness.estimates.entries()].sort(([a], [b]) => a - b)).toEqual([
+      [0, "root"],
+      [1, "a"],
+      [2, "b"],
+    ]);
+  });
+});