Add Function.memo

.changeset/cute-chairs-kneel.md

@@ -0,0 +1,5 @@
+---
+"effect": minor
+---
+
+Add Function.memo

packages/effect/src/Function.ts

@@ -1220,3 +1220,124 @@ export const hole: <T>() => T = unsafeCoerce(absurd)
  * @since 2.0.0
  */
 export const SK = <A, B>(_: A, b: B): B => b
+
+/**
+ * Empty ArgNode result sentinel.
+ *
+ * @internal
+ */
+const EMPTY = Symbol.for("effect/Function.memo")
+
+/**
+ * Memo helper class.
+ *
+ * @internal
+ */
+class ArgNode<A> {
+  /** The cached result of the args that led to this node, or `EMPTY`. */
+  r: A | typeof EMPTY = EMPTY
+
+  /** The primitive arg branch. */
+  private p: Map<unknown, ArgNode<A>> | null = null
+
+  /** The object arg branch. */
+  private o: WeakMap<object, ArgNode<A>> | null = null
+
+  /** Get the next node for an arg. If uninitialized, one will be created, set, and returned. */
+  get(arg: unknown): ArgNode<A> {
+    let next: ArgNode<A> | undefined
+    let isObject: boolean
+    switch (typeof arg) {
+      case "object":
+      // Fallthrough intended for primitive null case
+      case "function":
+        if (arg !== null) {
+          if (!this.o) this.o = new WeakMap()
+          next = this.o.get(arg)
+          isObject = true
+          break
+        }
+      default:
+        if (!this.p) this.p = new Map()
+        next = this.p.get(arg)
+        isObject = false
+    }
+
+    if (next) return next
+
+    const fresh = new ArgNode<A>()
+
+    if (isObject) {
+      this.o!.set(arg as object, fresh)
+    } else {
+      this.p!.set(arg, fresh)
+    }
+
+    return fresh
+  }
+}
+
+/**
+ * A global WeakMap for memoized functions and their `ArgNode`s, used so that functions wrapped multiple times will share results.
+ *
+ * @internal
+ */
+const fnRoots = new WeakMap<Function, ArgNode<unknown>>()
+
+/**
+ * Memoize a function, method, or getter, with any number of arguments. Repeated calls to the returned function with the same arguments will return cached values.
+ *
+ * Usage notes:
+ * - Memoized functions should be totally pure, and should return immutable values.
+ * - The cache size is unbounded, but internally a `WeakMap` is used when possible. To make the most of this, memoized functions should have object-type args at the start and primitive args at the end.
+ * - Works as a class method decorator.
+ *
+ * @example
+ * ```ts
+ * import { memo } from "effect/Function"
+ *
+ * const add = memo((x: number, y: number) => {
+ *   console.log("running add");
+ *   return x + y;
+ * });
+ *
+ * add(2, 3); // logs "running add", returns 5
+ * add(2, 3); // no log, returns cached 5
+ * add(2, 4); // logs "running add", returns 6
+ *
+ * // Expected console output:
+ * // running add
+ * // running add
+ * ```
+ *
+ * @since 3.20.0
+ */
+export const memo = <Args extends Array<unknown>, Return>(
+  fn: (...args: Args) => Return
+): (...args: Args) => Return => {
+  let root = fnRoots.get(fn) as ArgNode<Return> | undefined
+  if (!root) {
+    root = new ArgNode()
+    fnRoots.set(fn, root)
+  }
+
+  return function(this: unknown) {
+    let node = root.get(this)
+    for (let i = 0; i < arguments.length; i += 1) {
+      node = node.get(arguments[i])
+    }
+
+    if (node.r !== EMPTY) return node.r
+    node.r = fn.apply(this, arguments as unknown as Args)
+    return node.r
+  }
+}
+
+/**
+ * See {@link memo}. This is an alias that infers `This` and uses it in the returned function signature.
+ *
+ * @since 3.20.0
+ */
+export const memoThis: <This, Args extends Array<unknown>, Return>(
+  fn: (this: This, ...args: Args) => Return
+) => (this: This, ...args: Args) => Return = memo

packages/effect/test/Function.test.ts

@@ -1,4 +1,4 @@
-import { describe, it } from "@effect/vitest"
+import { describe, expect, it, vi } from "@effect/vitest"
 import { deepStrictEqual, strictEqual, throws } from "@effect/vitest/utils"
 import { Function, String } from "effect"
 
@@ -195,4 +195,89 @@ describe("Function", () => {
       deepStrictEqual(f.apply(null, ["_", "a", "b", "c", "d", "e", "f"] as any), "_abcde")
     })
   })
+
+  describe("memo", () => {
+    it("memoizes by primitive argument list", () => {
+      const fn = vi.fn((a: number, b: number) => a + b)
+      const memoFn = Function.memo(fn)
+
+      strictEqual(memoFn(1, 2), 3)
+      expect(fn).toHaveBeenCalledTimes(1)
+
+      // Cache hit (same args)
+      strictEqual(memoFn(1, 2), 3)
+      expect(fn).toHaveBeenCalledTimes(1)
+
+      // Different args -> cache miss
+      strictEqual(memoFn(2, 1), 3)
+      expect(fn).toHaveBeenCalledTimes(2)
+    })
+
+    it("memoizes zero-argument functions and shares cache per fn", () => {
+      const fn = vi.fn(() => 42)
+      const memoFnA = Function.memo(fn)
+      const memoFnB = Function.memo(fn) // wrappers share same per-fn cache
+
+      expect(memoFnA()).toBe(42)
+      expect(fn).toHaveBeenCalledTimes(1)
+
+      // Cache hit via same wrapper
+      expect(memoFnA()).toBe(42)
+      expect(fn).toHaveBeenCalledTimes(1)
+
+      // Cache hit via different wrapper for same original function
+      expect(memoFnB()).toBe(42)
+      expect(fn).toHaveBeenCalledTimes(1)
+    })
+
+    it("memoizes undefined results distinctly", () => {
+      const fn = vi.fn((_: number): void => undefined)
+      const memoFn = Function.memo(fn)
+
+      expect(memoFn(123)).toBeUndefined()
+      expect(fn).toHaveBeenCalledTimes(1)
+
+      // Cache hit despite undefined result
+      expect(memoFn(123)).toBeUndefined()
+      expect(fn).toHaveBeenCalledTimes(1)
+    })
+
+    it("uses object identity for object args", () => {
+      const a = { x: 1 }
+      const b = { x: 1 }
+
+      const fn = vi.fn((o: { x: number }) => o.x)
+      const memoFn = Function.memo(fn)
+
+      expect(memoFn(a)).toBe(1)
+      expect(fn).toHaveBeenCalledTimes(1)
+
+      // Cache hit for same object instance
+      expect(memoFn(a)).toBe(1)
+      expect(fn).toHaveBeenCalledTimes(1)
+
+      // Different object instance with same shape -> cache miss
+      expect(memoFn(b)).toBe(1)
+      expect(fn).toHaveBeenCalledTimes(2)
+    })
+
+    it("includes receiver (this) in the cache key", () => {
+      function plusOffset(this: { offset: number }, x: number) {
+        return x + this.offset
+      }
+      const fnThis = vi.fn(plusOffset)
+      // Casts only needed to appease TS 'this' typing for test invocations
+      const memoFn = Function.memo(fnThis)
+
+      const a = { offset: 1 }
+      const b = { offset: 100 }
+
+      expect(memoFn.call(a, 1)).toBe(2)
+      expect(fnThis).toHaveBeenCalledTimes(1)
+
+      // Different receiver with same args should not reuse prior result
+      expect(memoFn.call(b, 1)).toBe(101)
+      expect(fnThis).toHaveBeenCalledTimes(2)
+    })
+  })
 })