feat: support variable node height

modules/react-arborist/src/components/cursor.tsx

@@ -7,7 +7,7 @@ export function Cursor() {
   if (!cursor || cursor.type !== "line") return null;
   const indent = tree.indent;
   const top =
-    tree.rowHeight * cursor.index +
+    tree.rowTopPosition(cursor.index) +
     (tree.props.padding ?? tree.props.paddingTop ?? 0);
   const left = indent * cursor.level;
   const Cursor = tree.renderCursor;

modules/react-arborist/src/components/default-container.tsx

@@ -1,4 +1,4 @@
-import { FixedSizeList } from "react-window";
+import { VariableSizeList } from "react-window";
 import { useDataUpdates, useTreeApi } from "../context";
 import { focusNextElement, focusPrevElement } from "../utils";
 import { ListOuterElement } from "./list-outer-element";
@@ -217,7 +217,7 @@ export function DefaultContainer() {
       }}
     >
       {/* @ts-ignore */}
-      <FixedSizeList
+      <VariableSizeList
         className={tree.props.className}
         outerRef={tree.listEl}
         itemCount={tree.visibleNodes.length}
@@ -233,7 +233,7 @@ export function DefaultContainer() {
         ref={tree.list}
       >
         {RowContainer}
-      </FixedSizeList>
+      </VariableSizeList>
     </div>
   );
 }

modules/react-arborist/src/components/list-outer-element.tsx

@@ -5,7 +5,7 @@ import { Cursor } from "./cursor";
 
 export const ListOuterElement = forwardRef(function Outer(
   props: React.HTMLProps<HTMLDivElement>,
-  ref
+  ref,
 ) {
   const { children, ...rest } = props;
   const tree = useTreeApi();
@@ -29,7 +29,9 @@ const DropContainer = () => {
   return (
     <div
       style={{
-        height: tree.visibleNodes.length * tree.rowHeight,
+        height:
+          tree.rowTopPosition(tree.visibleNodes.length - 1) +
+          tree.rowHeight(tree.visibleNodes.length - 1),
         width: "100%",
         position: "absolute",
         left: "0",

modules/react-arborist/src/components/provider.tsx

@@ -6,7 +6,7 @@ import {
   useRef,
 } from "react";
 import { useSyncExternalStore } from "use-sync-external-store/shim";
-import { FixedSizeList } from "react-window";
+import { VariableSizeList } from "react-window";
 import {
   DataUpdatesContext,
   DndContext,
@@ -35,7 +35,7 @@ export function TreeProvider<T>({
   imperativeHandle,
   children,
 }: Props<T>) {
-  const list = useRef<FixedSizeList | null>(null);
+  const list = useRef<VariableSizeList | null>(null);
   const listEl = useRef<HTMLDivElement | null>(null);
   const store = useRef<Store<RootState, Actions>>(
     // @ts-ignore

modules/react-arborist/src/dnd/drag-hook.ts

@@ -37,6 +37,7 @@ export function useDragHook<T>(node: NodeApi<T>): ConnectDragSource {
           });
           tree.open(parentId);
         }
+        tree.list?.current?.resetAfterIndex(0);
         tree.dispatch(dnd.dragEnd());
       },
     }),

modules/react-arborist/src/interfaces/tree-api.ts

@@ -2,7 +2,11 @@ import { EditResult } from "../types/handlers";
 import { Identity, IdObj } from "../types/utils";
 import { TreeProps } from "../types/tree-props";
 import { MutableRefObject } from "react";
-import { Align, FixedSizeList, ListOnItemsRenderedProps } from "react-window";
+import {
+  Align,
+  ListOnItemsRenderedProps,
+  VariableSizeList,
+} from "react-window";
 import * as utils from "../utils";
 import { DefaultCursor } from "../components/default-cursor";
 import { DefaultRow } from "../components/default-row";
@@ -34,8 +38,8 @@ export class TreeApi<T> {
   constructor(
     public store: Store<RootState, Actions>,
     public props: TreeProps<T>,
-    public list: MutableRefObject<FixedSizeList | null>,
-    public listEl: MutableRefObject<HTMLDivElement | null>
+    public list: MutableRefObject<VariableSizeList | null>,
+    public listEl: MutableRefObject<HTMLDivElement | null>,
   ) {
     /* Changes here must also be made in update() */
     this.root = createRoot<T>(this);
@@ -79,10 +83,6 @@ export class TreeApi<T> {
     return this.props.indent ?? 24;
   }
 
-  get rowHeight() {
-    return this.props.rowHeight ?? 24;
-  }
-
   get overscanCount() {
     return this.props.overscanCount ?? 1;
   }
@@ -91,6 +91,33 @@ export class TreeApi<T> {
     return (this.props.searchTerm || "").trim();
   }
 
+  rowHeight = (index: number): number => {
+    if (!this.props.rowHeight) {
+      return 24;
+    }
+
+    const node = this.at(index);
+    if (!node) {
+      return 0;
+    }
+
+    return typeof this.props.rowHeight === "function"
+      ? this.props.rowHeight(node)
+      : this.props.rowHeight;
+  };
+
+  rowTopPosition = (index: number): number => {
+    let position = 0;
+    for (let i = 0; i < index; i++) {
+      position += this.rowHeight(i);
+    }
+    return position;
+  };
+
+  redrawList = (afterIndex?: number | undefined | null) => {
+    this.list.current?.resetAfterIndex(afterIndex ?? 0);
+  };
+
   get matchFn() {
     const match =
       this.props.searchMatch ??
@@ -194,7 +221,7 @@ export class TreeApi<T> {
       type?: "internal" | "leaf";
       parentId?: null | string;
       index?: null | number;
-    } = {}
+    } = {},
   ) {
     const parentId =
       opts.parentId === undefined
@@ -224,6 +251,9 @@ export class TreeApi<T> {
     const idents = Array.isArray(node) ? node : [node];
     const ids = idents.map(identify);
     const nodes = ids.map((id) => this.get(id)!).filter((n) => !!n);
+
+    this.redrawList(Math.min(...nodes.map((node) => node.rowIndex ?? 0)));
+
     await safeRun(this.props.onDelete, { nodes, ids });
   }
 
@@ -232,6 +262,7 @@ export class TreeApi<T> {
     this.resolveEdit({ cancelled: true });
     this.scrollTo(id);
     this.dispatch(edit(id));
+    this.redrawList(this.get(id)?.rowIndex);
     return new Promise((resolve) => {
       TreeApi.editPromise = resolve;
     });
@@ -247,12 +278,14 @@ export class TreeApi<T> {
     });
     this.dispatch(edit(null));
     this.resolveEdit({ cancelled: false, value });
+    this.redrawList(this.get(id)?.rowIndex);
     setTimeout(() => this.onFocus()); // Return focus to element;
   }
 
   reset() {
     this.dispatch(edit(null));
     this.resolveEdit({ cancelled: true });
+    this.redrawList();
     setTimeout(() => this.onFocus()); // Return focus to element;
   }
 
@@ -470,19 +503,21 @@ export class TreeApi<T> {
 
   /* Visibility */
 
-  open(identity: Identity) {
+  open(identity: Identity, redraw: boolean = true) {
     const id = identifyNull(identity);
     if (!id) return;
     if (this.isOpen(id)) return;
     this.dispatch(visibility.open(id, this.isFiltered));
+    redraw && this.redrawList(this.get(id)?.rowIndex);
     safeRun(this.props.onToggle, id);
   }
 
-  close(identity: Identity) {
+  close(identity: Identity, redraw: boolean = true) {
     const id = identifyNull(identity);
     if (!id) return;
     if (!this.isOpen(id)) return;
     this.dispatch(visibility.close(id, this.isFiltered));
+    redraw && this.redrawList(this.get(id)?.rowIndex);
     safeRun(this.props.onToggle, id);
   }
 
@@ -499,9 +534,10 @@ export class TreeApi<T> {
     let parent = node?.parent;
 
     while (parent) {
-      this.open(parent.id);
+      this.open(parent.id, false);
       parent = parent.parent;
     }
+    this.redrawList();
   }
 
   openSiblings(node: NodeApi<T>) {
@@ -512,23 +548,26 @@ export class TreeApi<T> {
       const isOpen = node.isOpen;
       for (let sibling of parent.children) {
         if (sibling.isInternal) {
-          isOpen ? this.close(sibling.id) : this.open(sibling.id);
+          isOpen ? this.close(sibling.id, false) : this.open(sibling.id, false);
         }
       }
+      this.redrawList();
       this.scrollTo(this.focusedNode);
     }
   }
 
   openAll() {
     utils.walk(this.root, (node) => {
-      if (node.isInternal) node.open();
+      if (node.isInternal) this.open(node, false);
     });
+    this.redrawList();
   }
 
   closeAll() {
     utils.walk(this.root, (node) => {
-      if (node.isInternal) node.close();
+      if (node.isInternal) this.close(node, false);
     });
+    this.redrawList();
   }
 
   /* Scrolling */

modules/react-arborist/src/types/tree-props.ts

@@ -7,6 +7,27 @@ import { NodeApi } from "../interfaces/node-api";
 import { OpenMap } from "../state/open-slice";
 import { useDragDropManager } from "react-dnd";
 
+export type RowHeightCalculatorParams<T> = Pick<
+  NodeApi<T>,
+  | "childIndex"
+  | "children"
+  | "data"
+  | "parent"
+  | "id"
+  | "rowIndex"
+  | "tree"
+  | "isRoot"
+  | "isLeaf"
+  | "isClosed"
+  | "isOpen"
+  | "isAncestorOf"
+  | "level"
+>;
+
+export type RowHeightCalculator<T> = (
+  params: RowHeightCalculatorParams<T>,
+) => number;
+
 export interface TreeProps<T> {
   /* Data Options */
   data?: readonly T[];
@@ -26,7 +47,7 @@ export interface TreeProps<T> {
   renderContainer?: ElementType<{}>;
 
   /* Sizes */
-  rowHeight?: number;
+  rowHeight?: number | RowHeightCalculator<T>;
   overscanCount?: number;
   width?: number | string;
   height?: number;