HUGE sidebar and typing performance improvements (#516)

packages/ui/src/components/tree/TreeItemList.tsx

@@ -1,5 +1,9 @@
+import type { Virtualizer } from "@tanstack/react-virtual";
+import { useVirtualizer } from "@tanstack/react-virtual";
+import { useAtomValue } from "jotai";
 import type { CSSProperties } from "react";
-import { Fragment } from "react";
+import { Fragment, useLayoutEffect, useMemo, useRef, useState } from "react";
+import { draggingIdsFamily } from "./atoms";
 import type { SelectableTreeNode } from "./common";
 import type { TreeProps } from "./Tree";
 import { TreeDropMarker } from "./TreeDropMarker";
@@ -22,9 +26,22 @@ export type TreeItemListProps<T extends { id: string }> = Pick<
     className?: string;
     forceDepth?: number;
     addTreeItemRef?: (item: T, n: TreeItemHandle | null) => void;
+    /**
+     * Enable virtualization by providing the scroll container. Rows are then
+     * windowed with @tanstack/react-virtual and only visible rows mount.
+     */
+    getScrollElement?: () => HTMLElement | null;
+    onVirtualizerReady?: (v: Virtualizer<HTMLElement, Element>) => void;
   };
 
-export function TreeItemList<T extends { id: string }>({
+export function TreeItemList<T extends { id: string }>(props: TreeItemListProps<T>) {
+  if (props.getScrollElement != null) {
+    return <VirtualTreeItemList {...props} getScrollElement={props.getScrollElement} />;
+  }
+  return <StaticTreeItemList {...props} />;
+}
+
+function StaticTreeItemList<T extends { id: string }>({
   className,
   getItemKey,
   nodes,
@@ -32,6 +49,8 @@ export function TreeItemList<T extends { id: string }>({
   treeId,
   forceDepth,
   addTreeItemRef,
+  getScrollElement: _getScrollElement,
+  onVirtualizerReady: _onVirtualizerReady,
   ...props
 }: TreeItemListProps<T>) {
   return (
@@ -53,3 +72,120 @@ export function TreeItemList<T extends { id: string }>({
     </ul>
   );
 }
+
+// Rows are --height-sm (2rem). Derive the pixel estimate from the actual root
+// font size so scroll math stays accurate under interface scaling.
+function estimateRowHeightPx() {
+  const rem = Number.parseFloat(getComputedStyle(document.documentElement).fontSize) || 16;
+  return 2 * rem;
+}
+
+function VirtualTreeItemList<T extends { id: string }>({
+  className,
+  getItemKey,
+  nodes,
+  style,
+  treeId,
+  forceDepth,
+  addTreeItemRef,
+  getScrollElement,
+  onVirtualizerReady,
+  ...props
+}: TreeItemListProps<T> & { getScrollElement: () => HTMLElement | null }) {
+  const listRef = useRef<HTMLUListElement>(null);
+
+  // Offset of the list within the scroll container (eg. container padding),
+  // so windowing and scrollToIndex targets aren't shifted by it
+  const [scrollMargin, setScrollMargin] = useState(0);
+  useLayoutEffect(() => {
+    const list = listRef.current;
+    const scroller = getScrollElement();
+    if (list == null || scroller == null) return;
+    const offset =
+      list.getBoundingClientRect().top - scroller.getBoundingClientRect().top + scroller.scrollTop;
+    setScrollMargin(offset);
+  }, [getScrollElement]);
+
+  const virtualizer = useVirtualizer({
+    count: nodes.length,
+    getScrollElement,
+    estimateSize: estimateRowHeightPx,
+    overscan: 10,
+    scrollMargin,
+  });
+
+  useLayoutEffect(() => {
+    onVirtualizerReady?.(virtualizer);
+  }, [virtualizer, onVirtualizerReady]);
+
+  const virtualItems = virtualizer.getVirtualItems();
+
+  // Rows being dragged stay mounted even after they scroll out of the window.
+  //
+  // dnd-kit keeps a live reference to the node being dragged and re-measures it whenever that
+  // node is replaced. Its scroll adjustment is a delta against a baseline captured when that rect
+  // was last measured, and re-measuring resets the baseline to wherever the list happens to be
+  // scrolled right then, discarding the scroll accumulated so far. Letting the dragged row unmount
+  // mid-drag therefore corrupts every coordinate dnd-kit derives, by a bit more each time it
+  // happens, which is why the drop indicator drifts further the longer a drag autoscrolls.
+  const draggingIds = useAtomValue(draggingIdsFamily(treeId));
+  const pinnedIndexes = useMemo(() => {
+    if (draggingIds.length === 0) return [];
+    const rendered = new Set(virtualItems.map((v) => v.index));
+    const out: number[] = [];
+    for (let i = 0; i < nodes.length; i++) {
+      const node = nodes[i];
+      if (node == null || rendered.has(i)) continue;
+      if (draggingIds.includes(node.node.item.id)) out.push(i);
+    }
+    return out;
+    // biome-ignore lint/correctness/useExhaustiveDependencies: keyed off the rendered range
+  }, [draggingIds, nodes, virtualItems]);
+
+  const renderRow = (index: number, start: number, measure: boolean) => {
+    const child = nodes[index];
+    if (child == null) return null;
+    return (
+      <div
+        // Key by item so window shifts don't remount rows unnecessarily
+        key={getItemKey(child.node.item)}
+        ref={measure ? virtualizer.measureElement : undefined}
+        data-index={index}
+        className="tree-row"
+        style={{
+          position: "absolute",
+          top: 0,
+          left: 0,
+          width: "100%",
+          transform: `translateY(${start - scrollMargin}px)`,
+        }}
+      >
+        <TreeItem
+          treeId={treeId}
+          setRef={addTreeItemRef}
+          node={child.node}
+          getItemKey={getItemKey}
+          depth={forceDepth == null ? child.depth : forceDepth}
+          {...props}
+        />
+        <TreeDropMarker node={child.node} treeId={treeId} index={index + 1} />
+      </div>
+    );
+  };
+
+  return (
+    <ul
+      ref={listRef}
+      style={{ ...style, height: `${virtualizer.getTotalSize()}px`, position: "relative" }}
+      className={className}
+    >
+      <TreeDropMarker node={null} treeId={treeId} index={0} />
+      {virtualItems.map((virtualItem) => renderRow(virtualItem.index, virtualItem.start, true))}
+      {/* NOTE: Not measured. Measuring an out-of-window row would write its size into the
+          virtualizer's cache under an index the window isn't tracking. */}
+      {pinnedIndexes.map((index) =>
+        renderRow(index, virtualizer.measurementsCache[index]?.start ?? 0, false),
+      )}
+    </ul>
+  );
+}