From a0716b427beb5440193d3bc49f490d764e48d897 Mon Sep 17 00:00:00 2001
From: Bilal Shafi <bilalshafidev@gmail.com>
Date: Fri, 13 Dec 2024 14:45:19 +0500
Subject: [PATCH] [DataGrid] Trigger row spanning computation on rows update
 (#15858)

---
 .../hooks/features/rows/useGridRowSpanning.ts | 20 ++++++++----
 .../src/tests/rowSpanning.DataGrid.test.tsx   | 32 +++++++++++++++++++
 2 files changed, 46 insertions(+), 6 deletions(-)

diff --git a/packages/x-data-grid/src/hooks/features/rows/useGridRowSpanning.ts b/packages/x-data-grid/src/hooks/features/rows/useGridRowSpanning.ts
index 1c29902b6e791..8ee70db58ad24 100644
--- a/packages/x-data-grid/src/hooks/features/rows/useGridRowSpanning.ts
+++ b/packages/x-data-grid/src/hooks/features/rows/useGridRowSpanning.ts
@@ -5,6 +5,7 @@ import { gridVisibleColumnDefinitionsSelector } from '../columns/gridColumnsSele
 import { useGridVisibleRows } from '../../utils/useGridVisibleRows';
 import { gridRenderContextSelector } from '../virtualization/gridVirtualizationSelectors';
 import { useGridSelector } from '../../utils/useGridSelector';
+import { gridRowTreeSelector } from './gridRowsSelector';
 import type { GridColDef } from '../../../models/colDef';
 import type { GridRowId, GridValidRowModel, GridRowEntry } from '../../../models/gridRows';
 import type { DataGridProcessedProps } from '../../../models/props/DataGridProps';
@@ -71,7 +72,7 @@ const computeRowSpanningState = (
 
     for (
       let index = rangeToProcess.firstRowIndex;
-      index <= rangeToProcess.lastRowIndex;
+      index < rangeToProcess.lastRowIndex;
       index += 1
     ) {
       const row = visibleRows[index];
@@ -187,7 +188,7 @@ export const rowSpanningStateInitializer: GridStateInitializer = (state, props,
     }
     const rangeToProcess = {
       firstRowIndex: 0,
-      lastRowIndex: Math.min(DEFAULT_ROWS_TO_PROCESS - 1, Math.max(rowIds.length - 1, 0)),
+      lastRowIndex: Math.min(DEFAULT_ROWS_TO_PROCESS, Math.max(rowIds.length, 0)),
     };
     const rows = rowIds.map((id) => ({
       id,
@@ -226,13 +227,14 @@ export const useGridRowSpanning = (
   const { range, rows: visibleRows } = useGridVisibleRows(apiRef, props);
   const renderContext = useGridSelector(apiRef, gridRenderContextSelector);
   const colDefs = useGridSelector(apiRef, gridVisibleColumnDefinitionsSelector);
+  const tree = useGridSelector(apiRef, gridRowTreeSelector);
   const processedRange = useLazyRef<RowRange, void>(() => {
     return Object.keys(apiRef.current.state.rowSpanning.spannedCells).length > 0
       ? {
           firstRowIndex: 0,
           lastRowIndex: Math.min(
-            DEFAULT_ROWS_TO_PROCESS - 1,
-            Math.max(apiRef.current.state.rows.dataRowIds.length - 1, 0),
+            DEFAULT_ROWS_TO_PROCESS,
+            Math.max(apiRef.current.state.rows.dataRowIds.length, 0),
           ),
         }
       : EMPTY_RANGE;
@@ -265,7 +267,7 @@ export const useGridRowSpanning = (
       const rangeToProcess = getUnprocessedRange(
         {
           firstRowIndex: renderContext.firstRowIndex,
-          lastRowIndex: Math.min(renderContext.lastRowIndex - 1, range.lastRowIndex),
+          lastRowIndex: Math.min(renderContext.lastRowIndex, range.lastRowIndex + 1),
         },
         processedRange.current,
       );
@@ -326,11 +328,17 @@ export const useGridRowSpanning = (
   const prevRenderContext = React.useRef(renderContext);
   const isFirstRender = React.useRef(true);
   const shouldResetState = React.useRef(false);
+  const previousTree = React.useRef(tree);
   React.useEffect(() => {
     const firstRender = isFirstRender.current;
     if (isFirstRender.current) {
       isFirstRender.current = false;
     }
+    if (tree !== previousTree.current) {
+      previousTree.current = tree;
+      updateRowSpanningState(true);
+      return;
+    }
     if (range && lastRange.current && isRowRangeUpdated(range, lastRange.current)) {
       lastRange.current = range;
       shouldResetState.current = true;
@@ -344,5 +352,5 @@ export const useGridRowSpanning = (
       return;
     }
     updateRowSpanningState();
-  }, [updateRowSpanningState, renderContext, range, lastRange]);
+  }, [updateRowSpanningState, renderContext, range, lastRange, tree]);
 };
diff --git a/packages/x-data-grid/src/tests/rowSpanning.DataGrid.test.tsx b/packages/x-data-grid/src/tests/rowSpanning.DataGrid.test.tsx
index ef0b006133d1c..f89153e6c0b4e 100644
--- a/packages/x-data-grid/src/tests/rowSpanning.DataGrid.test.tsx
+++ b/packages/x-data-grid/src/tests/rowSpanning.DataGrid.test.tsx
@@ -1,6 +1,7 @@
 import * as React from 'react';
 import { createRenderer, fireEvent, act } from '@mui/internal-test-utils';
 import { expect } from 'chai';
+import { spy } from 'sinon';
 import { DataGrid, useGridApiRef, DataGridProps, GridApi } from '@mui/x-data-grid';
 import { getCell, getActiveCell } from 'test/utils/helperFn';
 
@@ -247,5 +248,36 @@ describe('<DataGrid /> - Row spanning', () => {
     });
   });
 
+  describe('rows update', () => {
+    it('should update the row spanning state when the rows are updated', () => {
+      const rowSpanValueGetter = spy();
+      let rowSpanningStateUpdates = 0;
+      let spannedCells = {};
+      render(
+        <TestDataGrid
+          columns={[{ field: 'code', rowSpanValueGetter }]}
+          rows={[{ id: 1, code: 'A101' }]}
+          onStateChange={(newState) => {
+            const newSpannedCells = newState.rowSpanning.spannedCells;
+            if (newSpannedCells !== spannedCells) {
+              rowSpanningStateUpdates += 1;
+              spannedCells = newSpannedCells;
+            }
+          }}
+        />,
+      );
+
+      // First update by initializer
+      expect(rowSpanningStateUpdates).to.equal(1);
+
+      act(() => {
+        apiRef.current.setRows([{ id: 1, code: 'A101' }]);
+      });
+
+      // Second update on row update
+      expect(rowSpanningStateUpdates).to.equal(2);
+    });
+  });
+
   // TODO: Add tests for row reordering
 });