perf(cache): optimize _serializeArgList key generation

cache/_serialize_arg_list.ts

@@ -10,13 +10,20 @@ import type { MemoizationCache } from "./memoize.ts";
  * @param cache The cache for which the keys will be used.
  * @returns `getKey`, the function for getting cache keys.
  */
-
 export function _serializeArgList<Return>(
   cache: MemoizationCache<unknown, Return>,
 ): (this: unknown, ...args: unknown[]) => string {
+  // Three cooperating data structures track weak (reference-type) arguments:
+  //   1. weakKeyToKeySegmentCache: WeakMap from object/symbol → segment id
+  //      (e.g. `"{0}"`) so the same reference always maps to the same segment.
+  //   2. weakKeySegmentToKeyCache: Map from segment id → set of composite cache
+  //      keys that contain that segment, used by the finalization callback.
+  //   3. registry (FinalizationRegistry): when a weak key is garbage-collected,
+  //      looks up its segment in (2) and deletes all associated entries from the
+  //      caller-provided cache.
   const weakKeyToKeySegmentCache = new WeakMap<WeakKey, string>();
-  const weakKeySegmentToKeyCache = new Map<string, string[]>();
-  let i = 0;
+  const weakKeySegmentToKeyCache = new Map<string, Set<string>>();
+  let nextWeakKeyId = 0;
 
   const registry = new FinalizationRegistry<string>((keySegment) => {
     for (const key of weakKeySegmentToKeyCache.get(keySegment) ?? []) {
@@ -46,42 +53,43 @@ export function _serializeArgList<Return>(
         return JSON.stringify(arg);
       }
 
-      try {
-        assertWeakKey(arg);
-      } catch {
-        if (typeof arg === "symbol") {
+      if (typeof arg === "symbol") {
+        try {
+          new WeakRef(arg);
+        } catch {
           return `Symbol.for(${JSON.stringify(arg.description)})`;
         }
-        // Non-weak keys other than `Symbol.for(...)` are handled by the branches above.
+      }
+
+      try {
+        new WeakRef(arg as WeakKey);
+      } catch {
         throw new Error(
           "Should be unreachable: please open an issue at https://github.com/denoland/std/issues/new",
         );
       }
 
-      if (!weakKeyToKeySegmentCache.has(arg)) {
-        const keySegment = `{${i++}}`;
-        weakKeySegments.push(keySegment);
-        registry.register(arg, keySegment);
-        weakKeyToKeySegmentCache.set(arg, keySegment);
+      let keySegment = weakKeyToKeySegmentCache.get(arg as WeakKey);
+      if (keySegment === undefined) {
+        keySegment = `{${nextWeakKeyId++}}`;
+        registry.register(arg as WeakKey, keySegment);
+        weakKeyToKeySegmentCache.set(arg as WeakKey, keySegment);
       }
-
-      const keySegment = weakKeyToKeySegmentCache.get(arg)!;
       weakKeySegments.push(keySegment);
       return keySegment;
     });
 
     const key = keySegments.join(",");
 
     for (const keySegment of weakKeySegments) {
-      const keys = weakKeySegmentToKeyCache.get(keySegment) ?? [];
-      keys.push(key);
-      weakKeySegmentToKeyCache.set(keySegment, keys);
+      let keys = weakKeySegmentToKeyCache.get(keySegment);
+      if (keys === undefined) {
+        keys = new Set();
+        weakKeySegmentToKeyCache.set(keySegment, keys);
+      }
+      keys.add(key);
     }
 
     return key;
   };
 }
-
-function assertWeakKey(arg: unknown): asserts arg is WeakKey {
-  new WeakRef(arg as WeakKey);
-}

cache/_serialize_arg_list_test.ts

@@ -83,6 +83,17 @@ Deno.test("_serializeArgList() gives same results as SameValueZero algorithm", a
   });
 });
 
+Deno.test("_serializeArgList() serializes function arguments as weak keys", () => {
+  const getKey = _serializeArgList(new Map());
+  const fn1 = () => {};
+  const fn2 = () => {};
+
+  assertEquals(getKey(fn1), "undefined,{0}");
+  assertEquals(getKey(fn1), "undefined,{0}");
+  assertEquals(getKey(fn2), "undefined,{1}");
+  assertEquals(getKey(fn1, fn2), "undefined,{0},{1}");
+});
+
 Deno.test("_serializeArgList() discriminates on `this` arg", () => {
   const getKey = _serializeArgList(new Map());
   const obj1 = {};
@@ -94,6 +105,42 @@ Deno.test("_serializeArgList() discriminates on `this` arg", () => {
   assertEquals(getKey.call(obj1, obj2), "{0},{1}");
 });
 
+Deno.test("_serializeArgList() cleans up cache entries when finalization callback fires", () => {
+  let cleanupCallback: (keySegment: string) => void;
+
+  const OriginalFinalizationRegistry = FinalizationRegistry;
+  // deno-lint-ignore no-explicit-any
+  (globalThis as any).FinalizationRegistry = function (
+    cb: (keySegment: string) => void,
+  ) {
+    cleanupCallback = cb;
+    return { register() {}, unregister() {} };
+  };
+
+  try {
+    const cache = new Map();
+    const getKey = _serializeArgList(cache);
+
+    const obj = {};
+    const k1 = getKey(obj);
+    const k2 = getKey(obj, "x");
+    cache.set(k1, "v1");
+    cache.set(k2, "v2");
+    cache.set("unrelated", "v3");
+
+    assertEquals(cache.size, 3);
+
+    cleanupCallback!("{0}");
+
+    assertEquals(cache.size, 1);
+    assertEquals(cache.has("unrelated"), true);
+    assertEquals(cache.has(k1), false);
+    assertEquals(cache.has(k2), false);
+  } finally {
+    globalThis.FinalizationRegistry = OriginalFinalizationRegistry;
+  }
+});
+
 Deno.test("_serializeArgList() allows garbage collection for weak keys", async () => {
   // @ts-expect-error - Triggering true garbage collection is only available
   // with `--v8-flags="--expose-gc"`, so we mock `FinalizationRegistry` with