refactor(tree-view): improve virtualization dx

segunadebayo · segunadebayo · commit 57478b87d44e · 2025-11-30T21:39:18.000Z
diff --git a/.changeset/cyan-wolves-travel.md b/.changeset/cyan-wolves-travel.md
@@ -3,4 +3,4 @@
 "@zag-js/tabs": patch
 ---
 
-bugfix: `tabs` and `radio-group` machine no longer show the `indicator` prematurely
+Fix issue where `tabs` and `radio-group` machine no longer show the `indicator` prematurely
diff --git a/.changeset/tree-view-virtualization.md b/.changeset/tree-view-virtualization.md
@@ -0,0 +1,19 @@
+---
+"@zag-js/tree-view": minor
+---
+
+- Added `scrollToIndexFn` prop to enable keyboard navigation in virtualized trees
+- **[Breaking]:** `getVisibleNodes()` now returns `{ node, indexPath }[]` instead of `node[]`. Returning the index path
+  perhaps the most useful use of this function, hence the change.
+
+**Migration:**
+
+```tsx
+// Before
+const nodes = api.getVisibleNodes()
+nodes.forEach((node) => console.log(node.id))
+
+// After
+const visibleNodes = api.getVisibleNodes()
+visibleNodes.forEach(({ node }) => console.log(node.id))
+```
diff --git a/examples/next-ts/pages/tree-view-virtualized.tsx b/examples/next-ts/pages/tree-view-virtualized.tsx
@@ -0,0 +1,205 @@
+import { useVirtualizer, type Virtualizer } from "@tanstack/react-virtual"
+import { normalizeProps, useMachine } from "@zag-js/react"
+import { treeviewControls } from "@zag-js/shared"
+import * as tree from "@zag-js/tree-view"
+import { ChevronRightIcon, FileIcon, FolderIcon } from "lucide-react"
+import { useId, useRef } from "react"
+import { StateVisualizer } from "../components/state-visualizer"
+import { Toolbar } from "../components/toolbar"
+import { useControls } from "../hooks/use-controls"
+
+interface Node {
+  id: string
+  name: string
+  children?: Node[]
+}
+
+// Generate a large tree for virtualization demo
+function generateLargeTree(): Node {
+  const folders: Node[] = []
+
+  for (let i = 0; i < 50; i++) {
+    const children: Node[] = []
+    for (let j = 0; j < 20; j++) {
+      children.push({
+        id: `folder-${i}/file-${j}.ts`,
+        name: `file-${j}.ts`,
+      })
+    }
+    folders.push({
+      id: `folder-${i}`,
+      name: `folder-${i}`,
+      children,
+    })
+  }
+
+  return {
+    id: "ROOT",
+    name: "",
+    children: folders,
+  }
+}
+
+const collection = tree.collection<Node>({
+  nodeToValue: (node) => node.id,
+  nodeToString: (node) => node.name,
+  rootNode: generateLargeTree(),
+})
+
+const ROW_HEIGHT = 32
+
+interface TreeNodeProps {
+  node: Node
+  indexPath: number[]
+  api: tree.Api
+}
+
+const TreeNode = (props: TreeNodeProps) => {
+  const { node, indexPath, api } = props
+
+  const nodeProps = { indexPath, node }
+  const nodeState = api.getNodeState(nodeProps)
+
+  // Calculate indentation based on depth
+  const indent = nodeState.depth * 20
+
+  if (nodeState.isBranch) {
+    return (
+      <div
+        {...api.getBranchControlProps(nodeProps)}
+        style={{
+          paddingLeft: indent,
+          height: ROW_HEIGHT,
+          display: "flex",
+          alignItems: "center",
+          gap: "4px",
+        }}
+      >
+        <span {...api.getBranchIndicatorProps(nodeProps)}>
+          <ChevronRightIcon
+            size={16}
+            style={{
+              transform: nodeState.expanded ? "rotate(90deg)" : "rotate(0deg)",
+              transition: "transform 0.15s",
+            }}
+          />
+        </span>
+        <FolderIcon size={16} />
+        <span {...api.getBranchTextProps(nodeProps)}>{node.name}</span>
+      </div>
+    )
+  }
+
+  return (
+    <div
+      {...api.getItemProps(nodeProps)}
+      style={{
+        paddingLeft: indent,
+        height: ROW_HEIGHT,
+        display: "flex",
+        alignItems: "center",
+        gap: "4px",
+      }}
+    >
+      <FileIcon size={16} />
+      <span {...api.getItemTextProps(nodeProps)}>{node.name}</span>
+    </div>
+  )
+}
+
+export default function Page() {
+  const controls = useControls(treeviewControls)
+  const parentRef = useRef<HTMLDivElement>(null)
+  const virtualizerRef = useRef<Virtualizer<HTMLDivElement, Element>>(null)
+
+  const service = useMachine(tree.machine, {
+    id: useId(),
+    collection,
+    ...controls.context,
+    scrollToIndexFn(details) {
+      virtualizerRef.current?.scrollToIndex(details.index, { align: "auto" })
+    },
+  })
+
+  const api = tree.connect(service, normalizeProps)
+
+  // Get visible nodes (now returns { node, indexPath }[])
+  const visibleNodes = api.getVisibleNodes()
+
+  const virtualizer = useVirtualizer({
+    count: visibleNodes.length,
+    getScrollElement: () => parentRef.current,
+    estimateSize: () => ROW_HEIGHT,
+    overscan: 10,
+  })
+
+  // Keep ref updated for scrollToIndexFn
+  virtualizerRef.current = virtualizer
+
+  return (
+    <>
+      <main className="tree-view">
+        <div {...api.getRootProps()}>
+          <h3 {...api.getLabelProps()}>Virtualized Tree ({visibleNodes.length} visible nodes)</h3>
+          <div style={{ display: "flex", gap: "10px", marginBottom: "10px" }}>
+            <button onClick={() => api.collapse()}>Collapse All</button>
+            <button onClick={() => api.expand()}>Expand All</button>
+            {controls.context.selectionMode === "multiple" && (
+              <>
+                <button onClick={() => api.select()}>Select All</button>
+                <button onClick={() => api.deselect()}>Deselect All</button>
+              </>
+            )}
+          </div>
+
+          {/* Scrollable container */}
+          <div
+            ref={parentRef}
+            {...api.getTreeProps()}
+            style={{
+              height: 400,
+              overflow: "auto",
+              border: "1px solid #ccc",
+              borderRadius: "4px",
+            }}
+          >
+            {/* Total size container */}
+            <div
+              style={{
+                height: `${virtualizer.getTotalSize()}px`,
+                width: "100%",
+                position: "relative",
+              }}
+            >
+              {/* Only render visible items */}
+              {virtualizer.getVirtualItems().map((virtualItem) => {
+                const { node, indexPath } = visibleNodes[virtualItem.index]
+
+                return (
+                  <div
+                    key={node.id}
+                    data-index={virtualItem.index}
+                    style={{
+                      position: "absolute",
+                      top: 0,
+                      left: 0,
+                      width: "100%",
+                      height: `${virtualItem.size}px`,
+                      transform: `translateY(${virtualItem.start}px)`,
+                    }}
+                  >
+                    <TreeNode node={node} indexPath={indexPath} api={api} />
+                  </div>
+                )
+              })}
+            </div>
+          </div>
+        </div>
+      </main>
+
+      <Toolbar controls={controls.ui}>
+        <StateVisualizer state={service} omit={["collection"]} />
+      </Toolbar>
+    </>
+  )
+}
diff --git a/packages/machines/tree-view/src/index.ts b/packages/machines/tree-view/src/index.ts
@@ -22,8 +22,10 @@ export type {
   NodeProps,
   NodeState,
   NodeWithError,
+  ScrollToIndexDetails,
   SelectionChangeDetails,
   TreeLoadingStatus,
   TreeLoadingStatusMap,
   TreeNode,
+  VisibleNode,
 } from "./tree-view.types"
diff --git a/packages/machines/tree-view/src/tree-view.connect.ts b/packages/machines/tree-view/src/tree-view.connect.ts
@@ -87,7 +87,7 @@ export function connect<T extends PropTypes, V extends TreeNode = TreeNode>(
       send({ type: value ? "NODE.SELECT" : "SELECTED.ALL", value, isTrusted: false })
     },
     getVisibleNodes() {
-      return computed("visibleNodes").map(({ node }) => node)
+      return computed("visibleNodes")
     },
     focus(value) {
       dom.focusNode(scope, value)
diff --git a/packages/machines/tree-view/src/tree-view.machine.ts b/packages/machines/tree-view/src/tree-view.machine.ts
@@ -1,6 +1,7 @@
 import type { TreeNode } from "@zag-js/collection"
+import type { Params } from "@zag-js/core"
 import { createGuards, createMachine } from "@zag-js/core"
-import { getByTypeahead, setElementValue } from "@zag-js/dom-query"
+import { getByTypeahead, raf, setElementValue } from "@zag-js/dom-query"
 import { addOrRemove, diff, first, isArray, isEqual, last, remove, toArray, uniq } from "@zag-js/utils"
 import { collection } from "./tree-view.collection"
 import * as dom from "./tree-view.dom"
@@ -358,47 +359,63 @@ export const machine = createMachine<TreeViewSchema>({
       clearSelected({ context }) {
         context.set("selectedValue", [])
       },
-      focusTreeFirstNode({ prop, scope }) {
+      focusTreeFirstNode(params) {
+        const { prop, scope } = params
         const collection = prop("collection")
         const firstNode = collection.getFirstNode()
         const firstValue = collection.getNodeValue(firstNode)
-        dom.focusNode(scope, firstValue)
+        const scrolled = scrollToNode(params, firstValue)
+        if (scrolled) raf(() => dom.focusNode(scope, firstValue))
+        else dom.focusNode(scope, firstValue)
       },
       focusTreeLastNode(params) {
         const { prop, scope } = params
         const collection = prop("collection")
         const lastNode = collection.getLastNode(undefined, { skip: skipFn(params) })
         const lastValue = collection.getNodeValue(lastNode)
-        dom.focusNode(scope, lastValue)
+        const scrolled = scrollToNode(params, lastValue)
+        if (scrolled) raf(() => dom.focusNode(scope, lastValue))
+        else dom.focusNode(scope, lastValue)
       },
-      focusBranchFirstNode({ event, prop, scope }) {
+      focusBranchFirstNode(params) {
+        const { event, prop, scope } = params
         const collection = prop("collection")
         const branchNode = collection.findNode(event.id)
         const firstNode = collection.getFirstNode(branchNode)
         const firstValue = collection.getNodeValue(firstNode)
-        dom.focusNode(scope, firstValue)
+        const scrolled = scrollToNode(params, firstValue)
+        if (scrolled) raf(() => dom.focusNode(scope, firstValue))
+        else dom.focusNode(scope, firstValue)
       },
       focusTreeNextNode(params) {
         const { event, prop, scope } = params
         const collection = prop("collection")
         const nextNode = collection.getNextNode(event.id, { skip: skipFn(params) })
         if (!nextNode) return
         const nextValue = collection.getNodeValue(nextNode)
-        dom.focusNode(scope, nextValue)
+        const scrolled = scrollToNode(params, nextValue)
+        if (scrolled) raf(() => dom.focusNode(scope, nextValue))
+        else dom.focusNode(scope, nextValue)
       },
       focusTreePrevNode(params) {
         const { event, prop, scope } = params
         const collection = prop("collection")
         const prevNode = collection.getPreviousNode(event.id, { skip: skipFn(params) })
         if (!prevNode) return
         const prevValue = collection.getNodeValue(prevNode)
-        dom.focusNode(scope, prevValue)
+        const scrolled = scrollToNode(params, prevValue)
+        if (scrolled) raf(() => dom.focusNode(scope, prevValue))
+        else dom.focusNode(scope, prevValue)
       },
-      focusBranchNode({ event, prop, scope }) {
+      focusBranchNode(params) {
+        const { event, prop, scope } = params
         const collection = prop("collection")
         const parentNode = collection.getParentNode(event.id)
         const parentValue = parentNode ? collection.getNodeValue(parentNode) : undefined
-        dom.focusNode(scope, parentValue)
+        if (!parentValue) return
+        const scrolled = scrollToNode(params, parentValue)
+        if (scrolled) raf(() => dom.focusNode(scope, parentValue))
+        else dom.focusNode(scope, parentValue)
       },
       selectAllNodes({ context, prop }) {
         context.set("selectedValue", prop("collection").getValues())
@@ -416,7 +433,10 @@ export const machine = createMachine<TreeViewSchema>({
           key: event.key,
         })
 
-        dom.focusNode(scope, node?.id)
+        if (!node?.id) return
+        const scrolled = scrollToNode(params, node.id)
+        if (scrolled) raf(() => dom.focusNode(scope, node.id))
+        else dom.focusNode(scope, node.id)
       },
       toggleNodeSelection({ context, event }) {
         const selectedValue = addOrRemove(context.get("selectedValue"), event.id)
@@ -641,3 +661,26 @@ export const machine = createMachine<TreeViewSchema>({
     },
   },
 })
+
+function scrollToNode(params: Params<TreeViewSchema>, value: string): boolean {
+  const { prop, scope, computed } = params
+  const scrollToIndexFn = prop("scrollToIndexFn")
+  if (!scrollToIndexFn) return false
+
+  const collection = prop("collection")
+  const visibleNodes = computed("visibleNodes")
+
+  for (let i = 0; i < visibleNodes.length; i++) {
+    const { node, indexPath } = visibleNodes[i]
+    if (collection.getNodeValue(node) !== value) continue
+    scrollToIndexFn({
+      index: i,
+      node,
+      indexPath,
+      getElement: () => scope.getById(dom.getNodeId(scope, value)),
+    })
+    return true
+  }
+
+  return false
+}
diff --git a/packages/machines/tree-view/src/tree-view.props.ts b/packages/machines/tree-view/src/tree-view.props.ts
@@ -30,6 +30,7 @@ export const props = createProps<TreeViewProps>()([
   "onRenameStart",
   "onBeforeRename",
   "onRenameComplete",
+  "scrollToIndexFn",
 ])
 
 export const splitProps = createSplitProps<Partial<TreeViewProps>>(props)
diff --git a/packages/machines/tree-view/src/tree-view.types.ts b/packages/machines/tree-view/src/tree-view.types.ts
