add WeakCache implementation

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,24 @@ 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<K, V> {
+  new <K,V>(max?: number, dispose?: (value: V, key?: K) => void): CommonCache<K,V>;
+}
+
 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.
@@ -106,9 +120,10 @@ export type OptimisticWrapOptions<
   // If provided, the subscribe function should either return an unsubscribe
   // function or return nothing.
   subscribe?: (...args: TArgs) => void | (() => any);
+  Cache?: CommonCacheConstructor<TCacheKey, TResult>
 };
 
-const caches = new Set<Cache<any, AnyEntry>>();
+const caches = new Set<CommonCache<any, AnyEntry>>();
 
 export function wrap<
   TArgs extends any[],
@@ -120,6 +135,7 @@ export function wrap<
   makeCacheKey = defaultMakeCacheKey,
   keyArgs,
   subscribe,
+  Cache = StrongCache
 }: OptimisticWrapOptions<TArgs, TKeyArgs> = Object.create(null)) {
   const cache = new Cache<TCacheKey, Entry<TArgs, TResult>>(
     max,
@@ -168,7 +184,7 @@ export function wrap<
 
   Object.defineProperty(optimistic, "size", {
     get() {
-      return cache["map"].size;
+      return cache.size;
     },
     configurable: false,
     enumerable: false,

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);
+  });
+});

src/weakCache.ts

@@ -1,21 +1,48 @@
-interface Node<K, V> {
-  key: K;
+interface Node<K extends object, V> {
+  keyRef: WeakRef<K>;
   value: V;
   newer: Node<K, V> | null;
   older: Node<K, V> | null;
 }
 
-function defaultDispose() {}
-
-export class Cache<K = any, V = any> {
-  private map = new Map<K, Node<K, V>>();
+function noop() {}
+const defaultDispose = noop;
+
+const _WeakRef =
+  typeof WeakRef !== "undefined"
+    ? WeakRef
+    : (function <T>(value: T) {
+        return { deref: () => value } satisfies Omit<
+          WeakRef<any>,
+          typeof Symbol.toStringTag
+        >;
+      } as typeof WeakRef);
+const _WeakMap = typeof WeakMap !== "undefined" ? WeakMap : Map;
+const _FinalizationRegistry =
+  typeof FinalizationRegistry !== "undefined"
+    ? FinalizationRegistry
+    : (function <T>() {
+        return {
+          register: noop,
+          unregister: noop,
+        } satisfies Omit<FinalizationRegistry<T>, typeof Symbol.toStringTag>;
+      } as typeof FinalizationRegistry);
+
+export class WeakCache<K extends object = any, V = any> {
+  private map = new _WeakMap<K, Node<K, V>>();
+  private registry: FinalizationRegistry<Node<K, V>>;
   private newest: Node<K, V> | null = null;
   private oldest: Node<K, V> | null = null;
+  public size = 0;
 
   constructor(
     private max = Infinity,
-    public dispose: (value: V, key: K) => void = defaultDispose,
-  ) {}
+    public dispose: (value: V, key?: K) => void = defaultDispose
+  ) {
+    this.registry = new _FinalizationRegistry<Node<K, V>>(
+      this.deleteNode.bind(this)
+    );
+  }
 
   public has(key: K): boolean {
     return this.map.has(key);
@@ -57,14 +84,14 @@ export class Cache<K = any, V = any> {
   public set(key: K, value: V): V {
     let node = this.getNode(key);
     if (node) {
-      return node.value = value;
+      return (node.value = value);
     }
 
     node = {
-      key,
+      keyRef: new _WeakRef(key),
       value,
       newer: null,
-      older: this.newest
+      older: this.newest,
     };
 
     if (this.newest) {
@@ -74,38 +101,45 @@ export class Cache<K = any, V = any> {
     this.newest = node;
     this.oldest = this.oldest || node;
 
+    this.registry.register(key, node);
     this.map.set(key, node);
+    this.size++;
 
     return node.value;
   }
 
   public clean() {
-    while (this.oldest && this.map.size > this.max) {
-      this.delete(this.oldest.key);
+    while (this.oldest && this.size > this.max) {
+      this.deleteNode(this.oldest);
     }
   }
 
-  public delete(key: K): boolean {
-    const node = this.map.get(key);
-    if (node) {
-      if (node === this.newest) {
-        this.newest = node.older;
-      }
+  private deleteNode(node: Node<K, V>) {
+    if (node === this.newest) {
+      this.newest = node.older;
+    }
 
-      if (node === this.oldest) {
-        this.oldest = node.newer;
-      }
+    if (node === this.oldest) {
+      this.oldest = node.newer;
+    }
 
-      if (node.newer) {
-        node.newer.older = node.older;
-      }
+    if (node.newer) {
+      node.newer.older = node.older;
+    }
 
-      if (node.older) {
-        node.older.newer = node.newer;
-      }
+    if (node.older) {
+      node.older.newer = node.newer;
+    }
+    this.size--;
+    const key = node.keyRef.deref();
+    this.dispose(node.value, key);
+    if (key) this.map.delete(key);
+  }
 
-      this.map.delete(key);
-      this.dispose(node.value, key);
+  public delete(key: K): boolean {
+    const node = this.map.get(key);
+    if (node) {
+      this.deleteNode(node);
 
       return true;
     }

tsconfig.json

@@ -8,7 +8,7 @@
     "sourceMap": true,
     "declaration": true,
     "importHelpers": true,
-    "lib": ["es2015"],
+    "lib": ["es2015", "ES2021.WeakRef"],
     "types": ["node", "mocha"],
     "strict": true,
     "noImplicitAny": true,