diff --git a/packages/x-internals/src/lruMemoize/index.ts b/packages/x-internals/src/lruMemoize/index.ts index cce78de9bbac3..2cbbfce30ac56 100644 --- a/packages/x-internals/src/lruMemoize/index.ts +++ b/packages/x-internals/src/lruMemoize/index.ts @@ -1 +1,2 @@ -export { lruMemoize } from 'reselect'; +export { lruMemoize } from './lruMemoize'; +export type { LruMemoizeOptions } from './lruMemoize'; diff --git a/packages/x-internals/src/lruMemoize/lruMemoize.ts b/packages/x-internals/src/lruMemoize/lruMemoize.ts new file mode 100644 index 0000000000000..1b8c836bb48b9 --- /dev/null +++ b/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 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; +} diff --git a/packages/x-internals/src/store/createSelector.ts b/packages/x-internals/src/store/createSelector.ts index 85bbf520a6edf..f079318c5c2f9 100644 --- a/packages/x-internals/src/store/createSelector.ts +++ b/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 & { 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): CreateSelectorFunction => + (options?: MemoizedSelectorOptions): CreateSelectorFunction => (...inputs: any[]) => { type CacheKey = { id: number }; @@ -166,13 +181,13 @@ 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 @@ -180,7 +195,7 @@ export const createSelectorMemoizedWithOptions = 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); diff --git a/packages/x-internals/src/store/weakMapMemoize.ts b/packages/x-internals/src/store/weakMapMemoize.ts new file mode 100644 index 0000000000000..89ab6a7580bda --- /dev/null +++ b/packages/x-internals/src/store/weakMapMemoize.ts @@ -0,0 +1,63 @@ +interface CacheNode { + computed: boolean; + value: unknown; + objects: WeakMap | null; + primitives: Map | 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 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; +}