Support [un]subscribe API for lightweight deps.

src/dep.ts

@@ -0,0 +1,46 @@
+import { AnyEntry } from "./entry";
+import { OptimisticWrapOptions } from ".";
+import { parentEntrySlot } from "./context";
+import { Unsubscribable, maybeUnsubscribe } from "./helpers";
+
+export type OptimisticDependencyFunction<TKey> =
+  ((key: TKey) => void) & {
+    dirty: (key: TKey) => void;
+  };
+
+export type Dep<TKey> = Set<AnyEntry> & {
+  subscribe: OptimisticWrapOptions<[TKey]>["subscribe"];
+} & Unsubscribable;
+
+export function dep<TKey>(options?: {
+  subscribe: Dep<TKey>["subscribe"];
+}) {
+  const depsByKey = new Map<TKey, Dep<TKey>>();
+  const subscribe = options && options.subscribe;
+
+  function depend(key: TKey) {
+    const parent = parentEntrySlot.getValue();
+    if (parent) {
+      let dep = depsByKey.get(key);
+      if (!dep) {
+        depsByKey.set(key, dep = new Set as Dep<TKey>);
+      }
+      parent.dependOn(dep);
+      if (typeof subscribe === "function") {
+        maybeUnsubscribe(dep);
+        dep.unsubscribe = subscribe(key);
+      }
+    }
+  }
+
+  depend.dirty = function(key: TKey) {
+    const dep = depsByKey.get(key);
+    if (dep) {
+      dep.forEach(entry => entry.setDirty());
+      depsByKey.delete(key);
+      maybeUnsubscribe(dep);
+    }
+  };
+
+  return depend as OptimisticDependencyFunction<TKey>;
+}

src/entry.ts

@@ -1,5 +1,7 @@
 import { parentEntrySlot } from "./context";
 import { OptimisticWrapOptions } from "./index";
+import { Dep } from "./dep";
+import { maybeUnsubscribe, Unsubscribable } from "./helpers";
 
 const emptySetPool: Set<any>[] = [];
 const POOL_TARGET_SIZE = 100;
@@ -49,7 +51,7 @@ export class Entry<TArgs extends any[], TValue> {
   public static count = 0;
 
   public subscribe: OptimisticWrapOptions<TArgs>["subscribe"];
-  public unsubscribe?: () => any;
+  public unsubscribe: Unsubscribable["unsubscribe"];
 
   public readonly parents = new Set<AnyEntry>();
   public readonly childValues = new Map<AnyEntry, Value<any>>();
@@ -117,22 +119,22 @@ export class Entry<TArgs extends any[], TValue> {
     });
   }
 
-  private sets: Set<Set<AnyEntry>> | null = null;
+  private deps: Set<Dep<any>> | null = null;
 
-  public addToSet(entrySet: Set<AnyEntry>) {
-    entrySet.add(this);
-    if (! this.sets) {
-      this.sets = emptySetPool.pop() || new Set<Set<AnyEntry>>();
+  public dependOn(dep: Dep<any>) {
+    dep.add(this);
+    if (! this.deps) {
+      this.deps = emptySetPool.pop() || new Set<Set<AnyEntry>>();
     }
-    this.sets.add(entrySet);
+    this.deps.add(dep);
   }
 
-  public removeFromSets() {
-    if (this.sets) {
-      this.sets.forEach(set => set.delete(this));
-      this.sets.clear();
-      emptySetPool.push(this.sets);
-      this.sets = null;
+  public forgetDeps() {
+    if (this.deps) {
+      this.deps.forEach(dep => dep.delete(this));
+      this.deps.clear();
+      emptySetPool.push(this.deps);
+      this.deps = null;
     }
   }
 }
@@ -278,7 +280,7 @@ function forgetChildren(parent: AnyEntry) {
 
   // Remove this parent Entry from any sets to which it was added by the
   // addToSet method.
-  parent.removeFromSets();
+  parent.forgetDeps();
 
   // After we forget all our children, this.dirtyChildren must be empty
   // and therefore must have been reset to null.
@@ -310,11 +312,3 @@ function maybeSubscribe(entry: AnyEntry) {
   // function or that it succeeded.
   return true;
 }
-
-function maybeUnsubscribe(entry: AnyEntry) {
-  const { unsubscribe } = entry;
-  if (typeof unsubscribe === "function") {
-    entry.unsubscribe = void 0;
-    unsubscribe();
-  }
-}

src/helpers.ts

@@ -0,0 +1,11 @@
+export type Unsubscribable = {
+  unsubscribe?: void | (() => any);
+}
+
+export function maybeUnsubscribe(entryOrDep: Unsubscribable) {
+  const { unsubscribe } = entryOrDep;
+  if (typeof unsubscribe === "function") {
+    entryOrDep.unsubscribe = void 0;
+    unsubscribe();
+  }
+}

src/index.ts

@@ -16,6 +16,12 @@ export {
   asyncFromGen,
 } from "./context";
 
+// A lighter-weight dependency, similar to OptimisticWrapperFunction, except
+// with only one argument, no makeCacheKey, no wrapped function to recompute,
+// and no result value. Useful for representing dependency leaves in the graph
+// of computation. Subscriptions are supported.
+export { dep } from "./dep";
+
 // Since the Cache uses a Map internally, any value or object reference can
 // be safely used as a key, though common types include object and string.
 export type TCacheKey = any;
@@ -46,11 +52,6 @@ export type OptimisticWrapperFunction<
   dirty: (...args: TArgs) => void;
 };
 
-export 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.
@@ -61,7 +62,7 @@ export type OptimisticWrapOptions<TArgs extends any[]> = {
   makeCacheKey?: (...args: TArgs) => TCacheKey;
   // If provided, the subscribe function should either return an unsubscribe
   // function or return nothing.
-  subscribe?: (...args: TArgs) => (() => any) | undefined;
+  subscribe?: (...args: TArgs) => void | (() => any);
 };
 
 const caches = new Set<Cache<TCacheKey, AnyEntry>>();
@@ -126,28 +127,3 @@ 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/deps.ts

@@ -116,4 +116,68 @@ describe("OptimisticDependencyFunction<TKey>", () => {
     assert.strictEqual(parent("oyez"), 9);
     assert.strictEqual(parent("mlem"), 11);
   });
+
+  it("supports subscribing and unsubscribing", function () {
+    let subscribeCallCount = 0;
+    let unsubscribeCallCount = 0;
+    let parentCallCount = 0;
+
+    function check(counts: {
+      subscribe: number;
+      unsubscribe: number;
+      parent: number;
+    }) {
+      assert.strictEqual(counts.subscribe, subscribeCallCount);
+      assert.strictEqual(counts.unsubscribe, unsubscribeCallCount);
+      assert.strictEqual(counts.parent, parentCallCount);
+    }
+
+    const d = dep({
+      subscribe(key: string) {
+        ++subscribeCallCount;
+        return () => {
+          ++unsubscribeCallCount;
+        };
+      },
+    });
+
+    assert.strictEqual(subscribeCallCount, 0);
+    assert.strictEqual(unsubscribeCallCount, 0);
+
+    const parent = wrap((key: string) => {
+      d(key);
+      return ++parentCallCount;
+    });
+
+    assert.strictEqual(parent("rawr"), 1);
+    check({ subscribe: 1, unsubscribe: 0, parent: 1 });
+    assert.strictEqual(parent("rawr"), 1);
+    check({ subscribe: 1, unsubscribe: 0, parent: 1 });
+    assert.strictEqual(parent("blep"), 2);
+    check({ subscribe: 2, unsubscribe: 0, parent: 2 });
+    assert.strictEqual(parent("rawr"), 1);
+    check({ subscribe: 2, unsubscribe: 0, parent: 2 });
+    assert.strictEqual(parent("blep"), 2);
+    check({ subscribe: 2, unsubscribe: 0, parent: 2 });
+
+    d.dirty("blep");
+    check({ subscribe: 2, unsubscribe: 1, parent: 2 });
+    assert.strictEqual(parent("rawr"), 1);
+    check({ subscribe: 2, unsubscribe: 1, parent: 2 });
+    d.dirty("blep"); // intentionally redundant
+    check({ subscribe: 2, unsubscribe: 1, parent: 2 });
+    assert.strictEqual(parent("blep"), 3);
+    check({ subscribe: 3, unsubscribe: 1, parent: 3 });
+    assert.strictEqual(parent("blep"), 3);
+    check({ subscribe: 3, unsubscribe: 1, parent: 3 });
+
+    d.dirty("rawr");
+    check({ subscribe: 3, unsubscribe: 2, parent: 3 });
+    assert.strictEqual(parent("blep"), 3);
+    check({ subscribe: 3, unsubscribe: 2, parent: 3 });
+    assert.strictEqual(parent("rawr"), 4);
+    check({ subscribe: 4, unsubscribe: 2, parent: 4 });
+    assert.strictEqual(parent("blep"), 3);
+    check({ subscribe: 4, unsubscribe: 2, parent: 4 });
+  });
 });