[DataGrid] Refactor: Use arguments selector for checkbox props (#17683)

Author: MBilalShafi

packages/x-data-grid/src/components/columnSelection/GridCellCheckboxRenderer.tsx

@@ -8,8 +8,8 @@ import { forwardRef } from '@mui/x-internals/forwardRef';
 import { useGridApiContext } from '../../hooks/utils/useGridApiContext';
 import { useGridRootProps } from '../../hooks/utils/useGridRootProps';
 import { getDataGridUtilityClass } from '../../constants/gridClasses';
-import { objectShallowCompare, useGridSelector } from '../../hooks/utils/useGridSelector';
-import { getCheckboxPropsSelector } from '../../hooks/features/rowSelection/utils';
+import { useGridSelector } from '../../hooks/utils/useGridSelector';
+import { checkboxPropsSelector } from '../../hooks/features/rowSelection/utils';
 import type { DataGridProcessedProps } from '../../models/props/DataGridProps';
 import type { GridRowSelectionCheckboxParams } from '../../models/params/gridRowSelectionCheckboxParams';
 import type { GridRenderCellParams } from '../../models/params/gridCellParams';
@@ -89,15 +89,10 @@ const GridCellCheckboxForwardRef = forwardRef<HTMLInputElement, GridRenderCellPa
 
     const isSelectable = apiRef.current.isRowSelectable(id);
 
-    const checkboxPropsSelector = getCheckboxPropsSelector(
-      id,
-      rootProps.rowSelectionPropagation?.parents ?? false,
-    );
-    const { isIndeterminate, isChecked } = useGridSelector(
-      apiRef,
-      checkboxPropsSelector,
-      objectShallowCompare,
-    );
+    const { isIndeterminate, isChecked } = useGridSelector(apiRef, checkboxPropsSelector, {
+      groupId: id,
+      autoSelectParents: rootProps.rowSelectionPropagation?.parents ?? false,
+    });
 
     if (rowNode.type === 'footer' || rowNode.type === 'pinnedRow') {
       return null;

packages/x-data-grid/src/hooks/features/rowSelection/utils.ts

@@ -50,57 +50,55 @@ function getGridRowGroupSelectableDescendants(
   return descendants;
 }
 
-// TODO v8: Use `createSelectorV8`
-export function getCheckboxPropsSelector(groupId: GridRowId, autoSelectParents: boolean) {
-  return createSelector(
-    gridRowTreeSelector,
-    gridSortedRowIdsSelector,
-    gridFilteredRowsLookupSelector,
-    selectedIdsLookupSelector,
-    (rowTree, sortedRowIds, filteredRowsLookup, rowSelectionLookup) => {
-      const groupNode = rowTree[groupId];
-      if (!groupNode || groupNode.type !== 'group') {
-        return {
-          isIndeterminate: false,
-          isChecked: rowSelectionLookup[groupId] === groupId,
-        };
-      }
+export const checkboxPropsSelector = createSelector(
+  gridRowTreeSelector,
+  gridFilteredRowsLookupSelector,
+  gridRowSelectionManagerSelector,
+  (
+    rowTree,
+    filteredRowsLookup,
+    rowSelectionManager,
+    { groupId, autoSelectParents }: { groupId: GridRowId; autoSelectParents: boolean },
+  ) => {
+    const groupNode = rowTree[groupId];
+    if (!groupNode || groupNode.type !== 'group' || rowSelectionManager.has(groupId)) {
+      return {
+        isIndeterminate: false,
+        isChecked: rowSelectionManager.has(groupId),
+      };
+    }
 
-      if (rowSelectionLookup[groupId] === groupId) {
-        return {
-          isIndeterminate: false,
-          isChecked: true,
-        };
-      }
+    let hasSelectedDescendant = false;
+    let hasUnSelectedDescendant = false;
 
-      let selectableDescendantsCount = 0;
-      let selectedDescendantsCount = 0;
-      const startIndex = sortedRowIds.findIndex((id) => id === groupId) + 1;
-      for (
-        let index = startIndex;
-        index < sortedRowIds.length && rowTree[sortedRowIds[index]]?.depth > groupNode.depth;
-        index += 1
+    const traverseDescendants = (itemToTraverseId: GridRowId) => {
+      if (
+        filteredRowsLookup[itemToTraverseId] === false ||
+        // Perf: Skip checking the rest of the descendants if we already
+        // know that there is a selected and an unselected descendant
+        (hasSelectedDescendant && hasUnSelectedDescendant)
       ) {
-        const id = sortedRowIds[index];
-        if (filteredRowsLookup[id] !== false) {
-          selectableDescendantsCount += 1;
-          if (rowSelectionLookup[id] !== undefined) {
-            selectedDescendantsCount += 1;
-          }
-        }
+        return;
       }
-      return {
-        isIndeterminate:
-          selectedDescendantsCount > 0 &&
-          (selectedDescendantsCount < selectableDescendantsCount ||
-            rowSelectionLookup[groupId] === undefined),
-        isChecked: autoSelectParents
-          ? selectedDescendantsCount > 0
-          : rowSelectionLookup[groupId] === groupId,
-      };
-    },
-  );
-}
+      const node = rowTree[itemToTraverseId];
+      if (node?.type === 'group') {
+        node.children.forEach(traverseDescendants);
+      }
+      if (rowSelectionManager.has(itemToTraverseId)) {
+        hasSelectedDescendant = true;
+      } else {
+        hasUnSelectedDescendant = true;
+      }
+    };
+
+    traverseDescendants(groupId);
+
+    return {
+      isIndeterminate: hasSelectedDescendant && hasUnSelectedDescendant,
+      isChecked: autoSelectParents ? hasSelectedDescendant && !hasUnSelectedDescendant : false,
+    };
+  },
+);
 
 export function isMultipleRowSelectionEnabled(
   props: Pick<