From 511eed5ab74d050f532c5c6284d390c61605e8d9 Mon Sep 17 00:00:00 2001
From: Sidney Alcantara <sidney@sidney.me>
Date: Thu, 24 Nov 2022 17:27:15 +1100
Subject: [PATCH] fix column widths not updating in virtualized cells

---
 src/components/Table/TableBody.tsx         |  3 ++-
 src/components/Table/useVirtualization.tsx | 22 ++++++++++++++--------
 2 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/src/components/Table/TableBody.tsx b/src/components/Table/TableBody.tsx
index 74efe194..8dd53d73 100644
--- a/src/components/Table/TableBody.tsx
+++ b/src/components/Table/TableBody.tsx
@@ -54,6 +54,7 @@ export const TableBody = memo(function TableBody({
   rows,
   canEditCells,
   lastFrozen,
+  columnSizing,
 }: ITableBodyProps) {
   const [tableSchema] = useAtom(tableSchemaAtom, tableScope);
   const [selectedCell] = useAtom(selectedCellAtom, tableScope);
@@ -66,7 +67,7 @@ export const TableBody = memo(function TableBody({
     paddingBottom,
     paddingLeft,
     paddingRight,
-  } = useVirtualization(containerRef, leafColumns);
+  } = useVirtualization(containerRef, leafColumns, columnSizing);
 
   const rowHeight = tableSchema.rowHeight || DEFAULT_ROW_HEIGHT;
 
diff --git a/src/components/Table/useVirtualization.tsx b/src/components/Table/useVirtualization.tsx
index 90df0329..090c3b83 100644
--- a/src/components/Table/useVirtualization.tsx
+++ b/src/components/Table/useVirtualization.tsx
@@ -15,8 +15,8 @@ import {
   OUT_OF_ORDER_MARGIN,
   DEFAULT_COL_WIDTH,
 } from "./Table";
-import { TableRow } from "@src/types/table";
-import { Column } from "@tanstack/react-table";
+import type { TableRow } from "@src/types/table";
+import type { Column, ColumnSizingState } from "@tanstack/react-table";
 
 import { MIN_COL_WIDTH } from "./Table";
 
@@ -26,7 +26,8 @@ import { MIN_COL_WIDTH } from "./Table";
  */
 export function useVirtualization(
   containerRef: React.RefObject<HTMLDivElement>,
-  leafColumns: Column<TableRow, unknown>[]
+  leafColumns: Column<TableRow, unknown>[],
+  columnSizing: ColumnSizingState
 ) {
   const [tableSchema] = useAtom(tableSchemaAtom, tableScope);
   const [tableRows] = useAtom(tableRowsAtom, tableScope);
@@ -63,12 +64,17 @@ export function useVirtualization(
     paddingStart: TABLE_PADDING,
     paddingEnd: TABLE_PADDING,
     estimateSize: useCallback(
-      (index: number) =>
-        Math.max(
+      (index: number) => {
+        const columnDef = leafColumns[index].columnDef;
+        const schemaWidth = columnDef.size;
+        const localWidth = columnSizing[columnDef.id || ""];
+
+        return Math.max(
           MIN_COL_WIDTH,
-          leafColumns[index].columnDef.size || DEFAULT_COL_WIDTH
-        ),
-      [leafColumns]
+          localWidth || schemaWidth || DEFAULT_COL_WIDTH
+        );
+      },
+      [leafColumns, columnSizing]
     ),
     rangeExtractor: useCallback(
       (range: Range) => {