fix: isolate cache eviction registrations per entry

index.d.ts

@@ -28,7 +28,7 @@ console.log(cache.has('key'));
 */
 export default class WeakCache<
   K = string,
-  V extends Record<string, unknown> = Record<string, unknown>,
+  V extends object = Record<string, unknown>,
 > {
   /**
 	The approximate size of the cache. May include entries with collected values.

index.js

@@ -14,7 +14,10 @@ export default class WeakCache {
 	*/
   constructor(options = {}) {
     this.#onEvict = options.onEvict;
-    this.#registry = new FinalizationRegistry((key) => {
+    this.#registry = new FinalizationRegistry(({ key, reference }) => {
+      if (this.#map.get(key) !== reference) {
+        return;
+      }
       this.#map.delete(key);
       if (this.#onEvict) {
         this.#onEvict(key);
@@ -62,14 +65,12 @@ export default class WeakCache {
 
     const existing = this.#map.get(key);
     if (existing !== undefined) {
-      const oldValue = existing.deref();
-      if (oldValue !== undefined) {
-        this.#registry.unregister(oldValue);
-      }
+      this.#registry.unregister(existing);
     }
 
-    this.#map.set(key, new WeakRef(value));
-    this.#registry.register(value, key, value);
+    const reference = new WeakRef(value);
+    this.#map.set(key, reference);
+    this.#registry.register(value, { key, reference }, reference);
   }
 
   /**
@@ -101,10 +102,7 @@ export default class WeakCache {
   delete(key) {
     const reference = this.#map.get(key);
     if (reference !== undefined) {
-      const value = reference.deref();
-      if (value !== undefined) {
-        this.#registry.unregister(value);
-      }
+      this.#registry.unregister(reference);
     }
 
     return this.#map.delete(key);

index.test-d.ts

@@ -28,3 +28,14 @@ expectType<{ value: string } | undefined>(numberCache.get(1));
 // Errors for primitive values
 expectError(typedCache.set("key", "string"));
 expectError(typedCache.set("key", 42));
+
+const functions = new WeakCache<string, () => number>();
+functions.set("callback", () => 42);
+expectType<(() => number) | undefined>(functions.get("callback"));
+const arrays = new WeakCache<string, number[]>();
+arrays.set("items", [1, 2]);
+expectType<number[] | undefined>(arrays.get("items"));
+const dates = new WeakCache<string, Date>();
+dates.set("date", new Date());
+expectType<Date | undefined>(dates.get("date"));
+expectError(new WeakCache<string, string>());

regression.test.js

@@ -0,0 +1,81 @@
+import test from "ava";
+import WeakCache from "./index.js";
+
+function withRegistryFixture(run) {
+  const OriginalRegistry = globalThis.FinalizationRegistry;
+  const OriginalWeakRef = globalThis.WeakRef;
+  let registry;
+  const references = [];
+  class RegistryFixture {
+    registrations = [];
+    constructor(callback) {
+      this.callback = callback;
+      registry = this;
+    }
+    register(_target, held, token) {
+      this.registrations.push({ held, token });
+    }
+    unregister(token) {
+      const previous = this.registrations.length;
+      this.registrations = this.registrations.filter(
+        (entry) => entry.token !== token
+      );
+      return this.registrations.length !== previous;
+    }
+  }
+  class WeakRefFixture {
+    constructor(value) {
+      this.value = value;
+      references.push(this);
+    }
+    deref() {
+      return this.value;
+    }
+  }
+  globalThis.FinalizationRegistry = RegistryFixture;
+  globalThis.WeakRef = WeakRefFixture;
+  try {
+    run(() => registry, references);
+  } finally {
+    globalThis.FinalizationRegistry = OriginalRegistry;
+    globalThis.WeakRef = OriginalWeakRef;
+  }
+}
+
+test.serial(
+  "a collected old value cannot evict a newer entry with the same key",
+  (t) => {
+    withRegistryFixture((getRegistry, references) => {
+      const evicted = [];
+      const cache = new WeakCache({ onEvict: (key) => evicted.push(key) });
+      cache.set("key", { old: true });
+      const registry = getRegistry();
+      const oldHeld = registry.registrations[0].held;
+      references[0].value = undefined;
+      const fresh = { fresh: true };
+      cache.set("key", fresh);
+      registry.callback(oldHeld);
+      t.is(cache.get("key"), fresh);
+      t.deepEqual(evicted, []);
+    });
+  }
+);
+
+test.serial(
+  "deleting one alias keeps another alias registered for eviction",
+  (t) => {
+    withRegistryFixture((getRegistry) => {
+      const evicted = [];
+      const cache = new WeakCache({ onEvict: (key) => evicted.push(key) });
+      const shared = {};
+      cache.set("first", shared);
+      cache.set("second", shared);
+      cache.delete("first");
+      const registry = getRegistry();
+      t.is(registry.registrations.length, 1);
+      registry.callback(registry.registrations[0].held);
+      t.is(cache.size, 0);
+      t.deepEqual(evicted, ["second"]);
+    });
+  }
+);