[core] Replace reselect with local memoizers in x-internals

packages/x-internals/src/lruMemoize/index.ts

@@ -1 +1,2 @@
-export { lruMemoize } from 'reselect';
+export { lruMemoize } from './lruMemoize';
+export type { LruMemoizeOptions } from './lruMemoize';

packages/x-internals/src/lruMemoize/lruMemoize.ts

@@ -0,0 +1,92 @@
+type EqualityFn = (a: any, b: any) => boolean;
+
+export interface LruMemoizeOptions {
+  /**
+   * Function used to compare each new argument against the cached one.
+   * @default (a, b) => a === b
+   */
+  equalityCheck?: EqualityFn;
+  /**
+   * When provided and a new result is computed, the cache is searched for a result
+   * considered equal by this function. If found, the cached value is returned instead,
+   * preserving referential stability.
+   */
+  resultEqualityCheck?: EqualityFn;
+  /**
+   * The number of previous calls to keep in the cache.
+   * @default 1
+   */
+  maxSize?: number;
+}
+
+interface CacheEntry {
+  args: any[];
+  value: any;
+}
+
+const referenceEqualityCheck: EqualityFn = (a, b) => a === b;
+
+/**
+ * Memoizes a function using a least-recently-used cache.
+ * The arguments of the last `maxSize` calls are compared with the provided `equalityCheck`.
+ *
+ * Drop-in replacement for the subset of `lruMemoize` from `reselect` that MUI X relies on.
+ */
+export function lruMemoize<F extends (...args: any[]) => any>(
+  func: F,
+  equalityCheckOrOptions?: EqualityFn | LruMemoizeOptions,
+): F {
+  const options: LruMemoizeOptions =
+    typeof equalityCheckOrOptions === 'object'
+      ? equalityCheckOrOptions
+      : { equalityCheck: equalityCheckOrOptions };
+  const { equalityCheck = referenceEqualityCheck, maxSize = 1, resultEqualityCheck } = options;
+
+  let entries: CacheEntry[] = [];
+
+  const areArgsEqual = (cachedArgs: any[], args: any[]) => {
+    if (cachedArgs.length !== args.length) {
+      return false;
+    }
+    for (let i = 0; i < args.length; i += 1) {
+      if (!equalityCheck(cachedArgs[i], args[i])) {
+        return false;
+      }
+    }
+    return true;
+  };
+
+  function memoized(...args: any[]) {
+    for (let i = 0; i < entries.length; i += 1) {
+      const entry = entries[i];
+      if (areArgsEqual(entry.args, args)) {
+        if (i > 0) {
+          // Move the hit to the front of the cache
+          entries.splice(i, 1);
+          entries.unshift(entry);
+        }
+        return entry.value;
+      }
+    }
+
+    let value = func(...args);
+    if (resultEqualityCheck) {
+      const matchingEntry = entries.find((entry) => resultEqualityCheck(entry.value, value));
+      if (matchingEntry) {
+        value = matchingEntry.value;
+      }
+    }
+
+    entries.unshift({ args, value });
+    if (entries.length > maxSize) {
+      entries.pop();
+    }
+    return value;
+  }
+
+  memoized.clearCache = () => {
+    entries = [];
+  };
+
+  return memoized as unknown as F;
+}

packages/x-internals/src/store/createSelector.ts

@@ -1,24 +1,39 @@
-import {
-  lruMemoize,
-  createSelectorCreator,
-  OverrideMemoizeOptions,
-  UnknownMemoizer,
-} from 'reselect';
+import { lruMemoize, LruMemoizeOptions } from '../lruMemoize/lruMemoize';
+import { weakMapMemoize } from './weakMapMemoize';
 import type { CreateSelectorFunction } from './createSelectorType';
 
 export type { CreateSelectorFunction } from './createSelectorType';
 
 /* eslint-disable no-underscore-dangle */ // __cacheKey__
 
-const reselectCreateSelector = createSelectorCreator({
-  memoize: lruMemoize,
-  memoizeOptions: {
-    maxSize: 1,
-    equalityCheck: Object.is,
-  },
-});
+interface MemoizedSelectorOptions {
+  memoizeOptions?: LruMemoizeOptions;
+}
 
-type SelectorWithArgs = ReturnType<typeof reselectCreateSelector> & { selectorArgs: any[3] };
+/**
+ * Combines input selectors and a combiner into a single memoized selector,
+ * with the same memoization strategy as `reselect`:
+ * the combiner is memoized on its input values (latest call only),
+ * while the selector itself is memoized on its arguments with a weak cache.
+ */
+function createMemoizedSelector(
+  selectorsAndCombiner: Function[],
+  options?: MemoizedSelectorOptions,
+): SelectorWithArgs {
+  const combiner = selectorsAndCombiner[selectorsAndCombiner.length - 1];
+  const inputSelectors = selectorsAndCombiner.slice(0, -1);
+  const memoizedCombiner = lruMemoize(
+    combiner as (...args: any[]) => any,
+    options?.memoizeOptions ?? { equalityCheck: Object.is },
+  );
+
+  return weakMapMemoize((...args: any[]) => {
+    const values = inputSelectors.map((inputSelector) => inputSelector(...args));
+    return memoizedCombiner(...values);
+  }) as unknown as SelectorWithArgs;
+}
+
+type SelectorWithArgs = ((...args: any[]) => any) & { selectorArgs: any[3] };
 
 /**
  * Creates a selector function that can be used to derive values from the store's state.
@@ -134,7 +149,7 @@ export const createSelector = ((
 /* eslint-enable id-denylist */
 
 export const createSelectorMemoizedWithOptions =
-  (options?: OverrideMemoizeOptions<UnknownMemoizer>): CreateSelectorFunction =>
+  (options?: MemoizedSelectorOptions): CreateSelectorFunction =>
   (...inputs: any[]) => {
     type CacheKey = { id: number };
 
@@ -166,21 +181,21 @@ export const createSelectorMemoizedWithOptions =
     let fn = cache.get(cacheKey);
     if (!fn) {
       const selectors = inputs.length === 1 ? [((x: any) => x), combiner] : inputs
-      let reselectArgs = inputs;
+      let selectorsAndCombiner = inputs;
       const selectorArgs = [undefined, undefined, undefined];
       switch (argsLength) {
         case 0:
           break;
         case 1: {
-          reselectArgs = [
+          selectorsAndCombiner = [
             ...selectors.slice(0, -1),
             () => selectorArgs[0],
             combiner
           ];
           break;
         }
         case 2: {
-          reselectArgs = [
+          selectorsAndCombiner = [
             ...selectors.slice(0, -1),
             () => selectorArgs[0],
             () => selectorArgs[1],
@@ -189,7 +204,7 @@ export const createSelectorMemoizedWithOptions =
           break;
         }
         case 3: {
-          reselectArgs = [
+          selectorsAndCombiner = [
             ...selectors.slice(0, -1),
             () => selectorArgs[0],
             () => selectorArgs[1],
@@ -205,11 +220,7 @@ export const createSelectorMemoizedWithOptions =
               'Consider restructuring your selector to use fewer arguments.',
           );
       }
-      if (options) {
-        reselectArgs = [...reselectArgs, options];
-      }
-
-      fn = reselectCreateSelector(...(reselectArgs as any)) as unknown as SelectorWithArgs;
+      fn = createMemoizedSelector(selectorsAndCombiner as Function[], options);
       fn.selectorArgs = selectorArgs;
 
       cache.set(cacheKey, fn);

packages/x-internals/src/store/weakMapMemoize.ts

@@ -0,0 +1,63 @@
+interface CacheNode {
+  computed: boolean;
+  value: unknown;
+  objects: WeakMap<object, CacheNode> | null;
+  primitives: Map<unknown, CacheNode> | null;
+}
+
+const createCacheNode = (): CacheNode => ({
+  computed: false,
+  value: undefined,
+  objects: null,
+  primitives: null,
+});
+
+const isObjectLike = (value: unknown): value is object =>
+  (typeof value === 'object' && value !== null) || typeof value === 'function';
+
+/**
+ * Memoizes a function using a tree of caches keyed by its arguments.
+ * Object arguments are held weakly, so cached results can be garbage collected
+ * together with the arguments that produced them.
+ *
+ * Drop-in replacement for the subset of `weakMapMemoize` from `reselect` that MUI X relies on.
+ */
+export function weakMapMemoize<F extends (...args: any[]) => any>(func: F): F {
+  const root = createCacheNode();
+
+  function memoized(...args: any[]) {
+    let node = root;
+    for (let i = 0; i < args.length; i += 1) {
+      const arg = args[i];
+      if (isObjectLike(arg)) {
+        if (node.objects === null) {
+          node.objects = new WeakMap();
+        }
+        let child = node.objects.get(arg);
+        if (child === undefined) {
+          child = createCacheNode();
+          node.objects.set(arg, child);
+        }
+        node = child;
+      } else {
+        if (node.primitives === null) {
+          node.primitives = new Map();
+        }
+        let child = node.primitives.get(arg);
+        if (child === undefined) {
+          child = createCacheNode();
+          node.primitives.set(arg, child);
+        }
+        node = child;
+      }
+    }
+
+    if (!node.computed) {
+      node.value = func(...args);
+      node.computed = true;
+    }
+    return node.value;
+  }
+
+  return memoized as unknown as F;
+}