diff --git a/src/TreeDataGrid.tsx b/src/TreeDataGrid.tsx index 78c0ee6f41..c24945f53c 100644 --- a/src/TreeDataGrid.tsx +++ b/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 extends } interface PartialGroupRow { - readonly groups: GroupsMap; + // `undefined` for the groups on the last level, which only contain rows + readonly groups: GroupsMap | 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 = ReadonlyMap>; // only used while the rows are being grouped interface MutablePartialGroupRow { - readonly groups: MutableGroupsMap; + readonly groups: MutableGroupsMap | undefined; readonly childRows: TRow[]; startRowIndex: number; } @@ -61,6 +62,12 @@ interface GroupedRows { readonly rowsCount: number; } +interface FlattenedRows { + readonly rows: readonly (TRow | GroupRow)[]; + // 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(); @@ -138,19 +145,29 @@ export function TreeDataGrid({ // Group the rows in a single pass, each row is added to every group it belongs to const groups: MutableGroupsMap = 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({ 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({ return { groups, rowsCount }; }, [getRowGroupKey, groupBy, rawRows]); - const rows = useMemo((): readonly (R | GroupRow)[] => { + const { rows, groupRowIndexes } = useMemo((): FlattenedRows => { if (groups === undefined) { - return rawRows; + return { rows: rawRows, groupRowIndexes: [] }; } const flattenedRows: (R | GroupRow)[] = []; + const groupRowIndexes: number[] = []; // Only the expanded groups are traversed function expandGroups(groups: GroupsMap, parentId: string | undefined, level: number): void { @@ -206,11 +224,12 @@ export function TreeDataGrid({ 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({ expandGroups(groups, undefined, 0); - return flattenedRows; + return { rows: flattenedRows, groupRowIndexes }; }, [expandedGroupIds, getGroupId, groups, rawRows]); const rowHeight = useMemo(() => { @@ -241,42 +260,11 @@ export function TreeDataGrid({ return rawRowHeight; }, [rawRowHeight]); - const getParentRowAndIndex = useCallback( - (row: R | GroupRow) => { - 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) => { - 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) => (isGroupRow(row) ? row.id : rawRowKeyGetter(row)); + }, [rawRowKeyGetter]); const selectedRows = useMemo((): Maybe> => { if (rawSelectedRows == null) return null; @@ -306,7 +294,7 @@ export function TreeDataGrid({ 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({ 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; + if (!isGroupRow(row) || row.parentId === parentRow.id) { + return [parentRow, parentIdx] as const; + } + } + + return undefined; + } + function handleKeyDown(args: CellKeyDownArgs, event: CellKeyboardEvent) { rawOnCellKeyDown?.(args, event); if (event.isGridDefaultPrevented()) return; @@ -356,7 +374,7 @@ export function TreeDataGrid({ // 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({ const { startRowIndex } = row; return ( ({ ); } + 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, diff --git a/test/browser/TreeDataGrid.test.tsx b/test/browser/TreeDataGrid.test.tsx index c736920823..efbaac9de0 100644 --- a/test/browser/TreeDataGrid.test.tsx +++ b/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 => new Set(['Canada']) + ); + + return ( + + ); + } + + await page.render(); + 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);