Implement options.normalizeResult callback for wrap.

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

@@ -92,6 +92,7 @@ export type OptimisticWrapOptions<
   TArgs extends any[],
   TKeyArgs extends any[] = TArgs,
   TCacheKey = any,
+  TResult = any,
 > = {
   // The maximum number of cache entries that should be retained before the
   // cache begins evicting the oldest ones.
@@ -103,6 +104,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: TKeyArgs) => TCacheKey;
+  // 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);
@@ -117,7 +121,12 @@ export function wrap<
   TCacheKey = any,
 >(
   originalFunction: (...args: TArgs) => TResult,
-  options: OptimisticWrapOptions<TArgs, TKeyArgs> = Object.create(null),
+  options: OptimisticWrapOptions<
+    TArgs,
+    TKeyArgs,
+    TCacheKey,
+    TResult
+  > = Object.create(null),
 ) {
   const cache = new Cache<TCacheKey, Entry<TArgs, TResult>>(
     options.max || Math.pow(2, 16),
@@ -141,6 +150,7 @@ export function wrap<
     let entry = cache.get(key)!;
     if (!entry) {
       cache.set(key, entry = new Entry(originalFunction));
+      entry.normalizeResult = options.normalizeResult;
       entry.subscribe = options.subscribe;
       // Give the Entry the ability to trigger cache.delete(key), even though
       // the Entry itself does not know about key or cache.

src/tests/api.ts

@@ -5,8 +5,10 @@ import {
   defaultMakeCacheKey,
   OptimisticWrapperFunction,
 } 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>;
 
@@ -705,4 +707,141 @@ describe("optimism", function () {
     d.dirty("shared", "forget");
     assert.strictEqual(fib.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;
+}