Implement lightweight OptimisticDependencyFunction<TKey> primitive.

src/entry.ts

@@ -1,8 +1,7 @@
 import { parentEntrySlot } from "./context";
 import { OptimisticWrapOptions } from "./index";
 
-const reusableEmptyArray: AnyEntry[] = [];
-const emptySetPool: Set<AnyEntry>[] = [];
+const emptySetPool: Set<any>[] = [];
 const POOL_TARGET_SIZE = 100;
 
 // Since this package might be used browsers, we should avoid using the
@@ -51,7 +50,6 @@ export class Entry<TArgs extends any[], TValue> {
 
   public subscribe: OptimisticWrapOptions<TArgs>["subscribe"];
   public unsubscribe?: () => any;
-  public reportOrphan?: (this: Entry<TArgs, TValue>) => any;
 
   public readonly parents = new Set<AnyEntry>();
   public readonly childValues = new Map<AnyEntry, Value<any>>();
@@ -80,14 +78,7 @@ export class Entry<TArgs extends any[], TValue> {
   // (3) valueGet(this.value) is usually returned without recomputation.
   public recompute(): TValue {
     assert(! this.recomputing, "already recomputing");
-
-    if (! rememberParent(this) && maybeReportOrphan(this)) {
-      // The recipient of the entry.reportOrphan callback decided to dispose
-      // of this orphan entry by calling entry.dispose(), so we don't need to
-      // (and should not) proceed with the recomputation.
-      return void 0 as any;
-    }
-
+    rememberParent(this);
     return mightBeDirty(this)
       ? reallyRecompute(this)
       : valueGet(this.value);
@@ -98,14 +89,15 @@ export class Entry<TArgs extends any[], TValue> {
     this.dirty = true;
     this.value.length = 0;
     reportDirty(this);
+    forgetChildren(this);
     // We can go ahead and unsubscribe here, since any further dirty
     // notifications we receive will be redundant, and unsubscribing may
     // free up some resources, e.g. file watchers.
     maybeUnsubscribe(this);
   }
 
   public dispose() {
-    forgetChildren(this).forEach(maybeReportOrphan);
+    forgetChildren(this);
     maybeUnsubscribe(this);
 
     // Because this entry has been kicked out of the cache (in index.js),
@@ -124,6 +116,25 @@ export class Entry<TArgs extends any[], TValue> {
       forgetChild(parent, this);
     });
   }
+
+  private sets: Set<Set<AnyEntry>> | null = null;
+
+  public addToSet(entrySet: Set<AnyEntry>) {
+    entrySet.add(this);
+    if (! this.sets) {
+      this.sets = emptySetPool.pop() || new Set<Set<AnyEntry>>();
+    }
+    this.sets.add(entrySet);
+  }
+
+  public removeFromSets() {
+    if (this.sets) {
+      this.sets.forEach(set => set.delete(this));
+      this.sets.clear();
+      emptySetPool.push(this.sets);
+      this.sets = null;
+    }
+  }
 }
 
 function rememberParent(child: AnyEntry) {
@@ -146,10 +157,7 @@ function rememberParent(child: AnyEntry) {
 }
 
 function reallyRecompute(entry: AnyEntry) {
-  // Since this recomputation is likely to re-remember some of this
-  // entry's children, we forget our children here but do not call
-  // maybeReportOrphan until after the recomputation finishes.
-  const originalChildren = forgetChildren(entry);
+  forgetChildren(entry);
 
   // Set entry as the parent entry while calling recomputeNewValue(entry).
   parentEntrySlot.withValue(entry, recomputeNewValue, [entry]);
@@ -160,11 +168,6 @@ function reallyRecompute(entry: AnyEntry) {
     setClean(entry);
   }
 
-  // Now that we've had a chance to re-remember any children that were
-  // involved in the recomputation, we can safely report any orphan
-  // children that remain.
-  originalChildren.forEach(maybeReportOrphan);
-
   return valueGet(entry.value);
 }
 
@@ -264,35 +267,22 @@ function removeDirtyChild(parent: AnyEntry, child: AnyEntry) {
   }
 }
 
-// If the given entry has a reportOrphan method, and no remaining parents,
-// call entry.reportOrphan and return true iff it returns true. The
-// reportOrphan function should return true to indicate entry.dispose()
-// has been called, and the entry has been removed from any other caches
-// (see index.js for the only current example).
-function maybeReportOrphan(entry: AnyEntry) {
-  return entry.parents.size === 0 &&
-    typeof entry.reportOrphan === "function" &&
-    entry.reportOrphan() === true;
-}
-
 // Removes all children from this entry and returns an array of the
 // removed children.
 function forgetChildren(parent: AnyEntry) {
-  let children = reusableEmptyArray;
-
   if (parent.childValues.size > 0) {
-    children = [];
     parent.childValues.forEach((_value, child) => {
       forgetChild(parent, child);
-      children.push(child);
     });
   }
 
+  // Remove this parent Entry from any sets to which it was added by the
+  // addToSet method.
+  parent.removeFromSets();
+
   // After we forget all our children, this.dirtyChildren must be empty
   // and therefore must have been reset to null.
   assert(parent.dirtyChildren === null);
-
-  return children;
 }
 
 function forgetChild(parent: AnyEntry, child: AnyEntry) {

src/index.ts

@@ -46,14 +46,15 @@ export type OptimisticWrapperFunction<
   dirty: (...args: TArgs) => void;
 };
 
+type OptimisticDependencyFunction<TKey> =
+  ((key: TKey) => void) & {
+    dirty: (key: TKey) => void;
+  };
+
 export type OptimisticWrapOptions<TArgs extends any[]> = {
   // The maximum number of cache entries that should be retained before the
   // cache begins evicting the oldest ones.
   max?: number;
-  // If a wrapped function is "disposable," then its creator does not
-  // care about its return value, and it should be removed from the cache
-  // immediately when it no longer has any parents that depend on it.
-  disposable?: boolean;
   // 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.
@@ -77,19 +78,9 @@ export function wrap<
     entry => entry.dispose(),
   );
 
-  const disposable = !! options.disposable;
   const makeCacheKey = options.makeCacheKey || defaultMakeCacheKey;
 
   function optimistic(): TResult {
-    if (disposable && ! parentEntrySlot.hasValue()) {
-      // If there's no current parent computation, and this wrapped
-      // function is disposable (meaning we don't care about entry.value,
-      // just dependency tracking), then we can short-cut everything else
-      // in this function, because entry.recompute() is going to recycle
-      // the entry object without recomputing anything, anyway.
-      return void 0 as any;
-    }
-
     const key = makeCacheKey.apply(null, arguments as any);
     if (key === void 0) {
       return originalFunction.apply(null, arguments as any);
@@ -104,9 +95,6 @@ export function wrap<
       entry = new Entry<TArgs, TResult>(originalFunction, args);
       cache.set(key, entry);
       entry.subscribe = options.subscribe;
-      if (disposable) {
-        entry.reportOrphan = () => cache.delete(key);
-      }
     }
 
     const value = entry.recompute();
@@ -125,10 +113,7 @@ export function wrap<
       caches.clear();
     }
 
-    // If options.disposable is truthy, the caller of wrap is telling us
-    // they don't care about the result of entry.recompute(), so we should
-    // avoid returning the value, so it won't be accidentally used.
-    return disposable ? void 0 as any : value;
+    return value;
   }
 
   optimistic.dirty = function () {
@@ -141,3 +126,28 @@ export function wrap<
 
   return optimistic as OptimisticWrapperFunction<TArgs, TResult>;
 }
+
+export function dep<TKey>() {
+  const parentEntriesByKey = new Map<TKey, Set<AnyEntry>>();
+
+  function depend(key: TKey) {
+    const parent = parentEntrySlot.getValue();
+    if (parent) {
+      let parentEntrySet = parentEntriesByKey.get(key);
+      if (!parentEntrySet) {
+        parentEntriesByKey.set(key, parentEntrySet = new Set);
+      }
+      parent.addToSet(parentEntrySet);
+    }
+  }
+
+  depend.dirty = function(key: TKey) {
+    const parentEntrySet = parentEntriesByKey.get(key);
+    if (parentEntrySet) {
+      parentEntrySet.forEach(entry => entry.setDirty());
+      parentEntriesByKey.delete(key);
+    }
+  };
+
+  return depend as OptimisticDependencyFunction<TKey>;
+}

src/tests/api.ts

@@ -498,73 +498,6 @@ describe("optimism", function () {
     assert.strictEqual(child(), "parent");
   });
 
-  it("supports disposable wrapped functions", function () {
-    let dependCallCount = 0;
-    const depend = wrap(function (n?: number) {
-      return ++dependCallCount;
-    }, {
-      disposable: true
-    });
-
-    assert.strictEqual(typeof depend(), "undefined");
-    assert.strictEqual(dependCallCount, 0);
-
-    let parentCallCount = 0;
-    const parent = wrap(function () {
-      ++parentCallCount;
-      assert.strictEqual(typeof depend(1), "undefined");
-      assert.strictEqual(typeof depend(2), "undefined");
-    });
-
-    parent();
-    assert.strictEqual(parentCallCount, 1);
-    assert.strictEqual(dependCallCount, 2);
-
-    parent();
-    assert.strictEqual(parentCallCount, 1);
-    assert.strictEqual(dependCallCount, 2);
-
-    depend.dirty(1);
-    parent();
-    assert.strictEqual(parentCallCount, 2);
-    assert.strictEqual(dependCallCount, 3);
-
-    depend.dirty(2);
-    parent();
-    assert.strictEqual(parentCallCount, 3);
-    assert.strictEqual(dependCallCount, 4);
-
-    parent();
-    assert.strictEqual(parentCallCount, 3);
-    assert.strictEqual(dependCallCount, 4);
-
-    parent.dirty();
-    parent();
-    assert.strictEqual(parentCallCount, 4);
-    assert.strictEqual(dependCallCount, 4);
-
-    depend.dirty(1);
-    depend(1);
-    // No change to dependCallCount because depend is called outside of
-    // any parent computation, and depend is disposable.
-    assert.strictEqual(dependCallCount, 4);
-    depend(2);
-    assert.strictEqual(dependCallCount, 4);
-
-    depend.dirty(2);
-    depend(1);
-    // Again, no change because depend is disposable.
-    assert.strictEqual(dependCallCount, 4);
-    depend(2);
-    assert.strictEqual(dependCallCount, 4);
-
-    parent();
-    // Now, since both depend(1) and depend(2) are dirty, calling them in
-    // the context of the parent() computation results in two more
-    // increments of dependCallCount.
-    assert.strictEqual(dependCallCount, 6);
-  });
-
   it("is not confused by eviction during recomputation", function () {
     const fib: OptimisticWrapperFunction<[number], number> =
       wrap(function (n: number) {