Support options.keyArgs to complement options.makeCacheKey.

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