Add configurable WeakCache implementation based on FinalizationRegistry

package.json

@@ -32,7 +32,7 @@
     "tsc:es5": "tsc -p tsconfig.es5.json",
     "clean": "rimraf lib",
     "prepare": "npm run build",
-    "mocha": "mocha --require source-map-support/register --reporter spec --full-trace",
+    "mocha": "mocha --require source-map-support/register --reporter spec --full-trace -n expose-gc",
     "test:cjs": "npm run mocha -- lib/tests/bundle.cjs",
     "test:esm": "npm run mocha -- lib/tests/bundle.js",
     "test": "npm run test:esm && npm run test:cjs"

src/cache.ts

@@ -26,6 +26,10 @@ export class Cache<K = any, V = any> {
     return node && node.value;
   }
 
+  public get size() {
+    return this.map.size;
+  }
+
   private getNode(key: K): Node<K, V> | undefined {
     const node = this.map.get(key);
 

src/index.ts

@@ -1,6 +1,6 @@
 import { Trie } from "@wry/trie";
 
-import { Cache } from "./cache.js";
+import { Cache as StrongCache } from "./cache.js";
 import { Entry, AnyEntry } from "./entry.js";
 import { parentEntrySlot } from "./context.js";
 
@@ -88,10 +88,26 @@ export type OptimisticWrapperFunction<
   makeCacheKey: (...args: TKeyArgs) => TCacheKey;
 };
 
+export interface CommonCache<K,V> {
+  has(key: K): boolean;
+  get(key: K): V | undefined;
+  set(key: K, value: V): V;
+  delete(key: K): boolean;
+  clean(): void;
+  readonly size: number;
+}
+
+export interface CommonCacheConstructor<TCacheKey, TResult, TArgs extends any[]> {
+  new <K extends TCacheKey, V extends Entry<TArgs, TResult>>(max?: number, dispose?: (value: V, key?: K) => void): CommonCache<K,V>;
+}
+
+type NoInfer<T> = [T][T extends any ? 0 : never];
+
 export type OptimisticWrapOptions<
   TArgs extends any[],
   TKeyArgs extends any[] = TArgs,
   TCacheKey = any,
+  TResult = any,
 > = {
   // The maximum number of cache entries that should be retained before the
   // cache begins evicting the oldest ones.
@@ -102,13 +118,14 @@ export type OptimisticWrapOptions<
   // The makeCacheKey function takes the same arguments that were passed to
   // the wrapper function and returns a single value that can be used as a key
   // in a Map to identify the cached result.
-  makeCacheKey?: (...args: TKeyArgs) => TCacheKey;
+  makeCacheKey?: (...args: NoInfer<TKeyArgs>) => TCacheKey;
   // If provided, the subscribe function should either return an unsubscribe
   // function or return nothing.
   subscribe?: (...args: TArgs) => void | (() => any);
+  Cache?: CommonCacheConstructor<NoInfer<TCacheKey>, NoInfer<TResult>, NoInfer<TArgs>>
 };
 
-const caches = new Set<Cache<any, AnyEntry>>();
+const caches = new Set<CommonCache<any, AnyEntry>>();
 
 export function wrap<
   TArgs extends any[],
@@ -117,10 +134,11 @@ export function wrap<
   TCacheKey = any,
 >(originalFunction: (...args: TArgs) => TResult, {
   max = Math.pow(2, 16),
-  makeCacheKey = defaultMakeCacheKey,
+  makeCacheKey = (defaultMakeCacheKey as () => TCacheKey),
   keyArgs,
   subscribe,
-}: OptimisticWrapOptions<TArgs, TKeyArgs> = Object.create(null)) {
+  Cache = StrongCache
+}: OptimisticWrapOptions<TArgs, TKeyArgs, TCacheKey, TResult> = Object.create(null)) {
   const cache = new Cache<TCacheKey, Entry<TArgs, TResult>>(
     max,
     entry => entry.dispose(),
@@ -168,7 +186,7 @@ export function wrap<
 
   Object.defineProperty(optimistic, "size", {
     get() {
-      return cache["map"].size;
+      return cache.size;
     },
     configurable: false,
     enumerable: false,

src/tests/api.ts

@@ -7,6 +7,8 @@ import {
 } from "../index";
 import { wrapYieldingFiberMethods } from '@wry/context';
 import { dep } from "../dep";
+import { WeakCache } from "../weakCache";
+import { Cache } from "../cache";
 
 type NumThunk = OptimisticWrapperFunction<[], number>;
 
@@ -16,6 +18,36 @@ describe("optimism", function () {
     assert.strictEqual(typeof defaultMakeCacheKey, "function");
   });
 
+  it("can manually set the `Cache` to `WeakCache`", async () => {
+    let cache: WeakCache
+    wrap((obj: { value: string }) => obj.value, {
+      Cache: (class extends WeakCache {
+        constructor(...args: ConstructorParameters<typeof WeakCache>) {
+          super(...args)
+          cache = this
+        } 
+      }) as typeof WeakCache
+    });
+    // can't access the cache otherwise, so we can only test if the `Cache`
+    // argument constructor was called, not if it's actually used internally
+    assert.ok(cache! instanceof WeakCache);
+  });
+
+  it("can manually set the `Cache` to `Cache`", () => {
+    let cache: Cache
+    wrap((obj: { value: string }) => obj.value, {
+      Cache: (class extends Cache {
+        constructor(...args: ConstructorParameters<typeof Cache>) {
+          super(...args)
+          cache = this
+        } 
+      }) as typeof Cache
+    });
+    // can't access the cache otherwise, so we can only test if the `Cache`
+    // argument constructor was called, not if it's actually used internally
+    assert.ok(cache! instanceof Cache);
+  });
+
   it("works with single functions", function () {
     const test = wrap(function (x: string) {
       return x + salt;

src/tests/main.ts

@@ -5,3 +5,4 @@ import "./key-trie";
 import "./context";
 import "./exceptions";
 import "./performance";
+import "./weakCache";

src/tests/weakCache.ts

@@ -0,0 +1,117 @@
+import * as assert from "assert";
+import { WeakCache } from "../weakCache";
+
+describe("weak least-recently-used cache", function () {
+  it("can hold lots of elements", function () {
+    const cache = new WeakCache();
+    const count = 1000000;
+    const keys = [];
+
+    for (let i = 0; i < count; ++i) {
+      const key = {};
+      cache.set(key, String(i));
+      keys[i] = key;
+    }
+
+    cache.clean();
+
+    assert.strictEqual(cache.size, count);
+    assert.ok(cache.has(keys[0]));
+    assert.ok(cache.has(keys[count - 1]));
+    assert.strictEqual(cache.get(keys[43]), "43");
+  });
+
+  it("evicts excess old elements", function () {
+    const max = 10;
+    const evicted = [];
+    const cache = new WeakCache(max, (value, key) => {
+      assert.strictEqual(key.valueOf(), value.valueOf());
+      evicted.push(key);
+    });
+
+    const count = 100;
+    const keys = [];
+    for (let i = 0; i < count; ++i) {
+      const key = new String(i);
+      cache.set(key, String(i));
+      keys[i] = key;
+    }
+
+    cache.clean();
+
+    assert.strictEqual((cache as any).size, max);
+    assert.strictEqual(evicted.length, count - max);
+
+    for (let i = count - max; i < count; ++i) {
+      assert.ok(cache.has(keys[i]));
+    }
+  });
+
+  it("can cope with small max values", function () {
+    const cache = new WeakCache(2);
+    const keys = Array(10)
+      .fill(null)
+      .map((_, i) => new Number(i));
+
+    function check(...sequence: number[]) {
+      cache.clean();
+
+      let entry = (cache as any).newest;
+      const forwards = [];
+      while (entry) {
+        forwards.push(entry.keyRef.deref());
+        entry = entry.older;
+      }
+      assert.deepEqual(forwards.map(Number), sequence);
+
+      const backwards = [];
+      entry = (cache as any).oldest;
+      while (entry) {
+        backwards.push(entry.keyRef.deref());
+        entry = entry.newer;
+      }
+      backwards.reverse();
+      assert.deepEqual(backwards.map(Number), sequence);
+
+      sequence.forEach(function (n) {
+        assert.strictEqual((cache as any).map.get(keys[n]).value, n + 1);
+      });
+
+      if (sequence.length > 0) {
+        assert.strictEqual((cache as any).newest.keyRef.deref().valueOf(), sequence[0]);
+        assert.strictEqual(
+          (cache as any).oldest.keyRef.deref().valueOf(),
+          sequence[sequence.length - 1]
+        );
+      }
+    }
+
+    cache.set(keys[1], 2);
+    check(1);
+
+    cache.set(keys[2], 3);
+    check(2, 1);
+
+    cache.set(keys[3], 4);
+    check(3, 2);
+
+    cache.get(keys[2]);
+    check(2, 3);
+
+    cache.set(keys[4], 5);
+    check(4, 2);
+
+    assert.strictEqual(cache.has(keys[1]), false);
+    assert.strictEqual(cache.get(keys[2]), 3);
+    assert.strictEqual(cache.has(keys[3]), false);
+    assert.strictEqual(cache.get(keys[4]), 5);
+
+    cache.delete(keys[2]);
+    check(4);
+    cache.delete(keys[4]);
+    check();
+
+    assert.strictEqual((cache as any).newest, null);
+    assert.strictEqual((cache as any).oldest, null);
+  });
+});