[DataGrid] Fix memoized selectors with arguments (#15336) @MBilalShafi

Co-authored-by: Bilal Shafi <[email protected]>

.eslintrc.js

@@ -148,6 +148,11 @@ module.exports = {
         ),
       },
     ],
+    // TODO remove rule from here once it's merged in `@mui/monorepo/.eslintrc`
+    '@typescript-eslint/no-unused-vars': [
+      'error',
+      { vars: 'all', args: 'after-used', ignoreRestSiblings: true, argsIgnorePattern: '^_' },
+    ],
     // TODO move rule into the main repo once it has upgraded
     '@typescript-eslint/return-await': 'off',
     'no-restricted-imports': 'off',

packages/x-data-grid/src/components/GridDetailPanels.tsx

@@ -4,7 +4,6 @@ export interface GridDetailPanelsProps {
   virtualScroller: VirtualScroller;
 }
 
-// eslint-disable-next-line @typescript-eslint/no-unused-vars
 export function GridDetailPanels(_: GridDetailPanelsProps) {
   return null;
 }

packages/x-data-grid/src/components/GridPinnedRows.tsx

@@ -5,7 +5,6 @@ export interface GridPinnedRowsProps {
   virtualScroller: VirtualScroller;
 }
 
-// eslint-disable-next-line @typescript-eslint/no-unused-vars
 export function GridPinnedRows(_: GridPinnedRowsProps) {
   return null;
 }

packages/x-data-grid/src/hooks/utils/useGridSelector.ts

@@ -2,7 +2,7 @@ import * as React from 'react';
 import { fastObjectShallowCompare } from '@mui/x-internals/fastObjectShallowCompare';
 import { warnOnce } from '@mui/x-internals/warning';
 import type { GridApiCommon } from '../../models/api/gridApiCommon';
-import { OutputSelector, OutputSelectorV8 } from '../../utils/createSelector';
+import type { OutputSelector, OutputSelectorV8 } from '../../utils/createSelector';
 import { useLazyRef } from './useLazyRef';
 import { useOnMount } from './useOnMount';
 import type { GridCoreApi } from '../../models/api/gridCoreApi';
@@ -43,8 +43,25 @@ function applySelectorV8<Api extends GridApiCommon, Args, T>(
 
 const defaultCompare = Object.is;
 export const objectShallowCompare = fastObjectShallowCompare;
+const arrayShallowCompare = (a: any[], b: any[]) => {
+  if (a === b) {
+    return true;
+  }
+
+  return a.length === b.length && a.every((v, i) => v === b[i]);
+};
 
-const createRefs = () => ({ state: null, equals: null, selector: null }) as any;
+export const argsEqual = (prev: any, curr: any) => {
+  let fn = Object.is;
+  if (curr instanceof Array) {
+    fn = arrayShallowCompare;
+  } else if (curr instanceof Object) {
+    fn = objectShallowCompare;
+  }
+  return fn(prev, curr);
+};
+
+const createRefs = () => ({ state: null, equals: null, selector: null, args: null }) as any;
 
 // TODO v8: Remove this function
 export const useGridSelector = <Api extends GridApiCommon, T>(
@@ -98,7 +115,7 @@ export const useGridSelectorV8 = <Api extends GridApiCommon, Args, T>(
   apiRef: React.MutableRefObject<Api>,
   selector: Selector<Api, Args, T>,
   args: Args = undefined as Args,
-  equals: (a: T, b: T) => boolean = defaultCompare,
+  equals: <U = T>(a: U, b: U) => boolean = defaultCompare,
 ) => {
   if (process.env.NODE_ENV !== 'production') {
     if (!apiRef.current.state) {
@@ -114,6 +131,7 @@ export const useGridSelectorV8 = <Api extends GridApiCommon, Args, T>(
       state: T;
       equals: typeof equals;
       selector: typeof selector;
+      args: typeof args;
     },
     never
   >(createRefs);
@@ -127,13 +145,28 @@ export const useGridSelectorV8 = <Api extends GridApiCommon, Args, T>(
   refs.current.state = state;
   refs.current.equals = equals;
   refs.current.selector = selector;
+  const prevArgs = refs.current.args;
+  refs.current.args = args;
+
+  if (didInit && !argsEqual(prevArgs, args)) {
+    const newState = applySelectorV8(
+      apiRef,
+      refs.current.selector,
+      refs.current.args,
+      apiRef.current.instanceId,
+    ) as T;
+    if (!refs.current.equals(refs.current.state, newState)) {
+      refs.current.state = newState;
+      setState(newState);
+    }
+  }
 
   useOnMount(() => {
     return apiRef.current.store.subscribe(() => {
       const newState = applySelectorV8(
         apiRef,
         refs.current.selector,
-        args,
+        refs.current.args,
         apiRef.current.instanceId,
       ) as T;
       if (!refs.current.equals(refs.current.state, newState)) {

packages/x-data-grid/src/utils/createSelector.ts

@@ -2,6 +2,7 @@ import * as React from 'react';
 import { lruMemoize, createSelectorCreator, Selector, SelectorResultArray } from 'reselect';
 import { warnOnce } from '@mui/x-internals/warning';
 import type { GridCoreApi } from '../models/api/gridCoreApi';
+import { argsEqual } from '../hooks/utils/useGridSelector';
 
 type CacheKey = { id: number };
 
@@ -13,6 +14,10 @@ const reselectCreateSelector = createSelectorCreator({
   },
 });
 
+type GridCreateSelectorFunction = ReturnType<typeof reselectCreateSelector> & {
+  selectorArgs?: any;
+};
+
 // TODO v8: Remove this type
 export interface OutputSelector<State, Result> {
   (apiRef: React.MutableRefObject<{ state: State; instanceId: GridCoreApi['instanceId'] }>): Result;
@@ -24,7 +29,7 @@ export interface OutputSelector<State, Result> {
 export interface OutputSelectorV8<State, Args, Result> {
   (
     apiRef: React.MutableRefObject<{ state: State; instanceId: GridCoreApi['instanceId'] }>,
-    args: Args,
+    args?: Args,
   ): Result;
   (state: State, instanceId: GridCoreApi['instanceId']): Result;
   acceptsApiRef: boolean;
@@ -323,12 +328,28 @@ export const createSelectorMemoizedV8: CreateSelectorFunctionV8 = (...args: any)
     const cacheFn = cacheArgs?.get(args);
 
     if (cacheArgs && cacheFn) {
+      if (!argsEqual(cacheFn.selectorArgs, selectorArgs)) {
+        const reselectArgs =
+          selectorArgs !== undefined
+            ? [...args.slice(0, args.length - 1), () => selectorArgs, args[args.length - 1]]
+            : args;
+        const fn: GridCreateSelectorFunction = reselectCreateSelector(...reselectArgs);
+        fn.selectorArgs = selectorArgs;
+        cacheArgs.set(args, fn);
+        return fn(state, selectorArgs, cacheKey);
+      }
       // We pass the cache key because the called selector might have as
       // dependency another selector created with this `createSelector`.
       return cacheFn(state, selectorArgs, cacheKey);
     }
 
-    const fn = reselectCreateSelector(...args);
+    const reselectArgs =
+      selectorArgs !== undefined
+        ? [...args.slice(0, args.length - 1), () => selectorArgs, args[args.length - 1]]
+        : args;
+
+    const fn: GridCreateSelectorFunction = reselectCreateSelector(...reselectArgs);
+    fn.selectorArgs = selectorArgs;
 
     if (!cacheArgsInit) {
       cache.set(cacheKey, cacheArgs);

packages/x-data-grid/src/utils/utils.ts

@@ -209,11 +209,9 @@ export function deepClone(obj: Record<string, any>) {
   return JSON.parse(JSON.stringify(obj));
 }
 
-/* eslint-disable @typescript-eslint/no-unused-vars */
 /**
  * Mark a value as used so eslint doesn't complain. Use this instead
  * of a `eslint-disable-next-line react-hooks/exhaustive-deps` because
  * that hint disables checks on all values instead of just one.
  */
 export function eslintUseValue(_: any) {}
-/* eslint-enable @typescript-eslint/no-unused-vars */