feat: added memo()

Author: webdeveric

src/memo.ts

@@ -0,0 +1,48 @@
+import type { AnyFunction } from './types/common.js';
+
+export type CacheKeyFn<Fn extends AnyFunction, ReturnValue = unknown> = (args: Parameters<Fn>) => ReturnValue;
+
+export type DefaultCacheKeyFn<Fn extends AnyFunction> = CacheKeyFn<Fn, NonNullable<Parameters<Fn>[0]> | string>;
+
+export type MemoizedFn<Fn extends AnyFunction, CkFn extends CacheKeyFn<Fn>> = Fn & {
+  cache: Map<ReturnType<CkFn>, ReturnType<Fn>>;
+};
+
+/**
+ * Memoize a function
+ *
+ * @todo use `Map.getOrInsertComputed()` once it is widely available.
+ */
+export function memo<Fn extends AnyFunction, CkFn extends CacheKeyFn<Fn>>(
+  fn: Fn,
+  getCacheKey: CkFn,
+): MemoizedFn<Fn, CkFn>;
+
+export function memo<Fn extends AnyFunction>(fn: Fn): MemoizedFn<Fn, DefaultCacheKeyFn<Fn>>;
+
+export function memo<Fn extends AnyFunction, CkFn extends CacheKeyFn<Fn>>(
+  fn: Fn,
+  getCacheKey?: CkFn,
+): MemoizedFn<Fn, CkFn> {
+  const defaultGetCacheKey: DefaultCacheKeyFn<Fn> = (args) => args[0] ?? JSON.stringify(args);
+
+  const resolvedGetCacheKey = getCacheKey ?? defaultGetCacheKey;
+
+  const memoized = (...args: Parameters<Fn>): ReturnType<Fn> => {
+    const key = resolvedGetCacheKey(args);
+
+    if (memoized.cache.has(key)) {
+      return memoized.cache.get(key);
+    }
+
+    const value = fn(...args);
+
+    memoized.cache.set(key, value);
+
+    return value;
+  };
+
+  memoized.cache = new Map();
+
+  return memoized as MemoizedFn<Fn, CkFn>;
+}

test/memo.test.ts

@@ -0,0 +1,104 @@
+import { describe, expect, it, vi } from 'vitest';
+
+import { memo } from '../src/memo.js';
+
+describe('memo()', () => {
+  it('Memoizes functions', () => {
+    const mock = vi.fn((str: string) => str);
+
+    const fn = memo(mock);
+
+    expect(mock).not.toHaveBeenCalled();
+
+    fn('test');
+
+    expect(mock).toHaveBeenCalledTimes(1);
+
+    fn('test');
+
+    expect(mock).toHaveBeenCalledTimes(1);
+
+    fn('test2');
+
+    expect(mock).toHaveBeenCalledTimes(2);
+  });
+
+  it('Memoizes async functions', async () => {
+    const mock = vi.fn(async (str: string): Promise<string> => str);
+
+    const fn = memo(mock);
+
+    expect(mock).not.toHaveBeenCalled();
+
+    await fn('test');
+
+    expect(mock).toHaveBeenCalledTimes(1);
+
+    await fn('test');
+
+    expect(mock).toHaveBeenCalledTimes(1);
+
+    await fn('test2');
+
+    expect(mock).toHaveBeenCalledTimes(2);
+
+    expect(fn('test3')).toBe(fn('test3'));
+  });
+
+  it('Cache is available on the memoized function', () => {
+    const mock = vi.fn((str: string) => str);
+
+    const fn = memo(mock);
+
+    expect(mock).not.toHaveBeenCalled();
+
+    fn('test');
+
+    expect(mock).toHaveBeenCalledTimes(1);
+
+    fn('test');
+
+    expect(mock).toHaveBeenCalledTimes(1);
+
+    fn.cache.clear();
+
+    fn('test');
+
+    expect(mock).toHaveBeenCalledTimes(2);
+  });
+
+  it('Cache key is configurable', () => {
+    const mock = vi.fn((_id: number, name: string) => name);
+
+    const fn = memo(mock, ([id, name]) => `${id}-${name}-key`);
+
+    fn(1, 'test');
+
+    expect(fn.cache.has('1-test-key')).toBe(true);
+  });
+
+  it('Default getCacheKey()', () => {
+    const mock = vi.fn((_id: number | null | undefined, name: string) => name);
+
+    const fn = memo(mock);
+
+    expect(mock).not.toHaveBeenCalled();
+
+    fn(1, 'test');
+
+    expect(mock).toHaveBeenCalledTimes(1);
+
+    expect(fn.cache.has(1)).toBe(true);
+
+    fn(null, 'test');
+
+    expect(mock).toHaveBeenCalledTimes(2);
+
+    fn(undefined, 'test');
+
+    // `undefined` serialized to `null` when used in arrays.
+    expect(mock).toHaveBeenCalledTimes(2);
+
+    expect(fn.cache.has('[null,"test"]')).toBe(true);
+  });
+});