From 8f8af6536c0edf08b369133244befc98b803d215 Mon Sep 17 00:00:00 2001
From: Eunjae Lee <hey@eunjae.dev>
Date: Tue, 28 Jan 2025 17:37:06 +0100
Subject: [PATCH] fix: remove stale cache for column visibility state (#18944)

---
 .../filters/ColumnVisibilityButton.tsx        | 15 +--------------
 packages/features/data-table/lib/resizing.ts  | 19 +++++++++++++------
 2 files changed, 14 insertions(+), 20 deletions(-)

diff --git a/packages/features/data-table/components/filters/ColumnVisibilityButton.tsx b/packages/features/data-table/components/filters/ColumnVisibilityButton.tsx
index 9d9613b7be..938db7f65b 100644
--- a/packages/features/data-table/components/filters/ColumnVisibilityButton.tsx
+++ b/packages/features/data-table/components/filters/ColumnVisibilityButton.tsx
@@ -40,9 +40,6 @@ function ColumnVisibilityButtonComponent<TData>(
   const { t } = useLocale();
   const allColumns = table.getAllLeafColumns();
   const [open, setOpen] = useState(false);
-  const [visibleColumns, setVisibleColumns] = useState(
-    () => new Set(allColumns.filter((col) => col.getIsVisible()).map((col) => col.id))
-  );
 
   return (
     <Popover open={open} onOpenChange={setOpen}>
@@ -61,21 +58,12 @@ function ColumnVisibilityButtonComponent<TData>(
                 const canHide = column.getCanHide();
                 if (!column.columnDef.header || typeof column.columnDef.header !== "string" || !canHide)
                   return null;
-                const isVisible = visibleColumns.has(column.id);
+                const isVisible = column.getIsVisible();
                 return (
                   <CommandItem
                     key={column.id}
                     onSelect={() => {
                       column.toggleVisibility(!isVisible);
-                      setVisibleColumns((prev) => {
-                        const next = new Set(prev);
-                        if (isVisible) {
-                          next.delete(column.id);
-                        } else {
-                          next.add(column.id);
-                        }
-                        return next;
-                      });
                     }}>
                     <div
                       className={classNames(
@@ -95,7 +83,6 @@ function ColumnVisibilityButtonComponent<TData>(
             <CommandItem
               onSelect={() => {
                 allColumns.forEach((column) => column.toggleVisibility(true));
-                setVisibleColumns(new Set(allColumns.map((col) => col.id)));
               }}
               className={classNames(
                 "w-full justify-center text-center",
diff --git a/packages/features/data-table/lib/resizing.ts b/packages/features/data-table/lib/resizing.ts
index 3cf5f4dfe3..212803e218 100644
--- a/packages/features/data-table/lib/resizing.ts
+++ b/packages/features/data-table/lib/resizing.ts
@@ -101,16 +101,12 @@ export function usePersistentColumnResizing<TData>({
   const initialColumnSizing = useRef<ColumnSizingState>({});
   const resizedColumns = useRef<Set<string>>(new Set());
 
-  const onColumnSizingChange = useCallback(
+  const adjustColumnSizing = useCallback(
     (updater: ColumnSizingState | ((old: ColumnSizingState) => ColumnSizingState)) => {
       // `!identifier` is checked already in the `useEffect` hook,
       // but TS doesn't know that, and this won't happen.
       if (!identifier) return;
 
-      const isResizingColumn = table.getState().columnSizingInfo.isResizingColumn;
-      if (isResizingColumn) {
-        resizedColumns.current.add(isResizingColumn);
-      }
       table.setState((oldTableState) => {
         let newColumnSizing = typeof updater === "function" ? updater(oldTableState.columnSizing) : updater;
         newColumnSizing = getAdjustedColumnSizing({
@@ -132,7 +128,18 @@ export function usePersistentColumnResizing<TData>({
         };
       });
     },
-    [identifier, table]
+    [table, identifier, tableContainerRef]
+  );
+
+  const onColumnSizingChange = useCallback(
+    (updater: ColumnSizingState | ((old: ColumnSizingState) => ColumnSizingState)) => {
+      const isResizingColumn = table.getState().columnSizingInfo.isResizingColumn;
+      if (isResizingColumn) {
+        resizedColumns.current.add(isResizingColumn);
+      }
+      adjustColumnSizing(updater);
+    },
+    [identifier, table, adjustColumnSizing]
   );
 
   const debouncedContainerWidth = useDebouncedWidth(tableContainerRef);