Refactor virtualizer persisted keys (#3394)

Author: devongovett

packages/@react-aria/grid/src/useGridCell.ts

@@ -10,9 +10,9 @@
  * governing permissions and limitations under the License.
  */
 
-import {DOMAttributes, FocusableElement, Node as RSNode} from '@react-types/shared';
+import {DOMAttributes, FocusableElement} from '@react-types/shared';
 import {focusSafely, getFocusableTreeWalker} from '@react-aria/focus';
-import {GridCollection} from '@react-types/grid';
+import {GridCollection, GridNode} from '@react-types/grid';
 import {gridMap} from './utils';
 import {GridState} from '@react-stately/grid';
 import {isFocusVisible} from '@react-aria/interactions';
@@ -23,7 +23,7 @@ import {useSelectableItem} from '@react-aria/selection';
 
 export interface GridCellProps {
   /** An object representing the grid cell. Contains all the relevant information that makes up the grid cell. */
-  node: RSNode<unknown>,
+  node: GridNode<unknown>,
   /** Whether the grid cell is contained in a virtual scroller. */
   isVirtualized?: boolean,
   /** Whether the cell or its first focusable child element should be focused when the grid cell is focused. */
@@ -230,7 +230,7 @@ export function useGridCell<T, C extends GridCollection<T>>(props: GridCellProps
   });
 
   if (isVirtualized) {
-    gridCellProps['aria-colindex'] = node.index + 1; // aria-colindex is 1-based
+    gridCellProps['aria-colindex'] = (node.colIndex ?? node.index) + 1; // aria-colindex is 1-based
   }
 
   return {

packages/@react-aria/virtualizer/src/Virtualizer.tsx

@@ -145,14 +145,8 @@ export function useVirtualizer<T extends object, V, W>(props: VirtualizerOptions
     lastFocusedKey.current = focusedKey;
   }, [focusedKey, virtualizer.visibleRect.height, virtualizer, lastFocusedKey, scrollToItem]);
 
-  // Tell the virtualizer to persist the focusedKey and prevent it from being reused.
-  // This is to prevent the FocusRing from setting focus to the ListView when a
-  // focused ListViewItem is scrolled out of the Virtualizer visible rectangle.
-  let focusedKeySet = useMemo((): Set<Key> => (
-    focusedKey ? new Set([focusedKey]) : new Set()
-  ), [focusedKey]);
-
-  virtualizer.persistedKeys = focusedKeySet;
+  // Persist the focusedKey and prevent it from being removed from the DOM when scrolled out of view.
+  virtualizer.persistedKeys = useMemo(() => focusedKey ? new Set([focusedKey]) : new Set(), [focusedKey]);
 
   let onFocus = useCallback((e: FocusEvent) => {
     // If the focused item is scrolled out of view and is not in the DOM, the collection

packages/@react-spectrum/table/test/Table.test.js

@@ -67,6 +67,11 @@ for (let i = 1; i <= 100; i++) {
   manyItems.push({id: i, foo: 'Foo ' + i, bar: 'Bar ' + i, baz: 'Baz ' + i});
 }
 
+let manyColumns = [];
+for (let i = 1; i <= 100; i++) {
+  manyColumns.push({id: i, name: 'Column ' + i});
+}
+
 function ExampleSortTable() {
   let [sortDescriptor, setSortDescriptor] = React.useState({column: 'bar', direction: 'ascending'});
 
@@ -816,6 +821,23 @@ describe('TableView', function () {
       </TableView>
     );
 
+    let renderManyColumns = () => render(
+      <TableView aria-label="Table" selectionMode="multiple">
+        <TableHeader columns={manyColumns}>
+          {column =>
+            <Column>{column.name}</Column>
+          }
+        </TableHeader>
+        <TableBody items={manyItems}>
+          {item =>
+            (<Row key={item.foo}>
+              {key => <Cell>{item.foo + ' ' + key}</Cell>}
+            </Row>)
+          }
+        </TableBody>
+      </TableView>
+    );
+
     let focusCell = (tree, text) => act(() => getCell(tree, text).focus());
     let moveFocus = (key, opts = {}) => {fireEvent.keyDown(document.activeElement, {key, ...opts});};
 
@@ -1544,27 +1566,38 @@ describe('TableView', function () {
       });
 
       it('should scroll to a cell when it is focused off screen', function () {
-        let tree = renderMany();
+        let tree = renderManyColumns();
         let body = tree.getByRole('grid').childNodes[1];
 
-        let cell = getCell(tree, 'Baz 5');
+        let cell = getCell(tree, 'Foo 5 5');
         act(() => cell.focus());
         expect(document.activeElement).toBe(cell);
         expect(body.scrollTop).toBe(0);
 
         // When scrolling the focused item out of view, focus should remaind on the item,
         // virtualizer keeps focused items from being reused
         body.scrollTop = 1000;
+        body.scrollLeft = 1000;
         fireEvent.scroll(body);
 
         expect(body.scrollTop).toBe(1000);
         expect(document.activeElement).toBe(cell);
-        expect(tree.queryByText('Baz 5')).toBe(cell.firstElementChild);
+        expect(tree.queryByText('Foo 5 5')).toBe(cell.firstElementChild);
+
+        // Ensure we have the correct sticky cells in the right order.
+        let row = cell.closest('[role=row]');
+        let cells = within(row).getAllByRole(role => role === 'gridcell' || role === 'rowheader');
+        expect(cells).toHaveLength(18);
+        expect(cells[0]).toHaveAttribute('aria-colindex', '1'); // checkbox
+        expect(cells[1]).toHaveAttribute('aria-colindex', '2'); // rowheader
+        expect(cells[2]).toHaveAttribute('aria-colindex', '6'); // persisted
+        expect(cells[2]).toBe(cell);
+        expect(cells[3]).toHaveAttribute('aria-colindex', '14'); // first visible
 
         // Moving focus should scroll the new focused item into view
         moveFocus('ArrowLeft');
         expect(body.scrollTop).toBe(164);
-        expect(document.activeElement).toBe(getCell(tree, 'Bar 5'));
+        expect(document.activeElement).toBe(getCell(tree, 'Foo 5 4'));
       });
 
       it('should not scroll when a column header receives focus', function () {

packages/@react-stately/layout/src/ListLayout.ts

@@ -109,12 +109,6 @@ export class ListLayout<T> extends Layout<Node<T>> implements KeyboardDelegate,
             res.push(node.header);
           }
 
-          if (node.children) {
-            addNodes(node.children);
-          }
-        } else if (this.virtualizer.isPersistedKey(node)) {
-          res.push(node.layoutInfo);
-
           if (node.children) {
             addNodes(node.children);
           }
@@ -127,7 +121,7 @@ export class ListLayout<T> extends Layout<Node<T>> implements KeyboardDelegate,
   }
 
   isVisible(node: LayoutNode, rect: Rect) {
-    return node.layoutInfo.rect.intersects(rect) || node.layoutInfo.isSticky;
+    return node.layoutInfo.rect.intersects(rect) || node.layoutInfo.isSticky || this.virtualizer.isPersistedKey(node.layoutInfo.key);
   }
 
   validate(invalidationContext: InvalidationContext<Node<T>, unknown>) {

packages/@react-stately/layout/src/TableLayout.ts

@@ -24,6 +24,8 @@ export class TableLayout<T> extends ListLayout<T> {
   stickyColumnIndices: number[];
   wasLoading = false;
   isLoading = false;
+  lastPersistedKeys: Set<Key> = null;
+  persistedIndices: Map<Key, number[]> = new Map();
 
   constructor(options: ListLayoutOptions<T>) {
     super(options);
@@ -50,6 +52,7 @@ export class TableLayout<T> extends ListLayout<T> {
     let header = this.buildHeader();
     let body = this.buildBody(0);
     this.stickyColumnIndices = this.collection.columns.filter(c => c.props.isSelectionCell || this.collection.rowHeaderColumnKeys.has(c.key)).map(c => c.index);
+    this.lastPersistedKeys = null;
 
     body.layoutInfo.rect.width = Math.max(header.layoutInfo.rect.width, body.layoutInfo.rect.width);
     this.contentSize = new Size(body.layoutInfo.rect.width, body.layoutInfo.rect.maxY);
@@ -125,8 +128,9 @@ export class TableLayout<T> extends ListLayout<T> {
   // used to get the column widths when rendering to the DOM
   getColumnWidth_(node: GridNode<T>) {
     let colspan = node.colspan ?? 1;
+    let colIndex = node.colIndex ?? node.index;
     let width = 0;
-    for (let i = node.index; i < node.index + colspan; i++) {
+    for (let i = colIndex; i < colIndex + colspan; i++) {
       let column = this.collection.columns[i];
       width += this.getColumnWidth(column.key);
     }
@@ -274,6 +278,7 @@ export class TableLayout<T> extends ListLayout<T> {
   getVisibleLayoutInfos(rect: Rect) {
     let res: LayoutInfo[] = [];
 
+    this.buildPersistedIndices();
     for (let node of this.rootNodes) {
       res.push(node.layoutInfo);
       this.addVisibleLayoutInfos(res, node, rect);
@@ -298,24 +303,35 @@ export class TableLayout<T> extends ListLayout<T> {
       case 'rowgroup': {
         let firstVisibleRow = this.binarySearch(node.children, rect.topLeft, 'y');
         let lastVisibleRow = this.binarySearch(node.children, rect.bottomRight, 'y');
-        // Check to see if a persisted key exists before the visible rows
-        // This is for keeping focus on a row that scrolls out of view
-        for (let h = 0; h < firstVisibleRow; h++) {
-          if (this.virtualizer.isPersistedKey(node.children[h])) {
-            res.push(node.children[h].layoutInfo);
-            this.addVisibleLayoutInfos(res, node.children[h], rect);
-          }
+
+        // Add persisted rows before the visible rows.
+        let persistedRowIndices = this.persistedIndices.get(node.layoutInfo.key);
+        let persistIndex = 0;
+        while (
+          persistedRowIndices &&
+          persistIndex < persistedRowIndices.length &&
+          persistedRowIndices[persistIndex] < firstVisibleRow
+        ) {
+          let idx = persistedRowIndices[persistIndex];
+          res.push(node.children[idx].layoutInfo);
+          this.addVisibleLayoutInfos(res, node.children[idx], rect);
+          persistIndex++;
         }
+
         for (let i = firstVisibleRow; i <= lastVisibleRow; i++) {
+          // Skip persisted rows that overlap with visible cells.
+          while (persistedRowIndices && persistIndex < persistedRowIndices.length && persistedRowIndices[persistIndex] < i) {
+            persistIndex++;
+          }
+
           res.push(node.children[i].layoutInfo);
           this.addVisibleLayoutInfos(res, node.children[i], rect);
         }
-        // Check to see if a persisted key exists after the visible rows
-        for (let j = lastVisibleRow + 1; j < node.children.length; j++) {
-          if (this.virtualizer.isPersistedKey(node.children[j])) {
-            res.push(node.children[j].layoutInfo);
-            this.addVisibleLayoutInfos(res, node.children[j], rect);
-          }
+
+        // Add persisted rows after the visible rows.
+        while (persistedRowIndices && persistIndex < persistedRowIndices.length) {
+          let idx = persistedRowIndices[persistIndex++];
+          res.push(node.children[idx].layoutInfo);
         }
         break;
       }
@@ -324,35 +340,27 @@ export class TableLayout<T> extends ListLayout<T> {
         let firstVisibleCell = this.binarySearch(node.children, rect.topLeft, 'x');
         let lastVisibleCell = this.binarySearch(node.children, rect.topRight, 'x');
         let stickyIndex = 0;
-        // Check to see if a persisted key exists before the visible cells
-        // This is for keeping focus on a cell that scrolls out of view
-        for (let h = 0; h < firstVisibleCell; h++) {
-          if (this.virtualizer.isPersistedKey(node.children[h])) {
-            res.push(node.children[h].layoutInfo);
-          }
+
+        // Add persisted/sticky cells before the visible cells.
+        let persistedCellIndices = this.persistedIndices.get(node.layoutInfo.key) || this.stickyColumnIndices;
+        while (stickyIndex < persistedCellIndices.length && persistedCellIndices[stickyIndex] < firstVisibleCell) {
+          let idx = persistedCellIndices[stickyIndex];
+          res.push(node.children[idx].layoutInfo);
+          stickyIndex++;
         }
+
         for (let i = firstVisibleCell; i <= lastVisibleCell; i++) {
-          // Sticky columns and row headers are always in the DOM. Interleave these
-          // with the visible range so that they are in the right order.
-          if (stickyIndex < this.stickyColumnIndices.length) {
-            let idx = this.stickyColumnIndices[stickyIndex];
-            while (idx < i) {
-              res.push(node.children[idx].layoutInfo);
-              idx = this.stickyColumnIndices[stickyIndex++];
-            }
+          // Skip sticky cells that overlap with visible cells.
+          while (stickyIndex < persistedCellIndices.length && persistedCellIndices[stickyIndex] < i) {
+            stickyIndex++;
           }
 
           res.push(node.children[i].layoutInfo);
         }
-        // Check to see if a persisted key exists after the visible cells
-        for (let j = lastVisibleCell; j < node.children.length; j++) {
-          if (this.virtualizer.isPersistedKey(node.children[j])) {
-            res.push(node.children[j].layoutInfo);
-          }
-        }
 
-        while (stickyIndex < this.stickyColumnIndices.length) {
-          let idx = this.stickyColumnIndices[stickyIndex++];
+        // Add any remaining sticky cells after the visible cells.
+        while (stickyIndex < persistedCellIndices.length) {
+          let idx = persistedCellIndices[stickyIndex++];
           res.push(node.children[idx].layoutInfo);
         }
         break;
@@ -381,6 +389,46 @@ export class TableLayout<T> extends ListLayout<T> {
     return Math.max(0, Math.min(items.length - 1, low));
   }
 
+  buildPersistedIndices() {
+    if (this.virtualizer.persistedKeys === this.lastPersistedKeys) {
+      return;
+    }
+
+    this.lastPersistedKeys = this.virtualizer.persistedKeys;
+    this.persistedIndices.clear();
+
+    // Build a map of parentKey => indices of children to persist.
+    for (let key of this.virtualizer.persistedKeys) {
+      let layoutInfo = this.layoutInfos.get(key);
+
+      // Walk up ancestors so parents are also persisted if children are.
+      while (layoutInfo && layoutInfo.parentKey) {
+        let collectionNode = this.collection.getItem(layoutInfo.key);
+        let indices = this.persistedIndices.get(layoutInfo.parentKey);
+        if (!indices) {
+          // stickyColumnIndices are always persisted along with any cells from persistedKeys.
+          indices = collectionNode.type === 'cell' ? [...this.stickyColumnIndices] : [];
+          this.persistedIndices.set(layoutInfo.parentKey, indices);
+        }
+
+        let index = collectionNode.index;
+        if (layoutInfo.parentKey === 'body') {
+          index -= this.collection.headerRows.length;
+        }
+
+        if (!indices.includes(index)) {
+          indices.push(index);
+        }
+
+        layoutInfo = this.layoutInfos.get(layoutInfo.parentKey);
+      }
+    }
+
+    for (let indices of this.persistedIndices.values()) {
+      indices.sort((a, b) => a - b);
+    }
+  }
+
   getInitialLayoutInfo(layoutInfo: LayoutInfo) {
     let res = super.getInitialLayoutInfo(layoutInfo);
 

packages/@react-stately/table/src/TableCollection.ts

@@ -20,7 +20,7 @@ interface GridCollectionOptions {
 const ROW_HEADER_COLUMN_KEY = 'row-header-column-' + Math.random().toString(36).slice(2);
 
 function buildHeaderRows<T>(keyMap: Map<Key, GridNode<T>>, columnNodes: GridNode<T>[]): GridNode<T>[] {
-  let columns = [];
+  let columns: GridNode<T>[][] = [];
   let seen = new Map();
   for (let column of columnNodes) {
     let parentKey = column.parentKey;
@@ -104,7 +104,7 @@ function buildHeaderRows<T>(keyMap: Map<Key, GridNode<T>>, columnNodes: GridNode
         }
 
         item.level = i;
-        item.index = colIndex;
+        item.colIndex = colIndex;
         row.push(item);
       }