Merge pull request #46280 from callstack-internal/fix/memoize-constru…

…ctable-types

[No QA][TS] Memoize constructable types

src/libs/memoize/index.ts

@@ -1,9 +1,9 @@
 /* eslint-disable @typescript-eslint/no-unsafe-return */
-import type {Constructor} from 'type-fest';
+import type NonPartial from '@src/types/utils/NonPartial';
 import type {TakeFirst} from '@src/types/utils/TupleOperations';
 import ArrayCache from './cache/ArrayCache';
 import {MemoizeStats} from './stats';
-import type {ClientOptions, MemoizedFn, MemoizeFnPredicate, Stats} from './types';
+import type {Callable, ClientOptions, Constructable, IsomorphicFn, IsomorphicParameters, IsomorphicReturnType, MemoizedFn, Stats} from './types';
 import {getEqualityComparator, mergeOptions, truncateArgs} from './utils';
 
 /**
@@ -43,14 +43,18 @@ class Memoize {
  * @param opts - Options for the memoization layer, for more details see `ClientOptions` type.
  * @returns Memoized function with a cache API attached to it.
  */
-function memoize<Fn extends MemoizeFnPredicate, MaxArgs extends number = Parameters<Fn>['length'], Key = TakeFirst<Parameters<Fn>, MaxArgs>>(fn: Fn, opts?: ClientOptions<Fn, MaxArgs, Key>) {
+function memoize<Fn extends IsomorphicFn, MaxArgs extends number = NonPartial<IsomorphicParameters<Fn>>['length'], Key = TakeFirst<IsomorphicParameters<Fn>, MaxArgs>>(
+    fn: Fn,
+    opts?: ClientOptions<Fn, MaxArgs, Key>,
+) {
     const options = mergeOptions<Fn, MaxArgs, Key>(opts);
 
-    const cache = ArrayCache<Key, ReturnType<Fn>>({maxSize: options.maxSize, keyComparator: getEqualityComparator(options)});
+    const cache = ArrayCache<Key, IsomorphicReturnType<Fn>>({maxSize: options.maxSize, keyComparator: getEqualityComparator(options)});
 
     const stats = new MemoizeStats(options.monitor || Memoize.isMonitoringEnabled);
 
-    const memoized = function memoized(...args: Parameters<Fn>): ReturnType<Fn> {
+    const memoized = function memoized(...args: IsomorphicParameters<Fn>): IsomorphicReturnType<Fn> {
+        // Detect if memoized function was called with `new` keyword. If so we need to call the original function as constructor.
         const constructable = !!new.target;
 
         const truncatedArgs = truncateArgs(args, options.maxArgs);
@@ -63,9 +67,9 @@ function memoize<Fn extends MemoizeFnPredicate, MaxArgs extends number = Paramet
         const retrievalTimeStart = performance.now();
         const cached = cache.getSet(key, () => {
             const fnTimeStart = performance.now();
-            // If the function is constructable, we need to call it with the `new` keyword
-            // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
-            const result = constructable ? new (fn as unknown as Constructor<ReturnType<Fn>, Parameters<Fn>>)(...args) : fn(...args);
+
+            const result = (constructable ? new (fn as Constructable)(...args) : (fn as Callable)(...args)) as IsomorphicReturnType<Fn>;
+
             statsEntry.trackTime('fnTime', fnTimeStart);
             statsEntry.track('didHit', false);
 
@@ -78,7 +82,7 @@ function memoize<Fn extends MemoizeFnPredicate, MaxArgs extends number = Paramet
         statsEntry.save();
 
         return cached.value;
-    };
+    } as MemoizedFn<Fn, Key>;
 
     /**
      * Cache API attached to the memoized function. Currently there is an issue with typing cache keys, but the functionality works as expected.
@@ -90,7 +94,7 @@ function memoize<Fn extends MemoizeFnPredicate, MaxArgs extends number = Paramet
 
     Memoize.registerMemoized(options.monitoringName ?? fn.name, memoized);
 
-    return memoized as MemoizedFn<Fn, Key>;
+    return memoized;
 }
 
 export default memoize;

src/libs/memoize/types.ts

@@ -3,15 +3,23 @@ import type {Cache} from './cache/types';
 import type {MemoizeStats} from './stats';
 
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
-type MemoizeFnPredicate = (...args: any[]) => any;
+type Callable = (...args: any[]) => any;
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+type Constructable = new (...args: any[]) => any;
+
+type IsomorphicFn = Callable | Constructable;
+
+type IsomorphicParameters<Fn extends IsomorphicFn> = Fn extends Callable ? Parameters<Fn> : Fn extends Constructable ? ConstructorParameters<Fn> : never;
+
+type IsomorphicReturnType<Fn extends IsomorphicFn> = Fn extends Callable ? ReturnType<Fn> : Fn extends Constructable ? Fn : never;
 
 type KeyComparator<Key> = (k1: Key, k2: Key) => boolean;
 
 type InternalOptions = {
     cache: 'array';
 };
 
-type Options<Fn extends MemoizeFnPredicate, MaxArgs extends number, Key> = {
+type Options<Fn extends IsomorphicFn, MaxArgs extends number, Key> = {
     maxSize: number;
     equality: 'deep' | 'shallow' | KeyComparator<Key>;
     monitor: boolean;
@@ -24,15 +32,15 @@ type Options<Fn extends MemoizeFnPredicate, MaxArgs extends number, Key> = {
      * @param truncatedArgs - Tuple of arguments passed to the memoized function (truncated to `maxArgs`). Does not work with constructable (see description).
      * @returns - Key to use for caching
      */
-    transformKey?: (truncatedArgs: TakeFirst<Parameters<Fn>, MaxArgs>) => Key;
+    transformKey?: (truncatedArgs: TakeFirst<IsomorphicParameters<Fn>, MaxArgs>) => Key;
 } & InternalOptions;
 
-type ClientOptions<Fn extends MemoizeFnPredicate, MaxArgs extends number, Key> = Partial<Omit<Options<Fn, MaxArgs, Key>, keyof InternalOptions>>;
+type ClientOptions<Fn extends IsomorphicFn, MaxArgs extends number, Key> = Partial<Omit<Options<Fn, MaxArgs, Key>, keyof InternalOptions>>;
 
 type Stats = Pick<MemoizeStats, 'startMonitoring' | 'stopMonitoring'>;
 
-type MemoizedFn<Fn extends MemoizeFnPredicate, Key> = Fn & {
-    cache: Cache<Key, ReturnType<Fn>>;
+type MemoizedFn<Fn extends IsomorphicFn, Key> = Fn & {
+    cache: Cache<Key, IsomorphicReturnType<Fn>>;
 } & Stats;
 
-export type {Options, ClientOptions, MemoizeFnPredicate, Stats, KeyComparator, MemoizedFn};
+export type {Options, ClientOptions, IsomorphicFn, IsomorphicParameters, IsomorphicReturnType, Stats, KeyComparator, MemoizedFn, Callable, Constructable};

src/libs/memoize/utils.ts

@@ -1,9 +1,9 @@
 import {deepEqual, shallowEqual} from 'fast-equals';
 import type {TakeFirst} from '@src/types/utils/TupleOperations';
 import DEFAULT_OPTIONS from './const';
-import type {ClientOptions, KeyComparator, MemoizeFnPredicate, Options} from './types';
+import type {ClientOptions, IsomorphicFn, KeyComparator, Options} from './types';
 
-function getEqualityComparator<Fn extends MemoizeFnPredicate, MaxArgs extends number, Key>(opts: Options<Fn, MaxArgs, Key>): KeyComparator<Key> {
+function getEqualityComparator<Fn extends IsomorphicFn, MaxArgs extends number, Key>(opts: Options<Fn, MaxArgs, Key>): KeyComparator<Key> {
     // Use the custom equality comparator if it is provided
     if (typeof opts.equality === 'function') {
         return opts.equality;
@@ -16,13 +16,13 @@ function getEqualityComparator<Fn extends MemoizeFnPredicate, MaxArgs extends nu
     return deepEqual;
 }
 
-function mergeOptions<Fn extends MemoizeFnPredicate, MaxArgs extends number, Key>(options?: ClientOptions<Fn, MaxArgs, Key>): Options<Fn, MaxArgs, Key> {
+function mergeOptions<Fn extends IsomorphicFn, MaxArgs extends number, Key>(options?: ClientOptions<Fn, MaxArgs, Key>): Options<Fn, MaxArgs, Key> {
     if (!options) {
         return DEFAULT_OPTIONS;
     }
     return {...DEFAULT_OPTIONS, ...options};
 }
-function truncateArgs<T extends readonly unknown[], MaxArgs extends number = T['length']>(args: T, maxArgs?: MaxArgs): TakeFirst<T, MaxArgs> {
+function truncateArgs<T extends unknown[], MaxArgs extends number = T['length']>(args: T, maxArgs?: MaxArgs): TakeFirst<T, MaxArgs> {
     // Hot paths are declared explicitly to avoid the overhead of the slice method
 
     if (maxArgs === undefined) {
@@ -34,11 +34,11 @@ function truncateArgs<T extends readonly unknown[], MaxArgs extends number = T['
     }
 
     if (maxArgs === 0) {
-        return [];
+        return [] as unknown as TakeFirst<T, MaxArgs>;
     }
 
     if (maxArgs === 1) {
-        return [args[0]];
+        return [args[0]] as unknown as TakeFirst<T, MaxArgs>;
     }
 
     if (maxArgs === 2) {

src/types/utils/NonPartial.ts

@@ -0,0 +1,7 @@
+/**
+ * Utility type to make all properties of an object non-optional with `undefined` as a possible value. It fixes the issue with `Required` utility type, which makes all properties required.
+ * See https://github.com/microsoft/TypeScript/issues/31025 for more details.
+ */
+type NonPartial<T, RT extends T = Required<T>> = {[K in keyof RT]: K extends keyof T ? (undefined extends T[K] ? T[K] | undefined : T[K]) : never};
+
+export default NonPartial;

src/types/utils/TupleOperations.ts

@@ -1,12 +1,26 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
+import type NonPartial from './NonPartial';
+
+/**
+ * Fill a tuple with `N` elements from the rest of the tuple.
+ */
+type FillFromRest<T extends any[], N extends number, O extends any[] = []> = O['length'] extends N
+    ? O
+    : T extends [infer Head, ...infer Tail]
+    ? FillFromRest<Tail, N, [...O, Head]>
+    : T extends [...infer Rest]
+    ? FillFromRest<Rest, N, [...O, Rest[number]]>
+    : O;
 
 // Based on: https://stackoverflow.com/questions/67605122/obtain-a-slice-of-a-typescript-parameters-tuple
-type TupleSplit<T, N extends number, O extends readonly any[] = readonly []> = O['length'] extends N
-    ? [O, T]
-    : T extends readonly [infer F, ...infer R]
-    ? TupleSplit<readonly [...R], N, readonly [...O, F]>
-    : [O, T];
+/**
+ * Split a tuple into two parts: the first `N` elements and the rest.
+ */
+type TupleSplit<T, N extends number, O extends any[] = []> = O['length'] extends N ? [O, T] : T extends [infer F, ...infer R] ? TupleSplit<[...R], N, [...O, F]> : [O, T];
 
-type TakeFirst<T extends readonly any[], N extends number = T['length']> = TupleSplit<T, N>[0];
+/**
+ * Get the first `N` elements of a tuple. If `N` is not provided, it returns the whole tuple.
+ */
+type TakeFirst<T extends any[], N extends number = T['length']> = number extends N ? T : TupleSplit<NonPartial<FillFromRest<T, N>>, N>[0];
 
 export type {TupleSplit, TakeFirst};