diff --git a/lib/helpers.ts b/lib/helpers.ts index 49ade39..198a758 100644 --- a/lib/helpers.ts +++ b/lib/helpers.ts @@ -5,23 +5,39 @@ import path from 'node:path'; export const XCRUN_TIMEOUT = 15000; +type AnyFunction = (this: any, ...args: any[]) => any; + /** - * Memoizes function calls by caching results for serialized argument lists. + * Memoizes function calls by traversing a nested argument-keyed map. * * @param fn The function to memoize * @returns A memoized wrapper around the input function */ -export function memoize( - fn: (...args: Args) => Result, -): (...args: Args) => Result { - const cache = new Map(); - return (...args: Args): Result => { - const key = JSON.stringify(args); - if (!cache.has(key)) { - cache.set(key, fn(...args)); +export function memoize(fn: F): F { + const rootCache = new Map(); + const RESULT = Symbol('memoize.result'); + + function memoized(this: ThisParameterType, ...args: Parameters): ReturnType { + let currentCache = rootCache; + + for (const arg of args) { + if (!currentCache.has(arg)) { + currentCache.set(arg, new Map()); + } + currentCache = currentCache.get(arg); + } + + if (currentCache.has(RESULT)) { + return currentCache.get(RESULT); } - return cache.get(key) as Result; - }; + + const result = fn.apply(this, args); + currentCache.set(RESULT, result); + + return result; + } + + return memoized as F; } /** diff --git a/test/unit/helpers-specs.ts b/test/unit/helpers-specs.ts index 7366388..83febf2 100644 --- a/test/unit/helpers-specs.ts +++ b/test/unit/helpers-specs.ts @@ -32,5 +32,44 @@ describe('helpers', function () { expect(result2).to.equal(5); expect(callCount).to.equal(2); }); + + it('should support BigInt arguments', function () { + let callCount = 0; + const multiply = memoize((a: bigint, b: bigint) => { + callCount += 1; + return a * b; + }); + + const result1 = multiply(2n, 3n); + const result2 = multiply(2n, 3n); + const result3 = multiply(3n, 3n); + + expect(result1).to.equal(6n); + expect(result2).to.equal(6n); + expect(result3).to.equal(9n); + expect(callCount).to.equal(2); + }); + + it('should support circular object arguments', function () { + let callCount = 0; + const pickName = memoize((obj: {name: string; self?: unknown}) => { + callCount += 1; + return obj.name; + }); + + const circularArg = {name: 'first'} as {name: string; self?: unknown}; + circularArg.self = circularArg; + const secondCircularArg = {name: 'second'} as {name: string; self?: unknown}; + secondCircularArg.self = secondCircularArg; + + const result1 = pickName(circularArg); + const result2 = pickName(circularArg); + const result3 = pickName(secondCircularArg); + + expect(result1).to.equal('first'); + expect(result2).to.equal('first'); + expect(result3).to.equal('second'); + expect(callCount).to.equal(2); + }); }); });