Support options.keyArgs to complement options.makeCacheKey.

This optional function takes the raw arguments passed to an
OptimisticFunctionWrapper and simplifies/filters/reduces them down to an
array of arguments sufficient to identify the cache result. This
simplified array of arguments is then passed to options.makeCacheKey (or
defaultMakeCacheKey if no options.makeCacheKey was provided).

The returned array serves as the arguments list for any function that only
needs to compute a cache key (and thus does not need the full raw
arguments), such as wrapper.dirty and options.makeCacheKey. This makes it
possible to dirty a cached function using less information than would be
required to call it properly, which is useful in some cases.

In some sense, options.keyArgs and options.makeCacheKey divide up the work
of computing cache keys from raw arguments. Both functions are optional,
and any combination of the two that satisfies the type system is allowed.
If the sequence of arguments returned by options.keyArgs can be passed
directly to defaultMakeCacheKey, you may not need to define a custom
options.makeCacheKey function anymore.

src/index.ts

@@ -46,20 +46,27 @@ export { KeyTrie }
 export type OptimisticWrapperFunction<
   TArgs extends any[],
   TResult,
+  TKeyArgs extends any[] = TArgs,
 > = ((...args: TArgs) => TResult) & {
   // The .dirty(...) method of an optimistic function takes exactly the
   // same parameter types as the original function.
-  dirty: (...args: TArgs) => void;
+  dirty: (...args: TKeyArgs) => void;
 };
 
-export type OptimisticWrapOptions<TArgs extends any[]> = {
+export type OptimisticWrapOptions<
+  TArgs extends any[],
+  TKeyArgs extends any[] = TArgs,
+> = {
   // The maximum number of cache entries that should be retained before the
   // cache begins evicting the oldest ones.
   max?: number;
+  // Transform the raw arguments to some other type of array, which will then
+  // be passed to makeCacheKey.
+  keyArgs?: (...args: TArgs) => TKeyArgs;
   // 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: TArgs) => TCacheKey;
+  makeCacheKey?: (...args: TKeyArgs) => TCacheKey;
   // If provided, the subscribe function should either return an unsubscribe
   // function or return nothing.
   subscribe?: (...args: TArgs) => void | (() => any);
@@ -70,19 +77,21 @@ const caches = new Set<Cache<TCacheKey, AnyEntry>>();
 export function wrap<
   TArgs extends any[],
   TResult,
+  TKeyArgs extends any[] = TArgs,
 >(
   originalFunction: (...args: TArgs) => TResult,
-  options: OptimisticWrapOptions<TArgs> = Object.create(null),
+  options: OptimisticWrapOptions<TArgs, TKeyArgs> = Object.create(null),
 ) {
   const cache = new Cache<TCacheKey, Entry<TArgs, TResult>>(
     options.max || Math.pow(2, 16),
     entry => entry.dispose(),
   );
 
+  const keyArgs = options.keyArgs || ((...args: TArgs): TKeyArgs => args as any);
   const makeCacheKey = options.makeCacheKey || defaultMakeCacheKey;
 
   function optimistic(): TResult {
-    const key = makeCacheKey.apply(null, arguments as any);
+    const key = makeCacheKey.apply(null, keyArgs.apply(null, arguments as any));
     if (key === void 0) {
       return originalFunction.apply(null, arguments as any);
     }
@@ -118,12 +127,12 @@ export function wrap<
   }
 
   optimistic.dirty = function () {
-    const key = makeCacheKey.apply(null, arguments as any);
+    const key = makeCacheKey.apply(null, keyArgs.apply(null, arguments as any));
     const child = key !== void 0 && cache.get(key);
     if (child) {
       child.setDirty();
     }
   };
 
-  return optimistic as OptimisticWrapperFunction<TArgs, TResult>;
+  return optimistic as OptimisticWrapperFunction<TArgs, TResult, TKeyArgs>;
 }

src/tests/api.ts

@@ -470,6 +470,36 @@ describe("optimism", function () {
     check();
   });
 
+  it("supports options.keyArgs", function () {
+    const sumNums = wrap((...args: any[]) => ({
+      sum: args.reduce(
+        (sum, arg) => typeof arg === "number" ? arg + sum : sum,
+        0,
+      ) as number,
+    }), {
+      keyArgs(...args) {
+        return args.filter(arg => typeof arg === "number");
+      },
+    });
+
+    assert.strictEqual(sumNums().sum, 0);
+    assert.strictEqual(sumNums("asdf", true, sumNums).sum, 0);
+
+    const sumObj1 = sumNums(1, "zxcv", true, 2, false, 3);
+    assert.strictEqual(sumObj1.sum, 6);
+    // These results are === sumObj1 because the numbers involved are identical.
+    assert.strictEqual(sumNums(1, 2, 3), sumObj1);
+    assert.strictEqual(sumNums("qwer", 1, 2, true, 3, [3]), sumObj1);
+    assert.strictEqual(sumNums("backwards", 3, 2, 1).sum, 6);
+    assert.notStrictEqual(sumNums("backwards", 3, 2, 1), sumObj1);
+
+    sumNums.dirty(1, 2, 3);
+    const sumObj2 = sumNums(1, 2, 3);
+    assert.strictEqual(sumObj2.sum, 6);
+    assert.notStrictEqual(sumObj2, sumObj1);
+    assert.strictEqual(sumNums("a", 1, "b", 2, "c", 3), sumObj2);
+  });
+
   it("tolerates cycles when propagating dirty/clean signals", function () {
     let counter = 0;
     const dep = wrap(() => ++counter);