Fix column auto resize

src/HeaderRow.tsx

@@ -3,7 +3,7 @@ import { css } from '@linaria/core';
 import clsx from 'clsx';
 
 import { getColSpan } from './utils';
-import type { CalculatedColumn, Direction, Position } from './types';
+import type { CalculatedColumn, Direction, Position, ResizedWidth } from './types';
 import type { DataGridProps } from './DataGrid';
 import HeaderCell from './HeaderCell';
 import { cell, cellFrozen } from './style/cell';
@@ -17,7 +17,7 @@ type SharedDataGridProps<R, SR, K extends React.Key> = Pick<
 export interface HeaderRowProps<R, SR, K extends React.Key> extends SharedDataGridProps<R, SR, K> {
   rowIdx: number;
   columns: readonly CalculatedColumn<R, SR>[];
-  onColumnResize: (column: CalculatedColumn<R, SR>, width: number | 'max-content') => void;
+  onColumnResize: (column: CalculatedColumn<R, SR>, width: ResizedWidth) => void;
   selectCell: (position: Position) => void;
   lastFrozenColumnIndex: number;
   selectedCellIdx: number | undefined;

src/hooks/useColumnWidths.ts

@@ -1,7 +1,7 @@
-import { useLayoutEffect, useRef } from 'react';
+import { useLayoutEffect, useState } from 'react';
 import { flushSync } from 'react-dom';
 
-import type { CalculatedColumn, StateSetter } from '../types';
+import type { CalculatedColumn, ResizedWidth, StateSetter } from '../types';
 import type { DataGridProps } from '../DataGrid';
 
 export function useColumnWidths<R, SR>(
@@ -16,17 +16,27 @@ export function useColumnWidths<R, SR>(
   setMeasuredColumnWidths: StateSetter<ReadonlyMap<string, number>>,
   onColumnResize: DataGridProps<R, SR>['onColumnResize']
 ) {
-  const prevGridWidthRef = useRef(gridWidth);
+  const [columnToAutoResize, setColumnToAutoResize] = useState<{
+    readonly key: string;
+    readonly width: ResizedWidth;
+  } | null>(null);
+  const [prevGridWidth, setPreviousGridWidth] = useState(gridWidth);
   const columnsCanFlex: boolean = columns.length === viewportColumns.length;
   // Allow columns to flex again when...
   const ignorePreviouslyMeasuredColumns: boolean =
     // there is enough space for columns to flex and the grid was resized
-    columnsCanFlex && gridWidth !== prevGridWidthRef.current;
+    columnsCanFlex && gridWidth !== prevGridWidth;
   const newTemplateColumns = [...templateColumns];
   const columnsToMeasure: string[] = [];
 
   for (const { key, idx, width } of viewportColumns) {
-    if (
+    if (key === columnToAutoResize?.key) {
+      newTemplateColumns[idx] =
+        columnToAutoResize.width === 'max-content'
+          ? columnToAutoResize.width
+          : `${columnToAutoResize.width}px`;
+      columnsToMeasure.push(key);
+    } else if (
       typeof width === 'string' &&
       (ignorePreviouslyMeasuredColumns || !measuredColumnWidths.has(key)) &&
       !resizedColumnWidths.has(key)
@@ -38,12 +48,10 @@ export function useColumnWidths<R, SR>(
 
   const gridTemplateColumns = newTemplateColumns.join(' ');
 
-  useLayoutEffect(() => {
-    prevGridWidthRef.current = gridWidth;
-    updateMeasuredWidths(columnsToMeasure);
-  });
+  useLayoutEffect(updateMeasuredWidths);
 
-  function updateMeasuredWidths(columnsToMeasure: readonly string[]) {
+  function updateMeasuredWidths() {
+    setPreviousGridWidth(gridWidth);
     if (columnsToMeasure.length === 0) return;
 
     setMeasuredColumnWidths((measuredColumnWidths) => {
@@ -62,40 +70,54 @@ export function useColumnWidths<R, SR>(
 
       return hasChanges ? newMeasuredColumnWidths : measuredColumnWidths;
     });
-  }
 
-  function handleColumnResize(column: CalculatedColumn<R, SR>, nextWidth: number | 'max-content') {
-    const { key: resizingKey } = column;
-    const newTemplateColumns = [...templateColumns];
-    const columnsToMeasure: string[] = [];
-
-    for (const { key, idx, width } of viewportColumns) {
-      if (resizingKey === key) {
-        const width = typeof nextWidth === 'number' ? `${nextWidth}px` : nextWidth;
-        newTemplateColumns[idx] = width;
-      } else if (columnsCanFlex && typeof width === 'string' && !resizedColumnWidths.has(key)) {
-        newTemplateColumns[idx] = width;
-        columnsToMeasure.push(key);
-      }
+    if (columnToAutoResize !== null) {
+      const resizingKey = columnToAutoResize.key;
+      setResizedColumnWidths((resizedColumnWidths) => {
+        const oldWidth = resizedColumnWidths.get(resizingKey);
+        const newWidth = measureColumnWidth(gridRef, resizingKey);
+        if (newWidth !== undefined && oldWidth !== newWidth) {
+          const newResizedColumnWidths = new Map(resizedColumnWidths);
+          newResizedColumnWidths.set(resizingKey, newWidth);
+          return newResizedColumnWidths;
+        }
+        return resizedColumnWidths;
+      });
+      setColumnToAutoResize(null);
     }
+  }
 
-    gridRef.current!.style.gridTemplateColumns = newTemplateColumns.join(' ');
-    const measuredWidth =
-      typeof nextWidth === 'number' ? nextWidth : measureColumnWidth(gridRef, resizingKey)!;
+  function handleColumnResize(column: CalculatedColumn<R, SR>, nextWidth: ResizedWidth) {
+    const { key: resizingKey } = column;
 
-    // TODO: remove
-    // need flushSync to keep frozen column offsets in sync
-    // we may be able to use `startTransition` or even `requestIdleCallback` instead
     flushSync(() => {
-      setResizedColumnWidths((resizedColumnWidths) => {
-        const newResizedColumnWidths = new Map(resizedColumnWidths);
-        newResizedColumnWidths.set(resizingKey, measuredWidth);
-        return newResizedColumnWidths;
+      if (columnsCanFlex) {
+        // remeasure all the columns that can flex and are not resized by the user
+        setMeasuredColumnWidths((measuredColumnWidths) => {
+          const newMeasuredColumnWidths = new Map(measuredColumnWidths);
+          for (const { key, width } of viewportColumns) {
+            if (resizingKey !== key && typeof width === 'string' && !resizedColumnWidths.has(key)) {
+              newMeasuredColumnWidths.delete(key);
+            }
+          }
+          return newMeasuredColumnWidths;
+        });
+      }
+
+      setColumnToAutoResize({
+        key: resizingKey,
+        width: nextWidth
       });
-      updateMeasuredWidths(columnsToMeasure);
     });
 
-    onColumnResize?.(column, measuredWidth);
+    if (onColumnResize) {
+      const previousWidth = resizedColumnWidths.get(resizingKey);
+      const newWidth =
+        typeof nextWidth === 'number' ? nextWidth : measureColumnWidth(gridRef, resizingKey);
+      if (newWidth !== undefined && newWidth !== previousWidth) {
+        onColumnResize(column, newWidth);
+      }
+    }
   }
 
   return {

src/types.ts

@@ -322,3 +322,5 @@ export interface Renderers<TRow, TSummaryRow> {
 }
 
 export type Direction = 'ltr' | 'rtl';
+
+export type ResizedWidth = number | 'max-content';

test/browser/column/resizable.test.tsx

@@ -66,48 +66,54 @@ test('cannot not resize or auto resize column when resizable is not specified',
 test('should resize column when dragging the handle', async () => {
   const onColumnResize = vi.fn();
   setup<Row, unknown>({ columns, rows: [], onColumnResize });
-  const [, col2] = getHeaderCells();
   const grid = getGrid();
   expect(onColumnResize).not.toHaveBeenCalled();
   await expect.element(grid).toHaveStyle({ gridTemplateColumns: '100px 200px' });
+  const [, col2] = getHeaderCells();
   await resize({ column: col2, resizeBy: -50 });
   await expect.element(grid).toHaveStyle({ gridTemplateColumns: '100px 150px' });
   expect(onColumnResize).toHaveBeenCalledExactlyOnceWith(expect.objectContaining(columns[1]), 150);
 });
 
 test('should use the maxWidth if specified', async () => {
   setup<Row, unknown>({ columns, rows: [] });
-  const [, col2] = getHeaderCells();
   const grid = getGrid();
   await expect.element(grid).toHaveStyle({ gridTemplateColumns: '100px 200px ' });
+  const [, col2] = getHeaderCells();
   await resize({ column: col2, resizeBy: 1000 });
   await expect.element(grid).toHaveStyle({ gridTemplateColumns: '100px 400px' });
 });
 
 test('should use the minWidth if specified', async () => {
   setup<Row, unknown>({ columns, rows: [] });
-  const [, col2] = getHeaderCells();
   const grid = getGrid();
   await expect.element(grid).toHaveStyle({ gridTemplateColumns: '100px 200px' });
+  const [, col2] = getHeaderCells();
   await resize({ column: col2, resizeBy: -150 });
   await expect.element(grid).toHaveStyle({ gridTemplateColumns: '100px 100px' });
 });
 
 test('should auto resize column when resize handle is double clicked', async () => {
+  const onColumnResize = vi.fn();
   setup<Row, unknown>({
     columns,
     rows: [
       {
         col1: 1,
         col2: 'a'.repeat(50)
       }
-    ]
+    ],
+    onColumnResize
   });
-  const [, col2] = getHeaderCells();
   const grid = getGrid();
   await expect.element(grid).toHaveStyle({ gridTemplateColumns: '100px 200px' });
+  const [, col2] = getHeaderCells();
   await autoResize(col2);
   await expect.element(grid).toHaveStyle({ gridTemplateColumns: '100px 327.703px' });
+  expect(onColumnResize).toHaveBeenCalledExactlyOnceWith(
+    expect.objectContaining(columns[1]),
+    327.703125
+  );
 });
 
 test('should use the maxWidth if specified on auto resize', async () => {
@@ -120,9 +126,9 @@ test('should use the maxWidth if specified on auto resize', async () => {
       }
     ]
   });
-  const [, col2] = getHeaderCells();
   const grid = getGrid();
   await expect.element(grid).toHaveStyle({ gridTemplateColumns: '100px 200px' });
+  const [, col2] = getHeaderCells();
   await autoResize(col2);
   await expect.element(grid).toHaveStyle({ gridTemplateColumns: '100px 400px' });
 });
@@ -137,9 +143,57 @@ test('should use the minWidth if specified on auto resize', async () => {
       }
     ]
   });
-  const [, col2] = getHeaderCells();
   const grid = getGrid();
   await expect.element(grid).toHaveStyle({ gridTemplateColumns: '100px 200px' });
+  const [, col2] = getHeaderCells();
   await autoResize(col2);
   await expect.element(grid).toHaveStyle({ gridTemplateColumns: '100px 100px' });
 });
+
+test('should remeasure flex columns when resizing a column', async () => {
+  const onColumnResize = vi.fn();
+  setup<
+    {
+      readonly col1: string;
+      readonly col2: string;
+      readonly col3: string;
+    },
+    unknown
+  >({
+    columns: [
+      {
+        key: 'col1',
+        name: 'col1',
+        resizable: true
+      },
+      {
+        key: 'col2',
+        name: 'col2',
+        resizable: true
+      },
+      {
+        key: 'col3',
+        name: 'col3',
+        resizable: true
+      }
+    ],
+    rows: [
+      {
+        col1: 'a'.repeat(10),
+        col2: 'a'.repeat(10),
+        col3: 'a'.repeat(10)
+      }
+    ],
+    onColumnResize
+  });
+  const grid = getGrid();
+  await expect.element(grid).toHaveStyle({ gridTemplateColumns: '639.328px 639.328px 639.344px' });
+  const [col1] = getHeaderCells();
+  await autoResize(col1);
+  await expect.element(grid).toHaveStyle({ gridTemplateColumns: '79.1406px 919.422px 919.438px' });
+  expect(onColumnResize).toHaveBeenCalledOnce();
+  // onColumnResize is not called if width is not changed
+  await autoResize(col1);
+  await expect.element(grid).toHaveStyle({ gridTemplateColumns: '79.1406px 919.422px 919.438px' });
+  expect(onColumnResize).toHaveBeenCalledOnce();
+});