diff --git a/.changeset/cute-chairs-kneel.md b/.changeset/cute-chairs-kneel.md new file mode 100644 index 00000000000..e00e03eced1 --- /dev/null +++ b/.changeset/cute-chairs-kneel.md @@ -0,0 +1,5 @@ +--- +"effect": minor +--- + +Add Function.memo diff --git a/packages/effect/src/Function.ts b/packages/effect/src/Function.ts index 69d53fd2fdd..7e0d224a482 100644 --- a/packages/effect/src/Function.ts +++ b/packages/effect/src/Function.ts @@ -1220,3 +1220,124 @@ export const hole: () => T = unsafeCoerce(absurd) * @since 2.0.0 */ export const SK = (_: A, b: B): B => b + +/** + * Empty ArgNode result sentinel. + * + * @internal + */ +const EMPTY = Symbol.for("effect/Function.memo") + +/** + * Memo helper class. + * + * @internal + */ +class ArgNode { + /** 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> | null = null + + /** The object arg branch. */ + private o: WeakMap> | null = null + + /** Get the next node for an arg. If uninitialized, one will be created, set, and returned. */ + get(arg: unknown): ArgNode { + let next: ArgNode | 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() + + 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>() + +/** + * 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 = , Return>( + fn: (...args: Args) => Return +): (...args: Args) => Return => { + let root = fnRoots.get(fn) as ArgNode | 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: , Return>( + fn: (this: This, ...args: Args) => Return +) => (this: This, ...args: Args) => Return = memo diff --git a/packages/effect/test/Function.test.ts b/packages/effect/test/Function.test.ts index de67d89dc5a..53659145cfc 100644 --- a/packages/effect/test/Function.test.ts +++ b/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) + }) + }) })