Support options.normalizeResult callback for wrap (#283)

package.json

@@ -40,6 +40,7 @@
   "devDependencies": {
     "@types/mocha": "^10.0.1",
     "@types/node": "^20.2.5",
+    "@wry/equality": "^0.5.7",
     "mocha": "^10.2.0",
     "rimraf": "^5.0.0",
     "rollup": "^3.20.0",

rollup.config.js

@@ -5,6 +5,7 @@ const globals = {
   tslib: "tslib",
   assert: "assert",
   crypto: "crypto",
+  "@wry/equality": "wryEquality",
   "@wry/context": "wryContext",
   "@wry/trie": "wryTrie",
   "@wry/caches": "wryCaches",

src/entry.ts

@@ -50,6 +50,7 @@ export type AnyEntry = Entry<any, any>;
 export class Entry<TArgs extends any[], TValue> {
   public static count = 0;
 
+  public normalizeResult: OptimisticWrapOptions<TArgs, any, any, TValue>["normalizeResult"];
   public subscribe: OptimisticWrapOptions<TArgs>["subscribe"];
   public unsubscribe: Unsubscribable["unsubscribe"];
 
@@ -95,7 +96,6 @@ export class Entry<TArgs extends any[], TValue> {
   public setDirty() {
     if (this.dirty) return;
     this.dirty = true;
-    this.value.length = 0;
     reportDirty(this);
     // We can go ahead and unsubscribe here, since any further dirty
     // notifications we receive will be redundant, and unsubscribing may
@@ -191,15 +191,38 @@ function reallyRecompute(entry: AnyEntry, args: any[]) {
 
 function recomputeNewValue(entry: AnyEntry, args: any[]) {
   entry.recomputing = true;
-  // Set entry.value as unknown.
+
+  const { normalizeResult } = entry;
+  let oldValueCopy: Value<any> | undefined;
+  if (normalizeResult && entry.value.length === 1) {
+    oldValueCopy = valueCopy(entry.value);
+  }
+
+  // Make entry.value an empty array, representing an unknown value.
   entry.value.length = 0;
+
   try {
     // If entry.fn succeeds, entry.value will become a normal Value.
     entry.value[0] = entry.fn.apply(null, args);
+
+    // If we have a viable oldValueCopy to compare with the (successfully
+    // recomputed) new entry.value, and they are not already === identical, give
+    // normalizeResult a chance to pick/choose/reuse parts of oldValueCopy[0]
+    // and/or entry.value[0] to determine the final cached entry.value.
+    if (normalizeResult && oldValueCopy && !valueIs(oldValueCopy, entry.value)) {
+      try {
+        entry.value[0] = normalizeResult(entry.value[0], oldValueCopy[0]);
+      } catch {
+        // If normalizeResult throws, just use the newer value, rather than
+        // saving the exception as entry.value[1].
+      }
+    }
+
   } catch (e) {
-    // If entry.fn throws, entry.value will become exceptional.
+    // If entry.fn throws, entry.value will hold that exception.
     entry.value[1] = e;
   }
+
   // Either way, this line is always reached.
   entry.recomputing = false;
 }

src/index.ts

@@ -112,6 +112,9 @@ export type OptimisticWrapOptions<
   // 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: NoInfer<TKeyArgs>) => TCacheKey | undefined;
+  // Called when a new value is computed to allow efficient normalization of
+  // results over time, for example by returning older if equal(newer, older).
+  normalizeResult?: (newer: TResult, older: TResult) => TResult;
   // If provided, the subscribe function should either return an unsubscribe
   // function or return nothing.
   subscribe?: (...args: TArgs) => void | (() => any);
@@ -137,8 +140,9 @@ export function wrap<
   TCacheKey = any,
 >(originalFunction: (...args: TArgs) => TResult, {
   max = Math.pow(2, 16),
-  makeCacheKey = (defaultMakeCacheKey as () => TCacheKey),
   keyArgs,
+  makeCacheKey = (defaultMakeCacheKey as () => TCacheKey),
+  normalizeResult,
   subscribe,
   cache: cacheOption = StrongCache,
 }: OptimisticWrapOptions<TArgs, TKeyArgs, TCacheKey, TResult> = Object.create(null)) {
@@ -160,6 +164,7 @@ export function wrap<
     let entry = cache.get(key)!;
     if (!entry) {
       cache.set(key, entry = new Entry(originalFunction));
+      entry.normalizeResult = normalizeResult;
       entry.subscribe = subscribe;
       // Give the Entry the ability to trigger cache.delete(key), even though
       // the Entry itself does not know about key or cache.
@@ -195,8 +200,9 @@ export function wrap<
 
   Object.freeze(optimistic.options = {
     max,
-    makeCacheKey,
     keyArgs,
+    makeCacheKey,
+    normalizeResult,
     subscribe,
     cache,
   });

src/tests/api.ts

@@ -6,8 +6,10 @@ import {
   OptimisticWrapperFunction,
   CommonCache,
 } from "../index";
+import { equal } from '@wry/equality';
 import { wrapYieldingFiberMethods } from '@wry/context';
 import { dep } from "../dep";
+import { permutations } from "./test-utils";
 
 type NumThunk = OptimisticWrapperFunction<[], number>;
 
@@ -569,24 +571,49 @@ describe("optimism", function () {
     const keyArgs: () => [] = () => [];
     function makeCacheKey() { return "constant"; }
     function subscribe() {}
+    let normalizeCalls: [number, number][] = [];
+    function normalizeResult(newer: number, older: number) {
+      normalizeCalls.push([newer, older]);
+      return newer;
+    }
 
     let counter1 = 0;
     const wrapped = wrap(() => ++counter1, {
       max: 10,
       keyArgs,
       makeCacheKey,
+      normalizeResult,
       subscribe,
     });
     assert.strictEqual(wrapped.options.max, 10);
     assert.strictEqual(wrapped.options.keyArgs, keyArgs);
     assert.strictEqual(wrapped.options.makeCacheKey, makeCacheKey);
+    assert.strictEqual(wrapped.options.normalizeResult, normalizeResult);
     assert.strictEqual(wrapped.options.subscribe, subscribe);
 
+    assert.deepEqual(normalizeCalls, []);
+    assert.strictEqual(wrapped(), 1);
+    assert.deepEqual(normalizeCalls, []);
+    assert.strictEqual(wrapped(), 1);
+    assert.deepEqual(normalizeCalls, []);
+    wrapped.dirty();
+    assert.deepEqual(normalizeCalls, []);
+    assert.strictEqual(wrapped(), 2);
+    assert.deepEqual(normalizeCalls, [[2, 1]]);
+    assert.strictEqual(wrapped(), 2);
+    wrapped.dirty();
+    assert.strictEqual(wrapped(), 3);
+    assert.deepEqual(normalizeCalls, [[2, 1], [3, 2]]);
+    assert.strictEqual(wrapped(), 3);
+    assert.deepEqual(normalizeCalls, [[2, 1], [3, 2]]);
+    assert.strictEqual(wrapped(), 3);
+
     let counter2 = 0;
     const wrappedWithDefaults = wrap(() => ++counter2);
     assert.strictEqual(wrappedWithDefaults.options.max, Math.pow(2, 16));
     assert.strictEqual(wrappedWithDefaults.options.keyArgs, void 0);
     assert.strictEqual(typeof wrappedWithDefaults.options.makeCacheKey, "function");
+    assert.strictEqual(wrappedWithDefaults.options.normalizeResult, void 0);
     assert.strictEqual(wrappedWithDefaults.options.subscribe, void 0);
   });
 
@@ -801,4 +828,141 @@ describe("optimism", function () {
     d.dirty("shared", "forget");
     assert.strictEqual(size(), 0);
   });
+
+  describe("wrapOptions.normalizeResult", function () {
+    it("can normalize array results", function () {
+      const normalizeArgs: [number[], number[]][] = [];
+      const range = wrap((n: number) => {
+        let result = [];
+        for (let i = 0; i < n; ++i) {
+          result[i] = i;
+        }
+        return result;
+      }, {
+        normalizeResult(newer, older) {
+          normalizeArgs.push([newer, older]);
+          return equal(newer, older) ? older : newer;
+        },
+      });
+
+      const r3a = range(3);
+      assert.deepStrictEqual(r3a, [0, 1, 2]);
+      // Nothing surprising, just regular caching.
+      assert.strictEqual(r3a, range(3));
+
+      // Force range(3) to be recomputed below.
+      range.dirty(3);
+
+      const r3b = range(3);
+      assert.deepStrictEqual(r3b, [0, 1, 2]);
+
+      assert.strictEqual(r3a, r3b);
+
+      assert.deepStrictEqual(normalizeArgs, [
+        [r3b, r3a],
+      ]);
+      // Though r3a and r3b ended up ===, the normalizeResult callback should
+      // have been called with two !== arrays.
+      assert.notStrictEqual(
+        normalizeArgs[0][0],
+        normalizeArgs[0][1],
+      );
+    });
+
+    it("can normalize recursive array results", function () {
+      const range = wrap((n: number): number[] => {
+        if (n <= 0) return [];
+        return range(n - 1).concat(n - 1);
+      }, {
+        normalizeResult: (newer, older) => equal(newer, older) ? older : newer,
+      });
+
+      const ranges = [
+        range(0),
+        range(1),
+        range(2),
+        range(3),
+        range(4),
+      ];
+
+      assert.deepStrictEqual(ranges[0], []);
+      assert.deepStrictEqual(ranges[1], [0]);
+      assert.deepStrictEqual(ranges[2], [0, 1]);
+      assert.deepStrictEqual(ranges[3], [0, 1, 2]);
+      assert.deepStrictEqual(ranges[4], [0, 1, 2, 3]);
+
+      const perms = permutations(ranges[4]);
+      assert.strictEqual(perms.length, 4 * 3 * 2 * 1);
+
+      // For each permutation of the range sizes, check that strict equality
+      // holds for r[i] and range(i) for all i after dirtying each number.
+      let count = 0;
+      perms.forEach(perm => {
+        perm.forEach(toDirty => {
+          range.dirty(toDirty);
+          perm.forEach(i => {
+            assert.strictEqual(ranges[i], range(i));
+            ++count;
+          });
+        })
+      });
+      assert.strictEqual(count, perms.length * 4 * 4);
+    });
+
+    it("exceptions thrown by normalizeResult are ignored", function () {
+      const normalizeCalls: [string | number, string | number][] = [];
+
+      const maybeThrow = wrap((value: string | number, shouldThrow: boolean) => {
+        if (shouldThrow) throw value;
+        return value;
+      }, {
+        makeCacheKey(value, shouldThrow) {
+          return JSON.stringify({
+            // Coerce the value to a string so we can trigger normalizeResult
+            // using either 2 or "2" below.
+            value: String(value),
+            shouldThrow,
+          });
+        },
+        normalizeResult(a, b) {
+          normalizeCalls.push([a, b]);
+          throw new Error("from normalizeResult (expected)");
+        },
+      });
+
+      assert.strictEqual(maybeThrow(1, false), 1);
+      assert.strictEqual(maybeThrow(2, false), 2);
+
+      maybeThrow.dirty(2, false);
+      assert.strictEqual(maybeThrow("2", false), "2");
+      assert.strictEqual(maybeThrow(2, false), "2");
+      maybeThrow.dirty(2, false);
+      assert.strictEqual(maybeThrow(2, false), 2);
+      assert.strictEqual(maybeThrow("2", false), 2);
+
+      assert.throws(
+        () => maybeThrow(3, true),
+        error => error === 3,
+      );
+
+      assert.throws(
+        () => maybeThrow("3", true),
+        // Still 3 because the previous maybeThrow(3, true) exception is cached.
+        error => error === 3,
+      );
+
+      maybeThrow.dirty(3, true);
+      assert.throws(
+        () => maybeThrow("3", true),
+        error => error === "3",
+      );
+
+      // Even though the exception thrown by normalizeResult was ignored, check
+      // that it was in fact called (twice).
+      assert.deepStrictEqual(normalizeCalls, [
+        ["2", 2],
+        [2, "2"],
+      ]);
+    });
+  });
 });

src/tests/test-utils.ts

@@ -0,0 +1,14 @@
+export function permutations<T>(array: T[], start = 0): T[][] {
+  if (start === array.length) return [[]];
+  const item = array[start];
+  const results: T[][] = [];
+  permutations<T>(array, start + 1).forEach(perm => {
+    perm.forEach((_, i) => {
+      const copy = perm.slice(0);
+      copy.splice(i, 0, item);
+      results.push(copy);
+    });
+    results.push(perm.concat(item));
+  });
+  return results;
+}