From 66f3c4f797f1fd88cca9023ceeebe7781f4eaf41 Mon Sep 17 00:00:00 2001
From: Andrew Cherniavskii <andrew@mui.com>
Date: Fri, 1 Mar 2024 15:30:54 +0100
Subject: [PATCH] [DataGridPro] Fix lazy loading error when existing rows are
 passed to replaceRows (@martijn-basesoft) (#12272)

Co-authored-by: martijn-basesoft <martijn@basesoft.com>
Co-authored-by: Rom Grk <romgrk.cc@gmail.com>
---
 .../src/tests/lazyLoader.DataGridPro.test.tsx | 22 +++++++++++++++++++
 .../src/hooks/features/rows/useGridRows.ts    | 13 +++++++----
 2 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/packages/grid/x-data-grid-pro/src/tests/lazyLoader.DataGridPro.test.tsx b/packages/grid/x-data-grid-pro/src/tests/lazyLoader.DataGridPro.test.tsx
index b144e2d8640f..8cb39c7b55b0 100644
--- a/packages/grid/x-data-grid-pro/src/tests/lazyLoader.DataGridPro.test.tsx
+++ b/packages/grid/x-data-grid-pro/src/tests/lazyLoader.DataGridPro.test.tsx
@@ -171,4 +171,26 @@ describe('<DataGridPro /> - Lazy loader', () => {
     expect(apiRef.current.getRowNode(4)).to.not.equal(null);
     expect(apiRef.current.getRowNode(5)).to.not.equal(null);
   });
+
+  it('should update rows when `apiRef.current.updateRows` with data reversed', () => {
+    render(<TestLazyLoader rowCount={5} autoHeight={isJSDOM} />);
+
+    const newRows: GridRowModel[] = [
+      {
+        id: 3,
+        first: 'Jim',
+      },
+      {
+        id: 2,
+        first: 'Jack',
+      },
+      {
+        id: 1,
+        first: 'Mike',
+      },
+    ];
+
+    act(() => apiRef.current.unstable_replaceRows(0, newRows));
+    expect(getColumnValues(1)).to.deep.equal(['Jim', 'Jack', 'Mike']);
+  });
 });
diff --git a/packages/grid/x-data-grid/src/hooks/features/rows/useGridRows.ts b/packages/grid/x-data-grid/src/hooks/features/rows/useGridRows.ts
index 2562e1c3d932..1c7d90c55463 100644
--- a/packages/grid/x-data-grid/src/hooks/features/rows/useGridRows.ts
+++ b/packages/grid/x-data-grid/src/hooks/features/rows/useGridRows.ts
@@ -410,6 +410,7 @@ export const useGridRows = (
       const dataRowIdToIdLookup = { ...gridRowsDataRowIdToIdLookupSelector(apiRef) };
       const rootGroup = tree[GRID_ROOT_GROUP_ID] as GridGroupNode;
       const rootGroupChildren = [...rootGroup.children];
+      const seenIds = new Set<GridRowId>();
 
       for (let i = 0; i < newRows.length; i += 1) {
         const rowModel = newRows[i];
@@ -419,11 +420,13 @@ export const useGridRows = (
           'A row was provided without id when calling replaceRows().',
         );
 
-        const [replacedRowId] = rootGroupChildren.splice(firstRowToRender + i, 1, rowId);
+        const [removedRowId] = rootGroupChildren.splice(firstRowToRender + i, 1, rowId);
 
-        delete dataRowIdToModelLookup[replacedRowId];
-        delete dataRowIdToIdLookup[replacedRowId];
-        delete tree[replacedRowId];
+        if (!seenIds.has(removedRowId)) {
+          delete dataRowIdToModelLookup[removedRowId];
+          delete dataRowIdToIdLookup[removedRowId];
+          delete tree[removedRowId];
+        }
 
         const rowTreeNodeConfig: GridLeafNode = {
           id: rowId,
@@ -435,6 +438,8 @@ export const useGridRows = (
         dataRowIdToModelLookup[rowId] = rowModel;
         dataRowIdToIdLookup[rowId] = rowId;
         tree[rowId] = rowTreeNodeConfig;
+
+        seenIds.add(rowId);
       }
 
       tree[GRID_ROOT_GROUP_ID] = { ...rootGroup, children: rootGroupChildren };