context-state: factories, factories everywhere...

packages/context-state/src/combineStates.ts

@@ -1,3 +1,4 @@
+import { NestedMap } from "./internal/nested-map"
 import { of } from "rxjs"
 import {
   mapRecord,
@@ -17,29 +18,45 @@ type StringRecordNodeToNodeStringRecord<
 export const combineStates = <States extends StringRecord<StateNode<any>>>(
   states: States,
 ): StringRecordNodeToNodeStringRecord<States> => {
-  let inactiveStates = Object.keys(states).length
-  let emptyStates = 0
-  const activeStates = mapRecord(states, () => false)
-  const latestStates = mapRecord(states, () => null)
+  const instances = new NestedMap()
+  const nKeys = Object.keys(states).length
+  const _activeStates = mapRecord(states, () => false)
+  const _latestStates = mapRecord(states, () => null)
 
-  const [result, run] = detachedNode(() => of({ ...latestStates }))
+  const [result, run] = detachedNode((ctx) =>
+    of(mapRecord(states, (node) => ctx(node))),
+  )
 
   let latestValue: boolean | EMPTY_VALUE = false
   recordEntries(states).forEach(([key, node]) => {
-    children.get(node)!.add((isActive, value) => {
-      if (isActive !== activeStates[key]) {
-        inactiveStates += isActive ? -1 : +1
-        activeStates[key] = isActive
+    children.get(node)!.add((ctxKey, isActive, value) => {
+      let instance: any = instances.get(ctxKey)
+      if (!instance) {
+        instance = {
+          inactiveStates: nKeys,
+          activeStates: { ..._activeStates },
+          latestStates: { ..._latestStates },
+        }
+        instances.set(ctxKey, instance)
       }
 
-      if (value !== latestStates[key]) {
-        emptyStates +=
-          latestStates[key] === EMPTY_VALUE ? -1 : value === EMPTY_VALUE ? 1 : 0
-        latestStates[key] = value
+      if (isActive !== instance.activeStates[key]) {
+        instance.inactiveStates += isActive ? -1 : +1
+        instance.activeStates[key] = isActive
       }
 
-      latestValue = emptyStates === 0 ? !latestValue : EMPTY_VALUE
-      run(inactiveStates === 0, latestValue)
+      if (value !== instance.latestStates[key]) {
+        instance.emptyStates +=
+          instance.latestStates[key] === EMPTY_VALUE
+            ? -1
+            : value === EMPTY_VALUE
+            ? 1
+            : 0
+        instance.latestStates[key] = value
+      }
+
+      latestValue = instance.emptyStates === 0 ? !latestValue : EMPTY_VALUE
+      run(ctxKey, instance.inactiveStates === 0, instance.latestValue)
     })
   })
 

packages/context-state/src/create-root.ts

@@ -2,14 +2,16 @@ import { children } from "./internal"
 import { StateNode } from "./types"
 
 export interface RootNode extends StateNode<never> {
-  run(): () => void
+  run(rootKey?: any): () => void
 }
 
 export function createRoot(): RootNode {
-  const childRunners = new Set<(isActive: boolean, value: null) => void>()
-  const runChildren = (isActive: boolean) => {
+  const childRunners = new Set<
+    (key: any, isActive: boolean, value: null) => void
+  >()
+  const runChildren = (key: any, isActive: boolean) => {
     childRunners.forEach((cb) => {
-      cb(isActive, null)
+      cb(key, isActive, null)
     })
   }
 
@@ -20,11 +22,11 @@ export function createRoot(): RootNode {
     state$: () => {
       throw new Error("RootNode doesn't have value")
     },
-    run: () => {
+    run: (...rootKey) => {
       // Maybe more fancy with refcount, etc?
-      runChildren(true)
+      runChildren(rootKey, true)
       return () => {
-        runChildren(false)
+        runChildren(rootKey, false)
       }
     },
   }

packages/context-state/src/internal/children.ts

@@ -2,7 +2,7 @@ import { StateNode } from "../types"
 import { EMPTY_VALUE } from "./empty-value"
 
 export interface RunFn<P> {
-  (isActive: boolean, parentValue?: P | EMPTY_VALUE): void
+  (key: any, isActive: boolean, parentValue?: P | EMPTY_VALUE): void
 }
 
 export const children = new WeakMap<StateNode<any>, Set<RunFn<any>>>()

packages/context-state/src/internal/detached-node.ts

@@ -1,5 +1,5 @@
 import { Observable, ReplaySubject, Subscription } from "rxjs"
-import type { StateNode } from "../types"
+import type { StateNode, Ctx } from "../types"
 import {
   EMPTY_VALUE,
   inactiveContext,
@@ -10,36 +10,38 @@ import {
   children,
   RunFn,
 } from "./"
-
-export const ctx = <V>(node: StateNode<V>): V => {
-  const value = node.getValue()
-  if (value instanceof StatePromise) throw invalidContext()
-  return value
-}
-export type Ctx = typeof ctx
+import { NestedMap } from "./nested-map"
 
 export const detachedNode = <T, P>(
   getState$: (ctx: Ctx) => Observable<T>,
   equalityFn: (a: T, b: T) => boolean = Object.is,
 ): [StateNode<T>, RunFn<P>] => {
-  let subject: ReplaySubject<T> | null = null
-  let subscription: Subscription | null = null
-  let currentValue: EMPTY_VALUE | T = EMPTY_VALUE
-  let currentParentValue: EMPTY_VALUE | P = EMPTY_VALUE
-  let promise: DeferredPromise<T> | null = null
+  const instances = new NestedMap<
+    any,
+    {
+      subject: ReplaySubject<T>
+      subscription: Subscription | null
+      currentValue: EMPTY_VALUE | T
+      currentParentValue: EMPTY_VALUE
+      promise: DeferredPromise<T> | null
+    }
+  >()
 
   const result: StateNode<T> = {
-    getValue: () => {
-      if (!subject) throw inactiveContext()
+    getValue: (...key: any[]) => {
+      const instance = instances.get(key)
+      if (!instance) throw inactiveContext()
+      const { currentValue, promise } = instance
       if (currentValue !== EMPTY_VALUE) return currentValue
       if (promise) return promise.promise
-      promise = createDeferredPromise()
-      return promise.promise
+      instance.promise = createDeferredPromise()
+      return instance.promise.promise
     },
-    state$: () =>
+    state$: (...key: any[]) =>
       new Observable<T>((observer) => {
-        if (subject) return subject.subscribe(observer)
-        return observer.error(inactiveContext())
+        const instance = instances.get(key)
+        if (!instance) return observer.error(inactiveContext())
+        return instance.subject.subscribe(observer)
       }),
   }
 
@@ -52,89 +54,126 @@ export const detachedNode = <T, P>(
     })
   }
 
-  const run = (isActive: boolean, parentValue: any) => {
+  const run = (key: any[], isActive: boolean, parentValue: any) => {
+    let instance = instances.get(key)
     if (!isActive) {
-      const prevSubect = subject
-      const prevPromise = promise
-      subject = null
-      promise = null
+      if (!instance) return
+      instances.delete(key)
 
-      subscription?.unsubscribe()
-      subscription = null
+      instance.subscription?.unsubscribe()
 
-      currentValue = EMPTY_VALUE
-      currentParentValue = EMPTY_VALUE
-
-      runChildren(false, EMPTY_VALUE)
-      prevPromise?.rej(inactiveContext())
-      prevSubect?.complete()
+      runChildren(key, false)
+      instance.promise?.rej(inactiveContext())
+      instance.subject.complete()
       return
     }
 
     if (parentValue !== EMPTY_VALUE) {
       // an actual change of context
-      subscription?.unsubscribe()
-      currentValue = EMPTY_VALUE
-      currentParentValue = parentValue
-      subject = subject || new ReplaySubject<T>(1)
+      const hasPreviousValue = instance && instance.currentValue !== EMPTY_VALUE
+      if (!instance) {
+        instance = {
+          subject: new ReplaySubject<T>(1),
+          subscription: null,
+          currentValue: EMPTY_VALUE,
+          currentParentValue: EMPTY_VALUE,
+          promise: null,
+        }
+        instances.set(key, instance)
+      } else {
+        instance.subscription?.unsubscribe()
+        instance.currentValue = EMPTY_VALUE
+        instance.currentParentValue = parentValue
+      }
+      const actualInstance = instance
 
-      subscription = getState$(ctx).subscribe({
-        next(value) {
-          let prevValue = currentValue
-          currentValue = value
-          const prevPromise = promise
-          promise = null
-          if (prevValue === EMPTY_VALUE || !equalityFn(prevValue, value)) {
-            prevPromise?.res(value)
-            runChildren(true, value)
-            subject!.next(value)
-          }
-        },
-        error(err) {
-          const prevPromise = promise
-          const prevSubect = subject
+      const ctx = <V>(node: StateNode<V>): V => {
+        const value = (node as any).getValue(...key)
+        if (value instanceof StatePromise) throw invalidContext()
+        return value
+      }
 
-          subscription?.unsubscribe()
-          subscription = null
-          promise = null
-          subject = null
+      const onError = (err: any) => {
+        instances.delete(key)
+        const prevPromise = actualInstance.promise
+        const prevSubect = actualInstance.subject
 
-          currentValue = EMPTY_VALUE
-          currentParentValue = EMPTY_VALUE
+        actualInstance.subscription = null
+        actualInstance.promise = null
+        delete (actualInstance as any).subject
 
-          runChildren(false, EMPTY_VALUE)
-          prevPromise?.rej(err)
-          prevSubect?.error(err)
-        },
-        complete() {
-          subscription = null
-        },
-      })
+        actualInstance.currentValue = EMPTY_VALUE
+        actualInstance.currentParentValue = EMPTY_VALUE
 
-      if (subscription.closed) subscription = null
+        runChildren(key, false)
+        prevPromise?.rej(err)
+        prevSubect?.error(err)
+      }
 
-      if (currentValue === EMPTY_VALUE && subject) {
-        const prevSubect = subject
-        subject = new ReplaySubject<T>(1)
-        runChildren(true, EMPTY_VALUE)
-        prevSubect.complete()
+      let observable: Observable<any> | null = null
+      try {
+        observable = getState$(ctx)
+      } catch (e) {
+        onError(e)
+      }
+
+      actualInstance.subscription =
+        observable?.subscribe({
+          next(value) {
+            let prevValue = actualInstance.currentValue
+            actualInstance.currentValue = value
+            const prevPromise = actualInstance.promise
+            actualInstance.promise = null
+            if (prevValue === EMPTY_VALUE || !equalityFn(prevValue, value)) {
+              prevPromise?.res(value)
+              runChildren(key, true, value)
+              actualInstance.subject!.next(value)
+            }
+          },
+          error: onError,
+          complete() {
+            actualInstance.subscription = null
+          },
+        }) ?? null
+
+      if (actualInstance.subscription?.closed)
+        actualInstance.subscription = null
+
+      if (
+        actualInstance.currentValue === EMPTY_VALUE &&
+        actualInstance.subject
+      ) {
+        let prevSubect
+        if (hasPreviousValue) {
+          prevSubect = actualInstance.subject
+          actualInstance.subject = new ReplaySubject<T>(1)
+        }
+        runChildren(key, true, EMPTY_VALUE)
+        prevSubect?.complete()
       }
       return
     }
 
     // at this point parentValue is EMPTY_VALUE
-    if (currentParentValue === EMPTY_VALUE && subject) return
-
-    const prevSubect = subject
-    subject = new ReplaySubject<T>(1)
-
-    subscription?.unsubscribe()
-    subscription = null
-
-    currentValue = EMPTY_VALUE
-    currentParentValue = EMPTY_VALUE
-
-    runChildren(true, EMPTY_VALUE)
+    if (instance?.currentParentValue === EMPTY_VALUE) return
+    const prevSubect = instance?.subject
+    if (instance) {
+      instance.subject = new ReplaySubject<T>(1)
+      instance.subscription?.unsubscribe()
+      instance.subscription = null
+      instance.currentValue = EMPTY_VALUE
+      instance.currentParentValue = EMPTY_VALUE
+    } else {
+      instance = {
+        subject: new ReplaySubject<T>(1),
+        subscription: null,
+        currentValue: EMPTY_VALUE,
+        currentParentValue: EMPTY_VALUE,
+        promise: null,
+      }
+      instances.set(key, instance)
+    }
+    runChildren(key, true, EMPTY_VALUE)
     prevSubect?.complete()
   }
 

packages/context-state/src/internal/nested-map.ts

@@ -0,0 +1,56 @@
+export class NestedMap<K extends [], V extends Object> {
+  private root: Map<K, V>
+  private rootValue?: V
+  constructor() {
+    this.root = new Map()
+    this.rootValue = undefined
+  }
+
+  get(keys: K[]): V | undefined {
+    if (keys.length === 0) return this.rootValue
+    let current: any = this.root
+    for (let i = 0; i < keys.length; i++) {
+      current = current.get(keys[i])
+      if (!current) return undefined
+    }
+    return current
+  }
+
+  set(keys: K[], value: V): void {
+    if (keys.length === 0) {
+      this.rootValue = value
+      return
+    }
+    let current: Map<K, any> = this.root
+    let i
+    for (i = 0; i < keys.length - 1; i++) {
+      let nextCurrent = current.get(keys[i])
+      if (!nextCurrent) {
+        nextCurrent = new Map<K, any>()
+        current.set(keys[i], nextCurrent)
+      }
+      current = nextCurrent
+    }
+    current.set(keys[i], value)
+  }
+
+  delete(keys: K[]): void {
+    if (keys.length === 0) {
+      delete this.rootValue
+      return
+    }
+    const maps: Map<K, any>[] = [this.root]
+    let current: Map<K, any> = this.root
+
+    for (let i = 0; i < keys.length - 1; i++) {
+      maps.push((current = current.get(keys[i])))
+    }
+
+    let mapIdx = maps.length - 1
+    maps[mapIdx].delete(keys[mapIdx])
+
+    while (--mapIdx > -1 && maps[mapIdx].get(keys[mapIdx]).size === 0) {
+      maps[mapIdx].delete(keys[mapIdx])
+    }
+  }
+}

packages/context-state/src/route-state.ts

@@ -6,8 +6,8 @@ import {
   recordEntries,
 } from "./internal"
 import { of } from "rxjs"
-import { Ctx, substate } from "./substate"
-import { StateNode } from "./types"
+import { substate } from "./substate"
+import { StateNode, Ctx } from "./types"
 
 export const routeState = <
   T,
@@ -37,15 +37,19 @@ export const routeState = <
     recordEntries(routedState).map(([key, value]) => [key, value[1]]),
   )
 
-  const run = (isActive: boolean, value: keyof O | EMPTY_VALUE) => {
+  const run = (
+    ctxKey: any[],
+    isActive: boolean,
+    value: keyof O | EMPTY_VALUE,
+  ) => {
     if (!isActive || value === EMPTY_VALUE)
       runners.forEach((runner) => {
-        runner(false)
+        runner(ctxKey, false)
       })
 
     runners.forEach((runner, key) => {
-      if (key === value) runner(true, value)
-      else runner(false)
+      if (key === value) runner(ctxKey, true, value)
+      else runner(ctxKey, false)
     })
   }
 

packages/context-state/src/substate.test.ts

@@ -85,8 +85,7 @@ describe("subState", () => {
       )
 
       substate(root, (ctx) => of(ctx(branchB)))
-      // TODO Should the error happen on run? I think it should happen somewhere else where it can get captured consistently (see next test too, it's being swallowed)
-      expect(() => root.run()).toThrowError("Invalid Context")
+      expect(() => root.run()).not.toThrow()
     })
 
     it("becomes unactive after throws an error for an invalid accessed context", () => {

packages/context-state/src/substate.ts

@@ -1,18 +1,6 @@
 import { Observable } from "rxjs"
-import type { StateNode } from "./types"
-import {
-  invalidContext,
-  StatePromise,
-  children,
-  detachedNode,
-} from "./internal"
-
-export const ctx = <V>(node: StateNode<V>): V => {
-  const value = node.getValue()
-  if (value instanceof StatePromise) throw invalidContext()
-  return value
-}
-export type Ctx = typeof ctx
+import type { Ctx, StateNode } from "./types"
+import { children, detachedNode } from "./internal"
 
 export const substate = <T, P>(
   parent: StateNode<P>,

packages/context-state/src/types.ts

@@ -8,10 +8,12 @@ export declare type StringRecord<T> = {
 }
 
 export interface StateNode<T> {
-  getValue: () => T | StatePromise<T>
-  state$: () => Observable<T>
+  getValue: (...other: any[]) => T | StatePromise<T>
+  state$: (...other: any[]) => Observable<T>
 }
 
+export type Ctx = <V>(node: StateNode<V>) => V
+
 /*
 export type StateNodeFn<
   Key,