In TreeDataGrid, find the parent group row with a binary search and replace getOrInsertComputed

src/TreeDataGrid.tsx

@@ -1,4 +1,4 @@
-import { useCallback, useMemo } from 'react';
+import { useMemo } from 'react';
 import type { Key } from 'react';
 
 import { useLatestFunc } from './hooks';
@@ -39,7 +39,8 @@ export interface TreeDataGridProps<R, SR = unknown, K extends Key = Key> extends
 }
 
 interface PartialGroupRow<TRow> {
-  readonly groups: GroupsMap<TRow>;
+  // `undefined` for the groups on the last level, which only contain rows
+  readonly groups: GroupsMap<TRow> | undefined;
   // all the descendant rows of the group, not just the direct children
   readonly childRows: readonly TRow[];
   readonly startRowIndex: number;
@@ -49,7 +50,7 @@ type GroupsMap<TRow> = ReadonlyMap<string, PartialGroupRow<TRow>>;
 
 // only used while the rows are being grouped
 interface MutablePartialGroupRow<TRow> {
-  readonly groups: MutableGroupsMap<TRow>;
+  readonly groups: MutableGroupsMap<TRow> | undefined;
   readonly childRows: TRow[];
   startRowIndex: number;
 }
@@ -61,6 +62,12 @@ interface GroupedRows<TRow> {
   readonly rowsCount: number;
 }
 
+interface FlattenedRows<TRow> {
+  readonly rows: readonly (TRow | GroupRow<TRow>)[];
+  // indexes of the group rows in `rows`, in ascending order
+  readonly groupRowIndexes: readonly number[];
+}
+
 // Group rows are only ever created by `TreeDataGrid`, so every instance can share a single
 // set, which keeps `isGroupRow` stable and out of the dependency arrays below
 const allGroupRows = new WeakSet<object>();
@@ -138,19 +145,29 @@ export function TreeDataGrid<R, SR = unknown, K extends Key = Key>({
 
     // Group the rows in a single pass, each row is added to every group it belongs to
     const groups: MutableGroupsMap<R> = new Map();
+    const lastLevel = groupBy.length - 1;
 
     for (const row of rawRows) {
       let parentGroups = groups;
 
-      for (const columnKey of groupBy) {
-        const groupKey = getRowGroupKey(row, columnKey);
-        const group = parentGroups.getOrInsertComputed(groupKey, () => ({
-          groups: new Map(),
-          childRows: [],
-          startRowIndex: 0
-        }));
+      for (let level = 0; level <= lastLevel; level++) {
+        const groupKey = getRowGroupKey(row, groupBy[level]);
+        // `getOrInsertComputed` would be slower: V8 doesn't inline it like `get`,
+        // and its callback would be created for every row
+        let group = parentGroups.get(groupKey);
+        if (group === undefined) {
+          group = {
+            // the groups on the last level don't need a map, which saves one per group
+            groups: level === lastLevel ? undefined : new Map(),
+            childRows: [],
+            startRowIndex: 0
+          };
+          parentGroups.set(groupKey, group);
+        }
         group.childRows.push(row);
-        parentGroups = group.groups;
+        if (group.groups !== undefined) {
+          parentGroups = group.groups;
+        }
       }
     }
 
@@ -164,7 +181,7 @@ export function TreeDataGrid<R, SR = unknown, K extends Key = Key>({
         group.startRowIndex = startRowIndex + groupRowsCount;
         groupRowsCount++; // 1 for the group row
         groupRowsCount +=
-          group.groups.size === 0
+          group.groups === undefined
             ? // the group is on the last level, it only contains rows
               group.childRows.length
             : setStartRowIndexes(group.groups, startRowIndex + groupRowsCount);
@@ -178,12 +195,13 @@ export function TreeDataGrid<R, SR = unknown, K extends Key = Key>({
     return { groups, rowsCount };
   }, [getRowGroupKey, groupBy, rawRows]);
 
-  const rows = useMemo((): readonly (R | GroupRow<R>)[] => {
+  const { rows, groupRowIndexes } = useMemo((): FlattenedRows<R> => {
     if (groups === undefined) {
-      return rawRows;
+      return { rows: rawRows, groupRowIndexes: [] };
     }
 
     const flattenedRows: (R | GroupRow<R>)[] = [];
+    const groupRowIndexes: number[] = [];
 
     // Only the expanded groups are traversed
     function expandGroups(groups: GroupsMap<R>, parentId: string | undefined, level: number): void {
@@ -206,11 +224,12 @@ export function TreeDataGrid<R, SR = unknown, K extends Key = Key>({
           setSize: groups.size,
           startRowIndex
         };
+        groupRowIndexes.push(flattenedRows.length);
         flattenedRows.push(groupRow);
         allGroupRows.add(groupRow);
 
         if (isExpanded) {
-          if (childGroups.size === 0) {
+          if (childGroups === undefined) {
             // the group is on the last level, it only contains rows.
             // Push them one by one: spreading them would throw in groups with too many rows
             for (const row of childRows) {
@@ -225,7 +244,7 @@ export function TreeDataGrid<R, SR = unknown, K extends Key = Key>({
 
     expandGroups(groups, undefined, 0);
 
-    return flattenedRows;
+    return { rows: flattenedRows, groupRowIndexes };
   }, [expandedGroupIds, getGroupId, groups, rawRows]);
 
   const rowHeight = useMemo(() => {
@@ -241,42 +260,11 @@ export function TreeDataGrid<R, SR = unknown, K extends Key = Key>({
     return rawRowHeight;
   }, [rawRowHeight]);
 
-  const getParentRowAndIndex = useCallback(
-    (row: R | GroupRow<R>) => {
-      const rowIdx = rows.indexOf(row);
-      for (let i = rowIdx - 1; i >= 0; i--) {
-        const parentRow = rows[i];
-        if (isGroupRow(parentRow) && (!isGroupRow(row) || row.parentId === parentRow.id)) {
-          return [parentRow, i] as const;
-        }
-      }
-
-      return undefined;
-    },
-    [rows]
-  );
-
-  const rowKeyGetter = useCallback(
-    (row: R | GroupRow<R>) => {
-      if (isGroupRow(row)) {
-        return row.id;
-      }
-
-      if (typeof rawRowKeyGetter === 'function') {
-        return rawRowKeyGetter(row);
-      }
-
-      const parentRowAndIndex = getParentRowAndIndex(row);
-      if (parentRowAndIndex !== undefined) {
-        const { startRowIndex, childRows } = parentRowAndIndex[0];
-        const groupIndex = childRows.indexOf(row);
-        return startRowIndex + groupIndex + 1;
-      }
-
-      return rows.indexOf(row);
-    },
-    [getParentRowAndIndex, rawRowKeyGetter, rows]
-  );
+  // Without rawRowKeyGetter, the keys are computed in renderRow, where the row index is known
+  const rowKeyGetter = useMemo(() => {
+    if (typeof rawRowKeyGetter !== 'function') return undefined;
+    return (row: R | GroupRow<R>) => (isGroupRow(row) ? row.id : rawRowKeyGetter(row));
+  }, [rawRowKeyGetter]);
 
   const selectedRows = useMemo((): Maybe<ReadonlySet<Key>> => {
     if (rawSelectedRows == null) return null;
@@ -306,7 +294,7 @@ export function TreeDataGrid<R, SR = unknown, K extends Key = Key>({
 
     const newRawSelectedRows = new Set(rawSelectedRows);
     for (const row of rows) {
-      const key = rowKeyGetter(row);
+      const key = isGroupRow(row) ? row.id : rawRowKeyGetter(row);
       if (selectedRows?.has(key) && !newSelectedRows.has(key)) {
         if (isGroupRow(row)) {
           // unselect all children if the parent row is unselected
@@ -331,6 +319,36 @@ export function TreeDataGrid<R, SR = unknown, K extends Key = Key>({
     rawOnSelectedRowsChange(newRawSelectedRows);
   }
 
+  function getParentRowAndIndex(rowIdx: number) {
+    const row = rows[rowIdx];
+
+    // binary search for the position of the last group row before the row
+    let low = 0;
+    let high = groupRowIndexes.length - 1;
+    let position = -1;
+    while (low <= high) {
+      const middle = (low + high) >>> 1;
+      if (groupRowIndexes[middle] < rowIdx) {
+        position = middle;
+        low = middle + 1;
+      } else {
+        high = middle - 1;
+      }
+    }
+
+    // a row's parent is the closest group row before it,
+    // a group row's parent is the closest group row before it whose id is its parentId
+    for (let i = position; i >= 0; i--) {
+      const parentIdx = groupRowIndexes[i];
+      const parentRow = rows[parentIdx] as GroupRow<R>;
+      if (!isGroupRow(row) || row.parentId === parentRow.id) {
+        return [parentRow, parentIdx] as const;
+      }
+    }
+
+    return undefined;
+  }
+
   function handleKeyDown(args: CellKeyDownArgs<R, SR>, event: CellKeyboardEvent) {
     rawOnCellKeyDown?.(args, event);
     if (event.isGridDefaultPrevented()) return;
@@ -356,7 +374,7 @@ export function TreeDataGrid<R, SR = unknown, K extends Key = Key>({
 
     // If a group row is focused, and it is collapsed, move to the parent group row (if there is one).
     if (idx === -1 && event.key === leftKey && !row.isExpanded && row.level !== 0) {
-      const parentRowAndIndex = getParentRowAndIndex(row);
+      const parentRowAndIndex = getParentRowAndIndex(rowIdx);
       if (parentRowAndIndex !== undefined) {
         event.preventGridDefault();
         setActivePosition({ idx, rowIdx: parentRowAndIndex[1] }, { shouldFocus: true });
@@ -427,7 +445,7 @@ export function TreeDataGrid<R, SR = unknown, K extends Key = Key>({
       const { startRowIndex } = row;
       return (
         <GroupedRow
-          key={key}
+          key={row.id}
           {...rowProps}
           aria-rowindex={headerAndTopSummaryRowsCount + startRowIndex + 1}
           row={row}
@@ -437,15 +455,21 @@ export function TreeDataGrid<R, SR = unknown, K extends Key = Key>({
       );
     }
 
+    const { rowIdx } = rowProps;
+    let rowKey = key;
     let ariaRowIndex = rowProps['aria-rowindex'];
-    const parentRowAndIndex = getParentRowAndIndex(row);
+    const parentRowAndIndex = getParentRowAndIndex(rowIdx);
     if (parentRowAndIndex !== undefined) {
-      const { startRowIndex, childRows } = parentRowAndIndex[0];
-      const groupIndex = childRows.indexOf(row);
+      const [{ startRowIndex }, parentIdx] = parentRowAndIndex;
+      // the rows of a group directly follow its group row
+      const groupIndex = rowIdx - parentIdx - 1;
       ariaRowIndex = startRowIndex + headerAndTopSummaryRowsCount + groupIndex + 2;
+      if (rowKeyGetter === undefined) {
+        rowKey = startRowIndex + groupIndex + 1;
+      }
     }
 
-    return rawRenderRow(key, {
+    return rawRenderRow(rowKey, {
       ...rowProps,
       'aria-rowindex': ariaRowIndex,
       row,

test/browser/TreeDataGrid.test.tsx

@@ -249,6 +249,37 @@ test('should expand a group with more rows than a function call can take as argu
   await expect.element(page.getCell({ name: '0', exact: true })).toBeVisible();
 });
 
+test('should keep the row keys when an earlier group is expanded, without rowKeyGetter', async () => {
+  function Test() {
+    const [expandedGroupIds, setExpandedGroupIds] = useState(
+      (): ReadonlySet<unknown> => new Set(['Canada'])
+    );
+
+    return (
+      <TreeDataGrid
+        columns={[
+          { key: 'country', name: 'Country' },
+          { key: 'id', name: 'Id' }
+        ]}
+        rows={initialRows}
+        groupBy={['country']}
+        expandedGroupIds={expandedGroupIds}
+        onExpandedGroupIdsChange={setExpandedGroupIds}
+      />
+    );
+  }
+
+  await page.render(<Test />);
+  await testRowCount(4);
+  const idCell = page.getCell({ name: '3', exact: true });
+  const rowElement = getRowWithCell(idCell).element();
+
+  await userEvent.click(page.getCell({ name: 'USA' }));
+  await testRowCount(6);
+  // the row moved down, but its key didn't change, so React kept its element
+  expect(getRowWithCell(idCell).element()).toBe(rowElement);
+});
+
 test('should toggle group when group cell is clicked', async () => {
   await setup(['year']);
   await testRowCount(5);