[ENHANCEMENT] Compile TanStack Table v9 components with the React Compiler

components/src/Table/Table.test.tsx

@@ -12,8 +12,10 @@
 // limitations under the License.
 
 import type { RenderResult } from '@testing-library/react';
-import { render, screen, getAllByRole, within } from '@testing-library/react';
+import { fireEvent, render, screen, getAllByRole, within } from '@testing-library/react';
 import userEvent from '@testing-library/user-event';
+import type { ReactElement } from 'react';
+import { useState } from 'react';
 import { VirtuosoMockContext } from 'react-virtuoso';
 
 import type { TableCellConfigs, TableColumnConfig, TableProps } from './model/table-model';
@@ -124,6 +126,67 @@ const LARGE_TABLE_OVERALL_ROWS = 100;
 const LARGE_TABLE_VISIBLE_ROWS = 5;
 const LARGE_TABLE_DATA = generateMockTableData(LARGE_TABLE_OVERALL_ROWS);
 
+const STATEFUL_TABLE_DATA = generateMockTableData(3);
+const STATEFUL_TABLE_VIRTUOSO_CONTEXT = { viewportHeight: 600, itemHeight: MOCK_ITEM_HEIGHT };
+
+/**
+ * Feeds table state changes back into the table like a real consumer does, so
+ * tests cover re-renders driven by state rather than only the emitted events.
+ */
+function StatefulTable(props: Partial<TableProps<MockTableData>>): ReactElement {
+  const [rowSelection, setRowSelection] = useState<NonNullable<TableProps<MockTableData>['rowSelection']>>({});
+  const [sorting, setSorting] = useState<NonNullable<TableProps<MockTableData>['sorting']>>([]);
+
+  return (
+    <VirtuosoMockContext.Provider value={STATEFUL_TABLE_VIRTUOSO_CONTEXT}>
+      <Table
+        data={STATEFUL_TABLE_DATA}
+        columns={COLUMNS}
+        height={600}
+        width={300}
+        checkboxSelection
+        rowSelection={rowSelection}
+        onRowSelectionChange={setRowSelection}
+        sorting={sorting}
+        onSortingChange={setSorting}
+        {...props}
+      />
+    </VirtuosoMockContext.Provider>
+  );
+}
+
+const RESIZABLE_COLUMN_CONFIG = { enableResizing: true };
+const COLUMN_FILTER_TOOLBAR_CONFIG = { isColumnFilterEnabled: true };
+
+function getColumnHeader(index: number): HTMLElement {
+  const columnHeader = screen.getAllByRole('columnheader')[index];
+  if (!columnHeader) {
+    throw new Error(`Cannot find column header at index: ${index}.`);
+  }
+  return columnHeader;
+}
+
+function getHeaderCheckbox(): HTMLElement {
+  return within(getColumnHeader(0)).getByRole('checkbox');
+}
+
+function getSortButton(): HTMLElement {
+  return within(getColumnHeader(2)).getByRole('button');
+}
+
+function getSortIcon(): HTMLElement {
+  return within(getColumnHeader(2)).getByTestId('ArrowDownwardIcon');
+}
+
+function getResizeHandle(): HTMLElement {
+  // Query on each use because uncompiled re-renders can remount the header row.
+  const resizeHandle = within(getColumnHeader(2)).getByRole('separator').parentElement;
+  if (!resizeHandle) {
+    throw new Error('Cannot find resize handle');
+  }
+  return resizeHandle;
+}
+
 const renderLargeTable = (): void => {
   const height = LARGE_TABLE_VISIBLE_ROWS * MOCK_ITEM_HEIGHT;
 
@@ -1182,4 +1245,69 @@ describe('table', () => {
     expect(cell4.firstChild).not.toHaveStyle({ 'background-color': 'rgb(255, 0, 0)' });
     expect(cell4.firstChild).not.toHaveStyle({ color: 'green' });
   });
+
+  describe('when the parent feeds table state back in', () => {
+    test('header checkbox toggles between all and none on repeated clicks', () => {
+      render(<StatefulTable />);
+
+      for (let i = 0; i < 2; i++) {
+        userEvent.click(getHeaderCheckbox());
+        expect(getHeaderCheckbox()).toBeChecked();
+        screen.getAllByRole('checkbox').forEach((checkbox) => expect(checkbox).toBeChecked());
+
+        userEvent.click(getHeaderCheckbox());
+        expect(getHeaderCheckbox()).not.toBeChecked();
+        screen.getAllByRole('checkbox').forEach((checkbox) => expect(checkbox).not.toBeChecked());
+      }
+    });
+
+    test('sort indicator follows repeated header clicks', () => {
+      render(<StatefulTable />);
+
+      userEvent.click(getSortButton());
+      expect(getSortButton()).toHaveClass('Mui-active');
+      expect(getSortIcon()).toHaveClass('MuiTableSortLabel-iconDirectionDesc');
+
+      userEvent.click(getSortButton());
+      expect(getSortButton()).toHaveClass('Mui-active');
+      expect(getSortIcon()).toHaveClass('MuiTableSortLabel-iconDirectionAsc');
+
+      userEvent.click(getSortButton());
+      expect(getSortButton()).not.toHaveClass('Mui-active');
+    });
+
+    test('column visibility menu reflects toggled columns', () => {
+      render(<StatefulTable tableToolbarConfig={COLUMN_FILTER_TOOLBAR_CONFIG} />);
+      userEvent.click(screen.getByRole('button', { expanded: false }));
+      const menuItem = within(screen.getByRole('menu')).getByText('Value').closest('li');
+      if (!menuItem) {
+        throw new Error('Cannot find column visibility menu item');
+      }
+
+      userEvent.click(menuItem);
+      expect(within(menuItem).getByRole('checkbox')).not.toBeChecked();
+      expect(within(screen.getByRole('table', { hidden: true })).queryByText('Value')).not.toBeInTheDocument();
+
+      userEvent.click(menuItem);
+      expect(within(menuItem).getByRole('checkbox')).toBeChecked();
+      expect(within(screen.getByRole('table', { hidden: true })).getByText('Value')).toBeInTheDocument();
+    });
+
+    test('resizing a column updates its size and double clicking the handle resets it', () => {
+      const { container } = render(<StatefulTable defaultColumnConfig={RESIZABLE_COLUMN_CONFIG} />);
+      const getColumnSize = (): string =>
+        (container.querySelector('[style*="--col-value-size"]') as HTMLElement).style.getPropertyValue(
+          '--col-value-size',
+        );
+      const initialSize = getColumnSize();
+
+      fireEvent.mouseDown(getResizeHandle(), { clientX: 100 });
+      fireEvent.mouseMove(document, { clientX: 150 });
+      fireEvent.mouseUp(document, { clientX: 150 });
+      expect(getColumnSize()).toBe(String(Number(initialSize) + 50));
+
+      fireEvent.mouseDown(getResizeHandle(), { detail: 2 });
+      expect(getColumnSize()).toBe(initialSize);
+    });
+  });
 });

components/src/Table/Table.tsx

@@ -76,8 +76,6 @@ export function Table<TableData extends RowData>({
   defaultColumnConfig,
   ...otherProps
 }: TableProps<TableData>): ReactElement {
-  'use no memo'; // TanStack Table v8 uses a mutable instance that automatic memoization cannot track.
-
   const theme = useTheme();
 
   const hasSubRows = !!getSubRows;
@@ -236,13 +234,14 @@ export function Table<TableData extends RowData>({
     },
   });
 
+  // Use the row's stable core table rather than `table`, whose identity changes
+  // with every state update and would otherwise remount all virtualized rows.
   const handleRowClick = useCallback(
-    (e: React.MouseEvent<HTMLDivElement, MouseEvent>, rowId: string) => {
-      const row = table.getRow(rowId);
+    (e: React.MouseEvent<HTMLDivElement, MouseEvent>, row: Row<TableFeatures, TableData>) => {
       const isModifiedClick = e.metaKey || e.shiftKey;
-      handleRowSelectionEvent(table, row, isModifiedClick);
+      handleRowSelectionEvent(row.table, row, isModifiedClick);
     },
-    [handleRowSelectionEvent, table],
+    [handleRowSelectionEvent],
   );
 
   return (
@@ -252,15 +251,10 @@ export function Table<TableData extends RowData>({
       defaultColumnWidth={defaultColumnWidth}
       defaultColumnHeight={defaultColumnHeight}
       onRowClick={handleRowClick}
-      rows={table.getRowModel().rows}
-      columns={table.getVisibleFlatColumns()}
-      columnSizing={table.state.columnSizing}
-      columnResizing={table.state.columnResizing}
-      headers={table.getHeaderGroups()}
+      table={table}
       cellConfigs={cellConfigs}
       pagination={pagination}
       onPaginationChange={onPaginationChange}
-      rowCount={table.getRowCount()}
       toolbarConfig={{
         isSearchEnabled: tableToolbarConfig?.isSearchEnabled,
         globalFilter,

components/src/Table/TableToolbar.tsx

@@ -13,6 +13,7 @@
 
 import { Checkbox, IconButton, InputAdornment, ListItemText, Menu, MenuItem, Stack, Tooltip } from '@mui/material';
 import type { Column, RowData } from '@tanstack/react-table';
+import { Subscribe } from '@tanstack/react-table';
 import Close from 'mdi-material-ui/Close';
 import Magnify from 'mdi-material-ui/Magnify';
 import UnfoldLess from 'mdi-material-ui/UnfoldLessHorizontal';
@@ -87,8 +88,6 @@ export function TableToolbar<TableData extends RowData>({
   isAllExpanded,
   onExpandAllChange,
 }: TableToolbarProps<TableData>): ReactElement | null {
-  'use no memo'; // TanStack Table columns retain their identity when visibility changes.
-
   const [colMenuAnchor, setColMenuAnchor] = useState<null | HTMLElement>(null);
   const colMenuOpen = Boolean(colMenuAnchor);
   const [searchResetKey, setSearchResetKey] = useState(0);
@@ -177,13 +176,18 @@ export function TableToolbar<TableData extends RowData>({
                   onClick={column.getCanHide() ? column.getToggleVisibilityHandler() : undefined}
                   dense
                 >
-                  <Checkbox
-                    checked={column.getIsVisible()}
-                    disabled={!column.getCanHide()}
-                    size="small"
-                    disableRipple
-                    sx={{ p: 0, mr: 1 }}
-                  />
+                  {/* Columns keep their identity when visibility changes, so subscribe to it explicitly. */}
+                  <Subscribe source={column.table.atoms.columnVisibility}>
+                    {() => (
+                      <Checkbox
+                        checked={column.getIsVisible()}
+                        disabled={!column.getCanHide()}
+                        size="small"
+                        disableRipple
+                        sx={{ p: 0, mr: 1 }}
+                      />
+                    )}
+                  </Subscribe>
                   <ListItemText primary={label} />
                 </MenuItem>
               );

components/src/Table/VirtualizedTable.tsx

@@ -12,7 +12,7 @@
 // limitations under the License.
 
 import { Box, TablePagination, TableRow as MuiTableRow } from '@mui/material';
-import type { Column, columnResizingState, ColumnSizingState, HeaderGroup, Row, RowData } from '@tanstack/react-table';
+import type { ReactTable, Row, RowData } from '@tanstack/react-table';
 import { flexRender } from '@tanstack/react-table';
 import type { ReactElement } from 'react';
 import { useMemo, useRef } from 'react';
@@ -42,14 +42,14 @@ export type VirtualizedTableProps<TableData extends RowData> = Required<
   Pick<TableProps<TableData>, 'height' | 'width' | 'density' | 'defaultColumnHeight' | 'defaultColumnWidth'>
 > &
   Pick<TableProps<TableData>, 'onRowMouseOver' | 'onRowMouseOut' | 'pagination' | 'onPaginationChange'> & {
-    onRowClick: (e: React.MouseEvent<HTMLDivElement, MouseEvent>, id: string) => void;
-    rows: Array<Row<TableFeatures, TableData>>;
-    columns: Array<Column<TableFeatures, TableData, unknown>>;
-    headers: Array<HeaderGroup<TableFeatures, TableData>>;
-    columnSizing: ColumnSizingState;
-    columnResizing: columnResizingState;
+    onRowClick: (e: React.MouseEvent<HTMLDivElement, MouseEvent>, row: Row<TableFeatures, TableData>) => void;
+    /**
+     * The value returned by `useTable`. Its identity changes whenever table
+     * state or options change, unlike the row, column, and header objects it
+     * returns, which keep their identities across state changes.
+     */
+    table: ReactTable<TableFeatures, TableData>;
     cellConfigs?: TableCellConfigs;
-    rowCount: number;
     toolbarConfig: Pick<
       TableToolbarProps<TableData>,
       | 'isSearchEnabled'
@@ -76,16 +76,18 @@ export function VirtualizedTable<TableData extends RowData>({
   onRowClick,
   onRowMouseOver,
   onRowMouseOut,
-  rows,
-  columns,
-  headers,
+  table,
   cellConfigs,
   pagination,
   onPaginationChange,
-  rowCount,
   toolbarConfig,
 }: VirtualizedTableProps<TableData>): ReactElement {
-  'use no memo'; // TanStack Table row and column getters read mutable state behind stable identities.
+  // Rows, columns, and headers keep their identities when table state such as
+  // selection, sorting, or sizing changes. Anything that reads that state
+  // through their methods must call `table` directly in the same expression or
+  // callback, so the React Compiler invalidates it whenever `table` changes.
+  const rows = table.getRowModel().rows;
+  const columnCount = table.getVisibleFlatColumns().length;
 
   const virtuosoRef = useRef<TableVirtuosoHandle>(null);
   // Use a ref for these values because they are only needed for keyboard
@@ -106,7 +108,7 @@ export function VirtualizedTable<TableData extends RowData>({
 
     // We add 1 here for the header.
     maxRows: rows.length + 1,
-    maxColumns: columns.length,
+    maxColumns: columnCount,
   });
 
   const getFocusState = (cellPosition: TableCellPosition): TableCellProps['focusState'] => {
@@ -145,7 +147,7 @@ export function VirtualizedTable<TableData extends RowData>({
         return (
           <TableRow
             {...props}
-            onClick={(e) => onRowClick(e, row.id)}
+            onClick={(e) => onRowClick(e, row)}
             density={density}
             onMouseOver={(e) => {
               onRowMouseOver?.(e, rowEventOpts);
@@ -186,7 +188,7 @@ export function VirtualizedTable<TableData extends RowData>({
    * and pass the column sizes down as CSS variables to the <table> element.
    */
   const columnSizeVars: { [key: string]: number } = {};
-  headers.forEach((headerGroup) => {
+  table.getHeaderGroups().forEach((headerGroup) => {
     headerGroup.headers
       .filter((header) => header.column.getCanResize())
       .forEach((header) => {
@@ -209,7 +211,7 @@ export function VirtualizedTable<TableData extends RowData>({
         fixedHeaderContent={() => {
           return (
             <>
-              {headers.map((headerGroup) => {
+              {table.getHeaderGroups().map((headerGroup) => {
                 return (
                   <TableRow key={headerGroup.id} density={density}>
                     {headerGroup.headers.map((header, i, headers) => {
@@ -246,7 +248,7 @@ export function VirtualizedTable<TableData extends RowData>({
                             header.column.getCanResize()
                               ? {
                                   resizeHandler: header.getResizeHandler(),
-                                  resetSizeHandler: header.column.resetSize,
+                                  resetSizeHandler: () => header.column.resetSize(),
                                   isResizing: header.column.getIsResizing(),
                                 }
                               : undefined
@@ -267,8 +269,8 @@ export function VirtualizedTable<TableData extends RowData>({
             ? (): ReactElement => (
                 <MuiTableRow sx={{ backgroundColor: (theme) => theme.palette.background.default }}>
                   <TablePagination
-                    colSpan={columns.length}
-                    count={rowCount}
+                    colSpan={table.getVisibleFlatColumns().length}
+                    count={table.getRowCount()}
                     page={pagination.pageIndex}
                     rowsPerPage={pagination.pageSize}
                     onPageChange={handleChangePage}
@@ -279,7 +281,7 @@ export function VirtualizedTable<TableData extends RowData>({
             : undefined
         }
         itemContent={(index) => {
-          const row = rows[index];
+          const row = table.getRowModel().rows[index];
           if (!row) {
             return null;
           }