The caching decorator from the previous article was fine, but there was a problem with it: it was too specific, it assumed only one primitive argument for the decorated method, so it's abilities were limited. It also had no mechanism of clearing the cache, so results would be stored until the application is torn down.

Let's think of a better solution, that would be more flexible. Something that:

does not make assumptions about the key used for storing memoized result;
has a mechanism for clearing the cache;
can have debugging capabilities;
can use either Map or WeakMap for storage.

Since in js we can pass function as arguments, it allows us to have great flexibility: we can pass a function to extract unique id from arguments, which we later can use as a key to memoize a result, instead of using an argument as a key, which was the case with the previous decorator.

Hence, the payload for our new memoize decorator can be summed up in the following interface:

// ./models/memoize-payload.model.ts
export interface MemoizePayload {
  // The function that will determine a unique id for the provided arguments set, determined by used
  extractUniqueId: (...args: any[]) => any;
  // A flag to use WeakMap
  doUseWeakMap?: boolean;
  // If regular map is used, you can set timeout to clear its contents, optional
  clearCacheTimeout?: number;
  // For debug purposes you can pass an exta function for logging all actions
  debugReporter?: (message: string, state?: Map<any, unknown> | WeakMap<object, unknown> | unknown) => void;
}

Since now we're passing arguments to the decorator, it will no longer be a mere function, but a factory one, producing a decorator. The decorator function itself can be typed as:

// ./models/decorator.model.ts
export type MemoizeDecorator = (target: unknown, propertyKey: string, descriptor: PropertyDescriptor) => void;

Let's assume two use cases:

with timed cache clearing, this version doesn't need the boolean flag doUseWeakMap;
without cache clearing with optional WeakMap, this version does not need clearCacheTimeout argument.

Bearing that in mind, using function overload in typescript we can produce two signatures for the factory:

export function memoize(args: Omit<MemoizePayload, 'doUseWeakMap'>): MemoizeDecorator;
export function memoize(args: Omit<MemoizePayload, 'clearCacheTimeout'>): MemoizeDecorator;

Now for the implementation of the memoize decorator factory itself, with some debug messages passed to the optional debugReporter:

// ./memoize.decorator.ts
import { MemoizeDecorator } from './models/decorator.model';
import { MemoizePayload } from './models/memoize-payload.model';

export function memoize(args: Omit<MemoizePayload, 'doUseWeakMap'>): MemoizeDecorator;
export function memoize(args: Omit<MemoizePayload, 'clearCacheTimeout'>): MemoizeDecorator;

export function memoize({ extractUniqueId, clearCacheTimeout, doUseWeakMap, debugReporter }: MemoizePayload): MemoizeDecorator {
  return (target: unknown, propertyKey: string, descriptor: PropertyDescriptor): void => {
    let cacheTeardownTimer: ReturnType<typeof setTimeout>;

    let cache = initCache(doUseWeakMap);

    const startTeardownTimeout = !clearCacheTimeout
      ? null
      : () => {
          if (cacheTeardownTimer) {
            debugReporter?.('Clearing the cache timeout timer');
            clearTimeout(cacheTeardownTimer);
          }
          debugReporter?.(`Cache to be cleared in ${clearCacheTimeout}ms`);
          cacheTeardownTimer = setTimeout(() => {
            debugReporter?.('Clearing the current cache of', cache);
            cache = initCache(doUseWeakMap);
            debugReporter?.('Cache cleared: ', cache);
          }, clearCacheTimeout);
        };

    const originalMethod = descriptor.value;

    descriptor.value = function (...args: unknown[]) {
      startTeardownTimeout?.();

      const uniqueId: any = extractUniqueId(...args);
      debugReporter?.('Looking for a value with unique id of ', uniqueId);

      if (cache.has(uniqueId)) {
        const cachedResult = cache.get(uniqueId);
        debugReporter?.('Returning cached result', cachedResult);
        return cachedResult;
      }

      debugReporter?.('No cached result found');
      const result = originalMethod.apply(this, args);

      debugReporter?.('Storing a new entry in cache: ', { uniqueId, result });
      cache.set(uniqueId, result);
      debugReporter?.('Cache updated', cache);

      return result;
    };
  };
}

function initCache(doUseWeakMap?: boolean) {
  return doUseWeakMap ? new WeakMap<object, unknown>() : new Map<any, unknown>();
}

Looks good, but better test it. We can use jest for that, create a class with a method that does some calculations, running it several times and checking how much times does it take to run a decorated methods vs non-decorated.

First let's create a helper function to check the time:

function checkRuntime<T extends (...args: any) => any, C>(fun: T, iterations: number, countdown: C): number {
  const start = performance.now();
  for (let x = 0; x++ <= iterations; ) {
    fun(countdown);
  }
  const end = performance.now();
  return end - start;
}

It basically runs a function a number of times and tells how long it took, nothing special.

Since I have poor imagination, I will use a similar countdown class I used in the previous article:

interface CalculationPayload {
  id: number;
  n: number;
}

class ObjectCountdownCalculator {
  @memoize({
    extractUniqueId: (a) => a.id,
  })
  public static decoratedCountDownWithMap(a: CalculationPayload): number {
    return ObjectCountdownCalculator.countdown(a);
  }

  @memoize({
    extractUniqueId: (a) => a,
    doUseWeakMap: true,
  })
  public static decoratedCountdownWithWeakMap(a: CalculationPayload): number {
    return ObjectCountdownCalculator.countdown(a);
  }

  public static nonDecoratedCountDown(a: CalculationPayload): number {
    return ObjectCountdownCalculator.countdown(a);
  }

  private static countdown({ n }: CalculationPayload): number {
    let count = 0;
    while (count < n) {
      count += 1;
    }
    return count;
  }
}

So what's going on here? There's a private method doing all the work and 3 public ones calling it, two of them are decorated. With 3 public methods invoking the same internal one, we can run them and compare results to see how fast our memoization decorator is compared to regular method call.

Let's assume it's 5 times faster if there are 5 iteration calls:

it('The decorated version of the countdown with regular map should be at least 5 times faster', () => {
  // Arrange
  const countdown: CalculationPayload = {
    id: 1,
    n: 500000000,
  };
  const iterations = 5;

  // Act
  const decoratedVersionTime = checkRuntime(ObjectCountdownCalculator.decoratedCountDownWithMap, iterations, countdown);
  const nonDecoratedVersionTime = checkRuntime(ObjectCountdownCalculator.nonDecoratedCountDown, iterations, countdown);

  // Assert
  expect(Math.floor(nonDecoratedVersionTime / decoratedVersionTime)).toBeGreaterThanOrEqual(5);
});

And it actually is that fast. What's interesting, is that the WeakMap version turned out to be only 4 times faster. I wonder why though.

Anyway, that's about it.

You can find the code as

Memoize Decorator