perf(mosaic): re-render only the table rows whose selection changes

.changeset/mosaic-table-row-renders.md

@@ -0,0 +1,5 @@
+---
+'@clerk/mosaic': patch
+---
+
+Improve the performance of selecting rows in the `OrganizationProfile` member tables and the API keys table.

packages/mosaic/src/primitives/table/table-body.tsx

@@ -2,7 +2,7 @@
 
 import React from 'react';
 
-import { type ComponentProps, mergeProps, useRender } from '../utils';
+import { type ComponentProps, mergeProps, useRender, useStore } from '../utils';
 import { type TableItem, TableItemContext, useTableContext } from './table-context';
 
 export interface TableBodyProps<Item> extends Omit<ComponentProps<'tbody'>, 'children'> {
@@ -31,7 +31,8 @@ export const TableBody = React.forwardRef(TableBodyImpl) as TableBodyComponent;
 
 function TableBodyImpl<Item>(props: TableBodyProps<Item>, ref: React.ForwardedRef<HTMLTableSectionElement>) {
   const { render, items, itemToStringValue, itemToChildren, children, renderEmptyState, ...otherProps } = props;
-  const { expanded, syncBodyRows, columnCount, excludeFromNavigation } = useTableContext();
+  const { store, syncBodyRows, columnCount, excludeFromNavigation } = useTableContext();
+  const expanded = useStore(store, state => state.expanded);
   const tree = itemToChildren !== undefined;
 
   const rows = React.useMemo(

packages/mosaic/src/primitives/table/table-context.ts

@@ -1,5 +1,7 @@
 import { createContext, useContext, useId, useLayoutEffect, useRef, useState } from 'react';
 
+import type { Store } from '../utils';
+
 export type TableSelectionMode = 'none' | 'single' | 'multiple';
 
 export interface TableSort {
@@ -12,15 +14,19 @@ export interface TableBodyRows {
   tree: boolean;
 }
 
+export interface TableState {
+  selection: string[];
+  expanded: string[];
+  selectable: string[];
+}
+
+export type TableSelectAll = 'empty' | 'none' | 'some' | 'all';
+
 export interface TableContextValue {
   selectionMode: TableSelectionMode;
-  selection: string[];
+  store: Store<TableState>;
   sort: TableSort | null;
   tree: boolean;
-  allSelected: boolean;
-  someSelected: boolean;
-  canSelectAll: boolean;
-  expanded: string[];
   toggleSelected: (value: string, options?: { range?: boolean }) => void;
   toggleAllSelected: () => void;
   toggleExpanded: (value: string) => void;
@@ -71,6 +77,17 @@ export function useTableContext() {
   return ctx;
 }
 
+export function selectAllState({ selection, selectable }: TableState): TableSelectAll {
+  if (selectable.length === 0) {
+    return 'empty';
+  }
+  const selectedCount = selectable.filter(value => selection.includes(value)).length;
+  if (selectedCount === 0) {
+    return 'none';
+  }
+  return selectedCount === selectable.length ? 'all' : 'some';
+}
+
 export function useTableItem() {
   const item = useContext(TableItemContext);
   if (!item) {

packages/mosaic/src/primitives/table/table-root.tsx

@@ -3,13 +3,15 @@
 import React from 'react';
 
 import { useControllableState } from '../hooks/use-controllable-state';
-import { type ComponentProps, type DefaultProps, mergeProps, useRender } from '../utils';
+import { type ComponentProps, createStore, type DefaultProps, mergeProps, useRender } from '../utils';
 import {
+  selectAllState,
   type TableBodyRows,
   TableContext,
   type TableContextValue,
   type TableSelectionMode,
   type TableSort,
+  type TableState,
 } from './table-context';
 import { isTextEntry } from './table-dom';
 import { useGridNavigation } from './use-grid-navigation';
@@ -71,43 +73,61 @@ export const TableRoot = React.forwardRef<HTMLTableElement, TableRootProps>(func
     () => (selectionMode === 'none' ? [] : bodyRows.values.filter(value => !disabledValues.includes(value))),
     [selectionMode, bodyRows, disabledValues],
   );
-  const selectedCount = selectable.filter(value => selection.includes(value)).length;
-  const allSelected = selectable.length > 0 && selectedCount === selectable.length;
-  const someSelected = selectedCount > 0 && !allSelected;
+  const [store] = React.useState(() => createStore<TableState>({ selection, expanded, selectable }));
   const interactive = tree || selectionMode !== 'none' || interactiveCount > 0;
 
+  React.useLayoutEffect(() => {
+    store.set({ selection, expanded, selectable });
+  }, [store, selection, expanded, selectable]);
+
   const leaveCell = useGridNavigation(element, interactive, excluded);
 
-  const toggleSelected = (value: string, options?: { range?: boolean }) => {
-    if (!selectable.includes(value)) {
-      return;
-    }
-    const selected = !selection.includes(value);
-    if (selectionMode === 'single') {
-      setSelection(selected ? [value] : []);
-      return;
-    }
-    const anchor = options?.range && anchorRef.current !== null ? selectable.indexOf(anchorRef.current) : -1;
-    anchorRef.current = value;
-    const index = selectable.indexOf(value);
-    const changed = anchor === -1 ? [value] : selectable.slice(Math.min(anchor, index), Math.max(anchor, index) + 1);
-    setSelection(selected ? union(selection, changed) : selection.filter(other => !changed.includes(other)));
-  };
-
-  const toggleAllSelected = () => {
-    setSelection(allSelected ? selection.filter(value => !selectable.includes(value)) : union(selection, selectable));
-  };
-
-  const toggleExpanded = (value: string) => {
-    setExpanded(expanded.includes(value) ? expanded.filter(other => other !== value) : [...expanded, value]);
-  };
-
-  const sortBy = (column: string) => {
-    const direction = sort?.column === column && sort.direction === 'ascending' ? 'descending' : 'ascending';
-    const next: TableSort = { column, direction };
-    setSort(next);
-    onSortChange?.(next);
-  };
+  const toggleSelected = React.useCallback(
+    (value: string, options?: { range?: boolean }) => {
+      const { selection: current, selectable: rows } = store.get();
+      if (!rows.includes(value)) {
+        return;
+      }
+      const selected = !current.includes(value);
+      if (selectionMode === 'single') {
+        setSelection(selected ? [value] : []);
+        return;
+      }
+      const anchor = options?.range && anchorRef.current !== null ? rows.indexOf(anchorRef.current) : -1;
+      anchorRef.current = value;
+      const index = rows.indexOf(value);
+      const changed = anchor === -1 ? [value] : rows.slice(Math.min(anchor, index), Math.max(anchor, index) + 1);
+      setSelection(selected ? union(current, changed) : current.filter(other => !changed.includes(other)));
+    },
+    [store, selectionMode, setSelection],
+  );
+
+  const toggleAllSelected = React.useCallback(() => {
+    const state = store.get();
+    setSelection(
+      selectAllState(state) === 'all'
+        ? state.selection.filter(value => !state.selectable.includes(value))
+        : union(state.selection, state.selectable),
+    );
+  }, [store, setSelection]);
+
+  const toggleExpanded = React.useCallback(
+    (value: string) => {
+      const current = store.get().expanded;
+      setExpanded(current.includes(value) ? current.filter(other => other !== value) : [...current, value]);
+    },
+    [store, setExpanded],
+  );
+
+  const sortBy = React.useCallback(
+    (column: string) => {
+      const direction = sort?.column === column && sort.direction === 'ascending' ? 'descending' : 'ascending';
+      const next: TableSort = { column, direction };
+      setSort(next);
+      onSortChange?.(next);
+    },
+    [sort, setSort, onSortChange],
+  );
 
   const syncBodyRows = React.useCallback((next: TableBodyRows) => {
     setBodyRows(previous =>
@@ -117,28 +137,44 @@ export const TableRoot = React.forwardRef<HTMLTableElement, TableRootProps>(func
 
   const excludeFromNavigation = React.useCallback((node: Element) => excluded.add(node), [excluded]);
 
-  const contextValue: TableContextValue = {
-    selectionMode,
-    selection,
-    expanded,
-    sort,
-    tree,
-    allSelected,
-    someSelected,
-    canSelectAll: selectable.length > 0,
-    toggleSelected,
-    toggleAllSelected,
-    toggleExpanded,
-    sortBy,
-    syncBodyRows,
-    registerDisabled,
-    registerInteractive,
-    registerColumn,
-    rowHeaderColumns,
-    registerRowHeaderColumn,
-    columnCount,
-    excludeFromNavigation,
-  };
+  const contextValue = React.useMemo<TableContextValue>(
+    () => ({
+      selectionMode,
+      store,
+      sort,
+      tree,
+      toggleSelected,
+      toggleAllSelected,
+      toggleExpanded,
+      sortBy,
+      syncBodyRows,
+      registerDisabled,
+      registerInteractive,
+      registerColumn,
+      rowHeaderColumns,
+      registerRowHeaderColumn,
+      columnCount,
+      excludeFromNavigation,
+    }),
+    [
+      selectionMode,
+      store,
+      sort,
+      tree,
+      toggleSelected,
+      toggleAllSelected,
+      toggleExpanded,
+      sortBy,
+      syncBodyRows,
+      registerDisabled,
+      registerInteractive,
+      registerColumn,
+      rowHeaderColumns,
+      registerRowHeaderColumn,
+      columnCount,
+      excludeFromNavigation,
+    ],
+  );
 
   const defaultProps = {
     role: tree ? 'treegrid' : interactive ? 'grid' : undefined,
@@ -162,7 +198,7 @@ export const TableRoot = React.forwardRef<HTMLTableElement, TableRootProps>(func
         !isTextEntry(target)
       ) {
         event.preventDefault();
-        if (!allSelected) {
+        if (selectAllState({ selection, expanded, selectable }) !== 'all') {
           toggleAllSelected();
         }
       }

packages/mosaic/src/primitives/table/table-row.tsx

@@ -2,7 +2,7 @@
 
 import React, { type CSSProperties } from 'react';
 
-import { type ComponentProps, type DefaultProps, mergeProps, useRender } from '../utils';
+import { type ComponentProps, type DefaultProps, mergeProps, useRender, useStore } from '../utils';
 import { TableRowContext, type TableRowContextValue, useTableContext, useTableItem } from './table-context';
 import { focusIsInChildRows, isOwnClick, isRtl, openLink, type TableLink } from './table-dom';
 
@@ -18,15 +18,15 @@ export const TableRow = React.forwardRef<HTMLTableRowElement, TableRowProps>(fun
   const table = useTableContext();
   const item = useTableItem();
   const { value } = item;
-  const { registerDisabled, registerInteractive } = table;
+  const { store, registerDisabled, registerInteractive } = table;
   const elementRef = React.useRef<HTMLTableRowElement | null>(null);
 
   const actionable = !disabled && (onClick !== undefined || href !== undefined);
   const canSelect = !disabled && table.selectionMode !== 'none';
-  const selected = table.selection.includes(value);
+  const selected = useStore(store, state => state.selection.includes(value));
   const canExpand = table.tree && (item.hasChildren || expandable);
-  const expanded = canExpand && table.expanded.includes(value);
-  const hasSelection = table.selection.length > 0;
+  const isExpanded = useStore(store, state => state.expanded.includes(value));
+  const expanded = canExpand && isExpanded;
   const [rowHeaders, setRowHeaders] = React.useState<string[]>([]);
   const rowHeaderIds = rowHeaders.length > 0 ? rowHeaders.join(' ') : undefined;
 
@@ -82,6 +82,7 @@ export const TableRow = React.forwardRef<HTMLTableRowElement, TableRowProps>(fun
         return;
       }
       const pointer = event.detail !== 0;
+      const hasSelection = store.get().selection.length > 0;
       if (canSelect && pointer && (hasSelection || !actionable)) {
         toggleSelected({ range: event.shiftKey });
       } else if (actionable) {

packages/mosaic/src/primitives/table/table-selection-checkbox.tsx

@@ -2,8 +2,14 @@
 
 import React from 'react';
 
-import { type ComponentProps, type DefaultProps, mergeProps, useRender } from '../utils';
-import { TableHeaderContext, TableRowContext, useRowHeaderLabel, useTableContext } from './table-context';
+import { type ComponentProps, type DefaultProps, mergeProps, useRender, useStore } from '../utils';
+import {
+  selectAllState,
+  TableHeaderContext,
+  TableRowContext,
+  useRowHeaderLabel,
+  useTableContext,
+} from './table-context';
 
 export type TableSelectionCheckboxProps = ComponentProps<'input'>;
 
@@ -54,6 +60,7 @@ function useSelectionCheckboxState() {
   const table = useTableContext();
   const row = React.useContext(TableRowContext);
   const inHeader = React.useContext(TableHeaderContext);
+  const selectAll = useStore(table.store, state => (row ? 'empty' : selectAllState(state)));
 
   if (row) {
     return {
@@ -66,9 +73,9 @@ function useSelectionCheckboxState() {
   }
   return {
     enabled: inHeader && table.selectionMode === 'multiple',
-    checked: table.allSelected,
-    indeterminate: table.someSelected,
-    disabled: !table.canSelectAll,
+    checked: selectAll === 'all',
+    indeterminate: selectAll === 'some',
+    disabled: selectAll === 'empty',
     toggle: table.toggleAllSelected,
   };
 }

packages/mosaic/src/primitives/table/table.test.tsx

@@ -646,6 +646,31 @@ describe('Table', () => {
       expect(onSelectionChange).toHaveBeenCalledWith(['a']);
       expect(checkbox('Alpha')).not.toBeChecked();
     });
+
+    it('re-renders only the rows whose selection changes', async () => {
+      const user = userEvent.setup();
+      const renders: string[] = [];
+      render(
+        <LogsTable
+          selectionMode='multiple'
+          rowProps={log => ({
+            render: props => {
+              renders.push(log.id);
+              return <tr {...props} />;
+            },
+          })}
+        />,
+      );
+
+      renders.length = 0;
+      await user.click(checkbox('Bravo'));
+      expect(renders).toEqual(['b']);
+
+      renders.length = 0;
+      await user.click(checkbox('Charlie'));
+      expect(renders).toEqual(['c']);
+      expect(screen.getByRole('checkbox', { name: 'Select all' })).toHaveProperty('indeterminate', true);
+    });
   });
 
   describe('tree rows', () => {