From 44d13578d6dcdceea314eea93cdae86e946a7c50 Mon Sep 17 00:00:00 2001
From: Armin Mehinovic <4390250+arminmeh@users.noreply.github.com>
Date: Thu, 31 Oct 2024 19:42:26 +0100
Subject: [PATCH] [DataGridPremium] Fix incorrect rows selection count when
 selection propagation is enabled with row grouping (#15216)

---
 .../rowSelection.DataGridPremium.test.tsx     | 20 +++++++++++++-
 .../rowSelection/useGridRowSelection.ts       | 26 ++++++++-----------
 2 files changed, 30 insertions(+), 16 deletions(-)

diff --git a/packages/x-data-grid-premium/src/tests/rowSelection.DataGridPremium.test.tsx b/packages/x-data-grid-premium/src/tests/rowSelection.DataGridPremium.test.tsx
index d330b84ebd8c..88bb7b1dd69e 100644
--- a/packages/x-data-grid-premium/src/tests/rowSelection.DataGridPremium.test.tsx
+++ b/packages/x-data-grid-premium/src/tests/rowSelection.DataGridPremium.test.tsx
@@ -1,5 +1,5 @@
 import * as React from 'react';
-import { createRenderer, fireEvent } from '@mui/internal-test-utils';
+import { act, createRenderer, fireEvent } from '@mui/internal-test-utils';
 import { getCell } from 'test/utils/helperFn';
 import { expect } from 'chai';
 import {
@@ -122,6 +122,24 @@ describe('<DataGridPremium /> - Row selection', () => {
       expect(apiRef.current.getSelectedRows()).to.have.keys([0, 2]);
     });
 
+    // Context: https://github.com/mui/mui-x/issues/15206
+    it('should keep the correct selection items and the selection count when rows are updated', () => {
+      render(<Test />);
+
+      const expectedKeys = ['auto-generated-row-category1/Cat B', 3, 4];
+      const expectedCount = 3;
+
+      fireEvent.click(getCell(1, 0).querySelector('input')!);
+      expect(apiRef.current.getSelectedRows()).to.have.keys(expectedKeys);
+      expect(apiRef.current.state.rowSelection.length).to.equal(expectedCount);
+
+      act(() => {
+        apiRef.current.updateRows([...rows]);
+      });
+      expect(apiRef.current.getSelectedRows()).to.have.keys(expectedKeys);
+      expect(apiRef.current.state.rowSelection.length).to.equal(expectedCount);
+    });
+
     describe("prop: indeterminateCheckboxAction = 'select'", () => {
       it('should select all the children when selecting an indeterminate parent', () => {
         render(
diff --git a/packages/x-data-grid/src/hooks/features/rowSelection/useGridRowSelection.ts b/packages/x-data-grid/src/hooks/features/rowSelection/useGridRowSelection.ts
index 2c8515776a13..fb8b76b078a4 100644
--- a/packages/x-data-grid/src/hooks/features/rowSelection/useGridRowSelection.ts
+++ b/packages/x-data-grid/src/hooks/features/rowSelection/useGridRowSelection.ts
@@ -263,6 +263,7 @@ export const useGridRowSelection = (
         const removeRow = (rowId: GridRowId) => {
           newSelection.delete(rowId);
         };
+
         if (isSelected) {
           addRow(id);
           if (applyAutoSelection) {
@@ -309,13 +310,13 @@ export const useGridRowSelection = (
 
       const selectableIds = ids.filter((id) => apiRef.current.isRowSelectable(id));
 
-      let newSelection: GridRowId[];
+      let newSelection: Set<GridRowId>;
       if (resetSelection) {
         if (isSelected) {
-          newSelection = selectableIds;
+          newSelection = new Set(selectableIds);
           if (applyAutoSelection) {
             const addRow = (rowId: GridRowId) => {
-              newSelection.push(rowId);
+              newSelection.add(rowId);
             };
             selectableIds.forEach((id) => {
               findRowsToSelect(
@@ -329,23 +330,20 @@ export const useGridRowSelection = (
             });
           }
         } else {
-          newSelection = [];
+          newSelection = new Set();
         }
       } else {
-        // We clone the existing object to avoid mutating the same object returned by the selector to others part of the project
-        const selectionLookup = {
-          ...selectedIdsLookupSelector(apiRef),
-        };
+        newSelection = new Set(Object.values(selectedIdsLookupSelector(apiRef)));
         const addRow = (rowId: GridRowId) => {
-          selectionLookup[rowId] = rowId;
+          newSelection.add(rowId);
         };
         const removeRow = (rowId: GridRowId) => {
-          delete selectionLookup[rowId];
+          newSelection.delete(rowId);
         };
 
         selectableIds.forEach((id) => {
           if (isSelected) {
-            selectionLookup[id] = id;
+            newSelection.add(id);
             if (applyAutoSelection) {
               findRowsToSelect(
                 apiRef,
@@ -370,13 +368,11 @@ export const useGridRowSelection = (
             }
           }
         });
-
-        newSelection = Object.values(selectionLookup);
       }
 
-      const isSelectionValid = newSelection.length < 2 || canHaveMultipleSelection;
+      const isSelectionValid = newSelection.size < 2 || canHaveMultipleSelection;
       if (isSelectionValid) {
-        apiRef.current.setRowSelectionModel(newSelection);
+        apiRef.current.setRowSelectionModel(Array.from(newSelection));
       }
     },
     [