From cf888022886d643445fe179a90dc49f17609fa67 Mon Sep 17 00:00:00 2001
From: Aman Mahajan <amahajan@stratag.com>
Date: Mon, 2 Oct 2023 11:02:10 -0500
Subject: [PATCH] Fix drag handle position when frozen column is resized
 (#3352)

Fix drag handle posistion when frozen column is resized
---
 src/DataGrid.tsx                  | 18 +++++++++++++++---
 src/DragHandle.tsx                | 12 ++++++------
 src/hooks/useCalculatedColumns.ts | 11 ++++-------
 3 files changed, 25 insertions(+), 16 deletions(-)

diff --git a/src/DataGrid.tsx b/src/DataGrid.tsx
index f9a1b5b6b8..28e9ef41d8 100644
--- a/src/DataGrid.tsx
+++ b/src/DataGrid.tsx
@@ -283,6 +283,15 @@ function DataGrid<R, SR, K extends Key>(
   const [draggedOverRowIdx, setOverRowIdx] = useState<number | undefined>(undefined);
   const [scrollToPosition, setScrollToPosition] = useState<PartialPosition | null>(null);
 
+  const getColumnWidth = useCallback(
+    (column: CalculatedColumn<R, SR>) => {
+      return (
+        resizedColumnWidths.get(column.key) ?? measuredColumnWidths.get(column.key) ?? column.width
+      );
+    },
+    [measuredColumnWidths, resizedColumnWidths]
+  );
+
   const [gridRef, gridWidth, gridHeight] = useGridDimensions();
   const {
     columns,
@@ -297,8 +306,7 @@ function DataGrid<R, SR, K extends Key>(
   } = useCalculatedColumns({
     rawColumns,
     defaultColumnOptions,
-    measuredColumnWidths,
-    resizedColumnWidths,
+    getColumnWidth,
     scrollLeft,
     viewportWidth: gridWidth,
     enableVirtualization
@@ -830,11 +838,15 @@ function DataGrid<R, SR, K extends Key>(
       return;
     }
 
+    const column = columns[selectedPosition.idx];
+    const columnWidth = getColumnWidth(column);
+
     return (
       <DragHandle
         gridRowStart={headerAndTopSummaryRowsCount + selectedPosition.rowIdx + 1}
         rows={rows}
-        columns={columns}
+        column={column}
+        columnWidth={columnWidth}
         selectedPosition={selectedPosition}
         isCellEditable={isCellEditable}
         latestDraggedOverRowIdx={latestDraggedOverRowIdx}
diff --git a/src/DragHandle.tsx b/src/DragHandle.tsx
index 88c4a6dccb..fc8321515c 100644
--- a/src/DragHandle.tsx
+++ b/src/DragHandle.tsx
@@ -34,7 +34,8 @@ const cellDragHandleClassname = `rdg-cell-drag-handle ${cellDragHandle}`;
 
 interface Props<R, SR> extends Pick<DataGridProps<R, SR>, 'rows' | 'onRowsChange'> {
   gridRowStart: number;
-  columns: readonly CalculatedColumn<R, SR>[];
+  column: CalculatedColumn<R, SR>;
+  columnWidth: number | string;
   selectedPosition: SelectCellState;
   latestDraggedOverRowIdx: React.MutableRefObject<number | undefined>;
   isCellEditable: (position: Position) => boolean;
@@ -47,7 +48,8 @@ interface Props<R, SR> extends Pick<DataGridProps<R, SR>, 'rows' | 'onRowsChange
 export default function DragHandle<R, SR>({
   gridRowStart,
   rows,
-  columns,
+  column,
+  columnWidth,
   selectedPosition,
   latestDraggedOverRowIdx,
   isCellEditable,
@@ -58,7 +60,6 @@ export default function DragHandle<R, SR>({
   setDraggedOverRowIdx
 }: Props<R, SR>) {
   const { idx, rowIdx } = selectedPosition;
-  const column = columns[idx];
 
   function handleMouseDown(event: React.MouseEvent<HTMLDivElement>) {
     // keep the focus on the cell
@@ -99,7 +100,6 @@ export default function DragHandle<R, SR>({
   }
 
   function updateRows(startRowIdx: number, endRowIdx: number) {
-    const column = columns[idx];
     const sourceRow = rows[rowIdx];
     const updatedRows = [...rows];
     const indexes: number[] = [];
@@ -127,8 +127,8 @@ export default function DragHandle<R, SR>({
         ...style,
         gridRowStart,
         insetInlineStart:
-          style.insetInlineStart && typeof column.width === 'number'
-            ? `calc(${style.insetInlineStart} + ${column.width}px - var(--rdg-drag-handle-size))`
+          style.insetInlineStart && typeof columnWidth === 'number'
+            ? `calc(${style.insetInlineStart} + ${columnWidth}px - var(--rdg-drag-handle-size))`
             : undefined
       }}
       className={clsx(cellDragHandleClassname, column.frozen && cellDragHandleFrozenClassname)}
diff --git a/src/hooks/useCalculatedColumns.ts b/src/hooks/useCalculatedColumns.ts
index dfb8ef8363..0bc5b9ff26 100644
--- a/src/hooks/useCalculatedColumns.ts
+++ b/src/hooks/useCalculatedColumns.ts
@@ -32,16 +32,14 @@ interface CalculatedColumnsArgs<R, SR> {
   defaultColumnOptions: DataGridProps<R, SR>['defaultColumnOptions'];
   viewportWidth: number;
   scrollLeft: number;
-  measuredColumnWidths: ReadonlyMap<string, number>;
-  resizedColumnWidths: ReadonlyMap<string, number>;
+  getColumnWidth: (column: CalculatedColumn<R, SR>) => string | number;
   enableVirtualization: boolean;
 }
 
 export function useCalculatedColumns<R, SR>({
   rawColumns,
   defaultColumnOptions,
-  measuredColumnWidths,
-  resizedColumnWidths,
+  getColumnWidth,
   viewportWidth,
   scrollLeft,
   enableVirtualization
@@ -177,8 +175,7 @@ export function useCalculatedColumns<R, SR>({
     const templateColumns: string[] = [];
 
     for (const column of columns) {
-      let width =
-        resizedColumnWidths.get(column.key) ?? measuredColumnWidths.get(column.key) ?? column.width;
+      let width = getColumnWidth(column);
 
       if (typeof width === 'number') {
         width = clampColumnWidth(width, column);
@@ -205,7 +202,7 @@ export function useCalculatedColumns<R, SR>({
     }
 
     return { templateColumns, layoutCssVars, totalFrozenColumnWidth, columnMetrics };
-  }, [measuredColumnWidths, resizedColumnWidths, columns, lastFrozenColumnIndex]);
+  }, [getColumnWidth, columns, lastFrozenColumnIndex]);
 
   const [colOverscanStartIdx, colOverscanEndIdx] = useMemo((): [number, number] => {
     if (!enableVirtualization) {