Use state to handle column auto resize

Author: amanmahajan7

src/hooks/useColumnWidths.ts

@@ -1,5 +1,4 @@
-import { useLayoutEffect, useRef } from 'react';
-import { flushSync } from 'react-dom';
+import { useLayoutEffect, useRef, useState } from 'react';
 
 import type { CalculatedColumn, StateSetter } from '../types';
 import type { DataGridProps } from '../DataGrid';
@@ -16,6 +15,7 @@ export function useColumnWidths<R, SR>(
   setMeasuredColumnWidths: StateSetter<ReadonlyMap<string, number>>,
   onColumnResize: DataGridProps<R, SR>['onColumnResize']
 ) {
+  const [columnToAutoresize, setColumnToAutoResize] = useState<string | null>(null);
   const prevGridWidthRef = useRef(gridWidth);
   const columnsCanFlex: boolean = columns.length === viewportColumns.length;
   // Allow columns to flex again when...
@@ -26,7 +26,10 @@ export function useColumnWidths<R, SR>(
   const columnsToMeasure: string[] = [];
 
   for (const { key, idx, width } of viewportColumns) {
-    if (
+    if (key === columnToAutoresize) {
+      newTemplateColumns[idx] = 'max-content';
+      columnsToMeasure.push(key);
+    } else if (
       typeof width === 'string' &&
       (ignorePreviouslyMeasuredColumns || !measuredColumnWidths.has(key)) &&
       !resizedColumnWidths.has(key)
@@ -62,40 +65,53 @@ export function useColumnWidths<R, SR>(
 
       return hasChanges ? newMeasuredColumnWidths : measuredColumnWidths;
     });
+
+    if (columnToAutoresize !== null) {
+      setColumnToAutoResize(null);
+      setResizedColumnWidths((resizedColumnWidths) => {
+        const newResizedColumnWidths = new Map(resizedColumnWidths);
+        newResizedColumnWidths.set(
+          columnToAutoresize,
+          measureColumnWidth(gridRef, columnToAutoresize)!
+        );
+        return newResizedColumnWidths;
+      });
+    }
   }
 
   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;
+    for (const { key, width } of viewportColumns) {
+      if (
+        resizingKey !== key &&
+        columnsCanFlex &&
+        typeof width === 'string' &&
+        !resizedColumnWidths.has(key)
+      ) {
         columnsToMeasure.push(key);
       }
     }
 
-    gridRef.current!.style.gridTemplateColumns = newTemplateColumns.join(' ');
-    const measuredWidth =
-      typeof nextWidth === 'number' ? nextWidth : measureColumnWidth(gridRef, resizingKey)!;
+    setMeasuredColumnWidths((measuredColumnWidths) => {
+      const newMeasuredColumnWidths = new Map(measuredColumnWidths);
+      for (const columnKey of columnsToMeasure) {
+        newMeasuredColumnWidths.delete(columnKey);
+      }
+      return newMeasuredColumnWidths;
+    });
 
-    // TODO: remove
-    // need flushSync to keep frozen column offsets in sync
-    // we may be able to use `startTransition` or even `requestIdleCallback` instead
-    flushSync(() => {
+    if (typeof nextWidth === 'number') {
       setResizedColumnWidths((resizedColumnWidths) => {
         const newResizedColumnWidths = new Map(resizedColumnWidths);
-        newResizedColumnWidths.set(resizingKey, measuredWidth);
+        newResizedColumnWidths.set(resizingKey, nextWidth);
         return newResizedColumnWidths;
       });
-      updateMeasuredWidths(columnsToMeasure);
-    });
-
-    onColumnResize?.(column, measuredWidth);
+      onColumnResize?.(column, nextWidth);
+    } else {
+      setColumnToAutoResize(resizingKey);
+    }
   }
 
   return {