From 83f2922f41a1d85eb3213dc3d02fd3242ff38f43 Mon Sep 17 00:00:00 2001
From: Rom Grk <romgrk.cc@gmail.com>
Date: Fri, 19 Apr 2024 08:44:12 -0400
Subject: [PATCH] fix: flex column width

---
 .../src/hooks/features/columns/gridColumnsUtils.ts  | 13 +++++++++----
 .../src/hooks/features/columns/useGridColumns.tsx   |  9 +++++----
 2 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/packages/x-data-grid/src/hooks/features/columns/gridColumnsUtils.ts b/packages/x-data-grid/src/hooks/features/columns/gridColumnsUtils.ts
index 94ec07bb341a..9126960acde7 100644
--- a/packages/x-data-grid/src/hooks/features/columns/gridColumnsUtils.ts
+++ b/packages/x-data-grid/src/hooks/features/columns/gridColumnsUtils.ts
@@ -22,6 +22,7 @@ import { GridRowEntry } from '../../../models/gridRows';
 import { gridDensityFactorSelector } from '../density/densitySelector';
 import { gridHeaderFilteringEnabledSelector } from '../headerFiltering/gridHeaderFilteringSelectors';
 import { gridColumnGroupsHeaderMaxDepthSelector } from '../columnGrouping/gridColumnGroupsSelector';
+import type { GridDimensions } from '../dimensions/gridDimensionsApi';
 
 export const COLUMNS_DIMENSION_PROPERTIES = ['maxWidth', 'minWidth', 'width', 'flex'] as const;
 
@@ -161,7 +162,7 @@ export function computeFlexColumnsWidth({
  */
 export const hydrateColumnsWidth = (
   rawState: GridColumnsRawState,
-  viewportInnerWidth: number,
+  dimensions: GridDimensions | undefined,
 ): GridColumnsState => {
   const columnsLookup: GridColumnLookup = {};
   let totalFlexUnits = 0;
@@ -196,10 +197,14 @@ export const hydrateColumnsWidth = (
     columnsLookup[columnField] = newColumn;
   });
 
-  const initialFreeSpace = Math.max(viewportInnerWidth - widthAllocatedBeforeFlex, 0);
+  const availableWidth =
+    dimensions === undefined
+      ? 0
+      : dimensions.viewportOuterSize.width - (dimensions.hasScrollY ? dimensions.scrollbarSize : 0);
+  const initialFreeSpace = Math.max(availableWidth - widthAllocatedBeforeFlex, 0);
 
   // Allocate the remaining space to the flex columns
-  if (totalFlexUnits > 0 && viewportInnerWidth > 0) {
+  if (totalFlexUnits > 0 && availableWidth > 0) {
     const computedColumnWidths = computeFlexColumnsWidth({
       initialFreeSpace,
       totalFlexUnits,
@@ -395,7 +400,7 @@ export const createColumnsState = ({
 
   return hydrateColumnsWidth(
     columnsStateWithPortableColumns,
-    apiRef.current.getRootDimensions?.().viewportInnerSize.width ?? 0,
+    apiRef.current.getRootDimensions?.() ?? undefined,
   );
 };
 
diff --git a/packages/x-data-grid/src/hooks/features/columns/useGridColumns.tsx b/packages/x-data-grid/src/hooks/features/columns/useGridColumns.tsx
index 8b8702126ed9..852597ec4892 100644
--- a/packages/x-data-grid/src/hooks/features/columns/useGridColumns.tsx
+++ b/packages/x-data-grid/src/hooks/features/columns/useGridColumns.tsx
@@ -243,7 +243,7 @@ export function useGridColumns(
               [field]: newColumn,
             },
           },
-          apiRef.current.getRootDimensions().viewportInnerSize.width,
+          apiRef.current.getRootDimensions(),
         ),
       );
 
@@ -386,9 +386,10 @@ export function useGridColumns(
   useGridRegisterPipeProcessor(apiRef, 'restoreState', stateRestorePreProcessing);
   useGridRegisterPipeProcessor(apiRef, 'preferencePanel', preferencePanelPreProcessing);
 
-  /**
+  /*
    * EVENTS
    */
+
   const prevInnerWidth = React.useRef<number | null>(null);
   const handleGridSizeChange: GridEventListener<'viewportInnerSizeChange'> = (
     viewportInnerSize,
@@ -398,7 +399,7 @@ export function useGridColumns(
       setGridColumnsState(
         hydrateColumnsWidth(
           gridColumnsStateSelector(apiRef.current.state),
-          viewportInnerSize.width,
+          apiRef.current.getRootDimensions(),
         ),
       );
     }
@@ -423,7 +424,7 @@ export function useGridColumns(
 
   useGridRegisterPipeApplier(apiRef, 'hydrateColumns', hydrateColumns);
 
-  /**
+  /*
    * EFFECTS
    */
   // The effect do not track any value defined synchronously during the 1st render by hooks called after `useGridColumns`