Add benchmarks and optimize hot paths

package.json

@@ -40,6 +40,7 @@
     "test:watch": "vitest watch --project browser --project node",
     "visual": "vitest run --project visual --coverage.reportsDirectory='./coverage/visual'",
     "visual:update": "vitest run --project visual --update",
+    "bench": "vitest bench --project bench",
     "format": "oxfmt",
     "format:check": "oxfmt --check",
     "eslint": "eslint --max-warnings 0 --cache --cache-location .cache/eslint --cache-strategy content",

src/HeaderCell.tsx

@@ -1,4 +1,4 @@
-import { useRef, useState } from 'react';
+import { memo, useRef, useState } from 'react';
 import { flushSync } from 'react-dom';
 import { css } from 'ecij';
 
@@ -13,8 +13,7 @@ import {
   isCtrlKeyHeldDown,
   stopPropagation
 } from './utils';
-import type { CalculatedColumn, SortColumn } from './types';
-import type { HeaderRowProps } from './HeaderRow';
+import type { CalculatedColumn, Direction, Position, ResizedWidth, SortDirection } from './types';
 
 const cellSortableClassname = css`
   @layer rdg.HeaderCell {
@@ -64,37 +63,35 @@ const dragImageClassname = css`
   }
 `;
 
-type SharedHeaderRowProps<R, SR> = Pick<
-  HeaderRowProps<R, SR, React.Key>,
-  | 'sortColumns'
-  | 'onSortColumnsChange'
-  | 'setPosition'
-  | 'onColumnResize'
-  | 'onColumnResizeEnd'
-  | 'shouldFocusGrid'
-  | 'direction'
-  | 'onColumnsReorder'
->;
-
-export interface HeaderCellProps<R, SR> extends SharedHeaderRowProps<R, SR> {
+export interface HeaderCellProps<R, SR> {
   column: CalculatedColumn<R, SR>;
   colSpan: number | undefined;
   rowIdx: number;
   isCellActive: boolean;
+  sortDirection: SortDirection | undefined;
+  priority: number | undefined;
+  onSort: (column: CalculatedColumn<R, SR>, ctrlClick: boolean) => void;
+  onColumnResize: (column: CalculatedColumn<R, SR>, width: ResizedWidth) => void;
+  onColumnResizeEnd: () => void;
+  onColumnsReorder: ((sourceColumnKey: string, targetColumnKey: string) => void) | undefined | null;
+  setPosition: (position: Position) => void;
+  shouldFocusGrid: boolean;
+  direction: Direction;
   draggedColumnKey: string | undefined;
   setDraggedColumnKey: (draggedColumnKey: string | undefined) => void;
 }
 
-export default function HeaderCell<R, SR>({
+function HeaderCell<R, SR>({
   column,
   colSpan,
   rowIdx,
   isCellActive,
+  sortDirection,
+  priority,
+  onSort,
   onColumnResize,
   onColumnResizeEnd,
   onColumnsReorder,
-  sortColumns,
-  onSortColumnsChange,
   setPosition,
   shouldFocusGrid,
   direction,
@@ -107,11 +104,6 @@ export default function HeaderCell<R, SR>({
   const rowSpan = getHeaderCellRowSpan(column, rowIdx);
   // set the tabIndex to 0 when there is no active cell so grid can receive focus
   const { tabIndex, childTabIndex, onFocus } = useRovingTabIndex(shouldFocusGrid || isCellActive);
-  const sortIndex = sortColumns?.findIndex((sort) => sort.columnKey === column.key);
-  const sortColumn =
-    sortIndex !== undefined && sortIndex > -1 ? sortColumns![sortIndex] : undefined;
-  const sortDirection = sortColumn?.direction;
-  const priority = sortColumn !== undefined && sortColumns!.length > 1 ? sortIndex! + 1 : undefined;
   const ariaSort =
     sortDirection && !priority ? (sortDirection === 'ASC' ? 'ascending' : 'descending') : undefined;
   const { sortable, resizable, draggable } = column;
@@ -126,43 +118,6 @@ export default function HeaderCell<R, SR>({
     isOver && cellOverClassname
   );
 
-  function onSort(ctrlClick: boolean) {
-    if (onSortColumnsChange == null) return;
-    const { sortDescendingFirst } = column;
-    if (sortColumn === undefined) {
-      // not currently sorted
-      const nextSort: SortColumn = {
-        columnKey: column.key,
-        direction: sortDescendingFirst ? 'DESC' : 'ASC'
-      };
-      onSortColumnsChange(sortColumns && ctrlClick ? [...sortColumns, nextSort] : [nextSort]);
-    } else {
-      let nextSortColumn: SortColumn | undefined;
-      if (
-        (sortDescendingFirst === true && sortDirection === 'DESC') ||
-        (sortDescendingFirst !== true && sortDirection === 'ASC')
-      ) {
-        nextSortColumn = {
-          columnKey: column.key,
-          direction: sortDirection === 'ASC' ? 'DESC' : 'ASC'
-        };
-      }
-      if (ctrlClick) {
-        const nextSortColumns = [...sortColumns!];
-        if (nextSortColumn) {
-          // swap direction
-          nextSortColumns[sortIndex!] = nextSortColumn;
-        } else {
-          // remove sort
-          nextSortColumns.splice(sortIndex!, 1);
-        }
-        onSortColumnsChange(nextSortColumns);
-      } else {
-        onSortColumnsChange(nextSortColumn ? [nextSortColumn] : []);
-      }
-    }
-  }
-
   function handleFocus(event: React.FocusEvent<HTMLDivElement>) {
     onFocus?.(event);
     if (shouldFocusGrid) {
@@ -177,7 +132,7 @@ export default function HeaderCell<R, SR>({
 
   function onClick(event: React.MouseEvent<HTMLSpanElement>) {
     if (sortable) {
-      onSort(event.ctrlKey || event.metaKey);
+      onSort(column, event.ctrlKey || event.metaKey);
     }
   }
 
@@ -186,7 +141,7 @@ export default function HeaderCell<R, SR>({
     if (sortable && (key === ' ' || key === 'Enter')) {
       // prevent scrolling
       event.preventDefault();
-      onSort(event.ctrlKey || event.metaKey);
+      onSort(column, event.ctrlKey || event.metaKey);
     } else if (
       resizable &&
       isCtrlKeyHeldDown(event) &&
@@ -317,10 +272,14 @@ export default function HeaderCell<R, SR>({
   );
 }
 
-type ResizeHandleProps<R, SR> = Pick<
-  HeaderCellProps<R, SR>,
-  'direction' | 'column' | 'onColumnResize' | 'onColumnResizeEnd'
->;
+export default memo(HeaderCell) as <R, SR>(props: HeaderCellProps<R, SR>) => React.JSX.Element;
+
+interface ResizeHandleProps<R, SR> {
+  direction: Direction;
+  column: CalculatedColumn<R, SR>;
+  onColumnResize: (column: CalculatedColumn<R, SR>, width: ResizedWidth) => void;
+  onColumnResizeEnd: () => void;
+}
 
 function ResizeHandle<R, SR>({
   direction,

src/HeaderRow.tsx

@@ -1,14 +1,17 @@
-import { memo, useState } from 'react';
+import { memo, useMemo, useState } from 'react';
 import { css } from 'ecij';
 
+import { useLatestFunc } from './hooks';
 import { classnames } from './utils';
 import type {
   CalculatedColumn,
   Direction,
   IterateOverViewportColumnsForRow,
   Maybe,
   Position,
-  ResizedWidth
+  ResizedWidth,
+  SortColumn,
+  SortDirection
 } from './types';
 import type { DataGridProps } from './DataGrid';
 import HeaderCell from './HeaderCell';
@@ -52,6 +55,12 @@ const headerRow = css`
 
 export const headerRowClassname = `rdg-header-row ${headerRow}`;
 
+interface SortInfo {
+  direction: SortDirection;
+  priority: number | undefined;
+  index: number;
+}
+
 function HeaderRow<R, SR, K extends React.Key>({
   headerRowClass,
   rowIdx,
@@ -69,26 +78,86 @@ function HeaderRow<R, SR, K extends React.Key>({
   const [draggedColumnKey, setDraggedColumnKey] = useState<string>();
   const isPositionOnRow = activeCellIdx === -1;
 
+  const sortMap = useMemo(() => {
+    if (!sortColumns?.length) return undefined;
+    const map = new Map<string, SortInfo>();
+    for (let i = 0; i < sortColumns.length; i++) {
+      const sc = sortColumns[i];
+      map.set(sc.columnKey, {
+        direction: sc.direction,
+        priority: sortColumns.length > 1 ? i + 1 : undefined,
+        index: i
+      });
+    }
+    return map;
+  }, [sortColumns]);
+
+  function handleSort(column: CalculatedColumn<R, SR>, ctrlClick: boolean) {
+    if (onSortColumnsChange == null) return;
+    const { sortDescendingFirst } = column;
+    const sortInfo = sortMap?.get(column.key);
+    const currentSortDirection = sortInfo?.direction;
+
+    if (currentSortDirection === undefined) {
+      // not currently sorted
+      const nextSort: SortColumn = {
+        columnKey: column.key,
+        direction: sortDescendingFirst ? 'DESC' : 'ASC'
+      };
+      onSortColumnsChange(sortColumns && ctrlClick ? [...sortColumns, nextSort] : [nextSort]);
+    } else {
+      let nextSortColumn: SortColumn | undefined;
+      if (
+        (sortDescendingFirst === true && currentSortDirection === 'DESC') ||
+        (sortDescendingFirst !== true && currentSortDirection === 'ASC')
+      ) {
+        nextSortColumn = {
+          columnKey: column.key,
+          direction: currentSortDirection === 'ASC' ? 'DESC' : 'ASC'
+        };
+      }
+      if (ctrlClick) {
+        const nextSortColumns = [...sortColumns!];
+        if (nextSortColumn) {
+          // swap direction
+          nextSortColumns[sortInfo!.index] = nextSortColumn;
+        } else {
+          // remove sort
+          nextSortColumns.splice(sortInfo!.index, 1);
+        }
+        onSortColumnsChange(nextSortColumns);
+      } else {
+        onSortColumnsChange(nextSortColumn ? [nextSortColumn] : []);
+      }
+    }
+  }
+
+  const handleSortLatest = useLatestFunc(handleSort);
+
   const cells = iterateOverViewportColumnsForRow(activeCellIdx, { type: 'HEADER' })
-    .map(([column, isCellActive, colSpan], index) => (
-      <HeaderCell<R, SR>
-        key={column.key}
-        column={column}
-        colSpan={colSpan}
-        rowIdx={rowIdx}
-        isCellActive={isCellActive}
-        onColumnResize={onColumnResize}
-        onColumnResizeEnd={onColumnResizeEnd}
-        onColumnsReorder={onColumnsReorder}
-        onSortColumnsChange={onSortColumnsChange}
-        sortColumns={sortColumns}
-        setPosition={setPosition}
-        shouldFocusGrid={shouldFocusGrid && index === 0}
-        direction={direction}
-        draggedColumnKey={draggedColumnKey}
-        setDraggedColumnKey={setDraggedColumnKey}
-      />
-    ))
+    .map(([column, isCellActive, colSpan], index) => {
+      const sortInfo = sortMap?.get(column.key);
+      return (
+        <HeaderCell<R, SR>
+          key={column.key}
+          column={column}
+          colSpan={colSpan}
+          rowIdx={rowIdx}
+          isCellActive={isCellActive}
+          onColumnResize={onColumnResize}
+          onColumnResizeEnd={onColumnResizeEnd}
+          onColumnsReorder={onColumnsReorder}
+          sortDirection={sortInfo?.direction}
+          priority={sortInfo?.priority}
+          onSort={handleSortLatest}
+          setPosition={setPosition}
+          shouldFocusGrid={shouldFocusGrid && index === 0}
+          direction={direction}
+          draggedColumnKey={draggedColumnKey}
+          setDraggedColumnKey={setDraggedColumnKey}
+        />
+      );
+    })
     .toArray();
 
   return (

src/hooks/useCalculatedColumns.ts

@@ -166,9 +166,9 @@ export function useCalculatedColumns<R, SR>({
     templateColumns: readonly string[];
     layoutCssVars: Readonly<Record<string, string>>;
     totalFrozenColumnWidth: number;
-    columnMetrics: ReadonlyMap<CalculatedColumn<R, SR>, ColumnMetric>;
+    columnMetrics: readonly ColumnMetric[];
   } => {
-    const columnMetrics = new Map<CalculatedColumn<R, SR>, ColumnMetric>();
+    const columnMetrics: ColumnMetric[] = [];
     let left = 0;
     let totalFrozenColumnWidth = 0;
     const templateColumns: string[] = [];
@@ -184,20 +184,19 @@ export function useCalculatedColumns<R, SR>({
         width = column.minWidth;
       }
       templateColumns.push(`${width}px`);
-      columnMetrics.set(column, { width, left });
+      columnMetrics.push({ width, left });
       left += width;
     }
 
     if (lastFrozenColumnIndex !== -1) {
-      const columnMetric = columnMetrics.get(columns[lastFrozenColumnIndex])!;
+      const columnMetric = columnMetrics[lastFrozenColumnIndex];
       totalFrozenColumnWidth = columnMetric.left + columnMetric.width;
     }
 
     const layoutCssVars: Record<string, string> = {};
 
     for (let i = 0; i <= lastFrozenColumnIndex; i++) {
-      const column = columns[i];
-      layoutCssVars[`--rdg-frozen-left-${column.idx}`] = `${columnMetrics.get(column)!.left}px`;
+      layoutCssVars[`--rdg-frozen-left-${columns[i].idx}`] = `${columnMetrics[i].left}px`;
     }
 
     return { templateColumns, layoutCssVars, totalFrozenColumnWidth, columnMetrics };
@@ -222,7 +221,7 @@ export function useCalculatedColumns<R, SR>({
     // get the first visible non-frozen column index
     let colVisibleStartIdx = firstUnfrozenColumnIdx;
     while (colVisibleStartIdx < lastColIdx) {
-      const { left, width } = columnMetrics.get(columns[colVisibleStartIdx])!;
+      const { left, width } = columnMetrics[colVisibleStartIdx];
       // if the right side of the columnn is beyond the left side of the available viewport,
       // then it is the first column that's at least partially visible
       if (left + width > viewportLeft) {
@@ -234,7 +233,7 @@ export function useCalculatedColumns<R, SR>({
     // get the last visible non-frozen column index
     let colVisibleEndIdx = colVisibleStartIdx;
     while (colVisibleEndIdx < lastColIdx) {
-      const { left, width } = columnMetrics.get(columns[colVisibleEndIdx])!;
+      const { left, width } = columnMetrics[colVisibleEndIdx];
       // if the right side of the column is beyond or equal to the right side of the available viewport,
       // then it the last column that's at least partially visible, as the previous column's right side is not beyond the viewport.
       if (left + width >= viewportRight) {

src/hooks/useColumnWidths.ts

@@ -28,12 +28,13 @@ export function useColumnWidths<R, SR>(
     columnsCanFlex &&
     // there is enough space for columns to flex and the grid was resized
     gridWidth !== prevGridWidth;
-  const newTemplateColumns = [...templateColumns];
+  let newTemplateColumns: string[] | undefined;
   const columnsToMeasure: string[] = [];
 
   for (const { key, idx, width } of viewportColumns) {
     const columnWidth = columnWidths.get(key);
     if (key === columnToAutoResize?.key) {
+      newTemplateColumns ??= [...templateColumns];
       newTemplateColumns[idx] =
         columnToAutoResize.width === 'max-content'
           ? columnToAutoResize.width
@@ -47,12 +48,13 @@ export function useColumnWidths<R, SR>(
         columnsToMeasureOnResize?.has(key) === true ||
         columnWidth === undefined)
     ) {
+      newTemplateColumns ??= [...templateColumns];
       newTemplateColumns[idx] = width;
       columnsToMeasure.push(key);
     }
   }
 
-  const gridTemplateColumns = newTemplateColumns.join(' ');
+  const gridTemplateColumns = (newTemplateColumns ?? templateColumns).join(' ');
 
   useLayoutEffect(updateMeasuredAndResizedWidths);