move common subscription logic into a helper function

packages/context-state2/src/combineStates.ts

@@ -4,6 +4,7 @@ import {
   getInternals,
   mapRecord,
   recordEntries,
+  trackParentChanges,
 } from "./internal"
 import { NestedMap, Wildcard } from "./internal/nested-map"
 import { StateNode } from "./types"
@@ -69,36 +70,34 @@ export const combineStates = <
     }
   }
 
-  recordEntries(internalStates).forEach(([key, node]) => {
-    for (let instance of node.getInstances()) {
-      activeInstances[key].set(
-        node.keysOrder.map((k) => instance.key[k]),
-        true,
-      )
-    }
-
-    node.instanceChange$.subscribe((change) => {
-      if (change.type === "added") {
-        activeInstances[key].set(
-          node.keysOrder.map((k) => change.key[k]),
+  recordEntries(states).forEach(([nodeKey, node]) => {
+    trackParentChanges(node, {
+      onAdded(key, isInitial) {
+        const nodeKeyOrder = getInternals(node).keysOrder
+        activeInstances[nodeKey].set(
+          nodeKeyOrder.map((k) => key[k]),
           true,
         )
-        addInstances(
-          keysOrder.map((k) =>
-            node.keysOrder.includes(k) ? change.key[k] : Wildcard,
-          ),
-        )
-      } else if (change.type === "removed") {
-        activeInstances[key].set(
-          node.keysOrder.map((k) => change.key[k]),
+        if (!isInitial) {
+          addInstances(
+            keysOrder.map((k) =>
+              nodeKeyOrder.includes(k) ? key[k] : Wildcard,
+            ),
+          )
+        }
+      },
+      onActive() {},
+      onReset() {},
+      onRemoved(key) {
+        const nodeKeyOrder = getInternals(node).keysOrder
+        activeInstances[nodeKey].set(
+          nodeKeyOrder.map((k) => key[k]),
           true,
         )
         removeInstances(
-          keysOrder.map((k) =>
-            node.keysOrder.includes(k) ? change.key[k] : Wildcard,
-          ),
+          keysOrder.map((k) => (nodeKeyOrder.includes(k) ? key[k] : Wildcard)),
         )
-      }
+      },
     })
   })
   addInstances()

packages/context-state2/src/internal/index.ts

@@ -6,3 +6,4 @@ export * from "./nested-map"
 export * from "./internals"
 export * from "./state-node"
 export * from "./state-instance"
+export * from "./parent-changes"

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

@@ -1,7 +1,7 @@
 export const Wildcard = Symbol("Wildcard")
 export type Wildcard = typeof Wildcard
 
-export class NestedMap<K, V extends Object> {
+export class NestedMap<K, V> {
   private root: Map<K, any>
   private rootValue?: V
   constructor() {

packages/context-state2/src/internal/parent-changes.ts

@@ -0,0 +1,56 @@
+import { Subscription } from "rxjs"
+import { StateNode } from "../types"
+import { getInternals } from "./internals"
+import { NestedMap } from "./nested-map"
+
+type Storage<T> = {
+  readonly value: T
+  setValue: (value: T) => void
+  remove: () => T
+}
+export function trackParentChanges<T, K extends Record<string, any>, R>(
+  parent: StateNode<T, K>,
+  tracker: {
+    onAdded: (key: K, isInitial: boolean) => R
+    onActive: (key: K, storage: Storage<R>) => void
+    onReset: (key: K, storage: Storage<R>) => void
+    onRemoved: (key: K, storage: Storage<R>) => void
+  },
+): Subscription {
+  const internalParent = getInternals(parent)
+
+  const storage = new NestedMap<K[keyof K], R>()
+  const getStorage = (key: K): Storage<R> => {
+    const orderedKey = internalParent.keysOrder.map((k) => key[k])
+    const value = storage.get(orderedKey)!
+    return {
+      value,
+      setValue: (value) => storage.set(orderedKey, value),
+      remove: () => {
+        storage.delete(orderedKey)
+        return value
+      },
+    }
+  }
+
+  for (let instance of internalParent.getInstances()) {
+    getStorage(instance.key).setValue(tracker.onAdded(instance.key, true))
+  }
+  for (let instance of internalParent.getInstances()) {
+    tracker.onActive(instance.key, getStorage(instance.key))
+  }
+
+  return internalParent.instanceChange$.subscribe((change) => {
+    if (change.type === "added") {
+      getStorage(change.key).setValue(tracker.onAdded(change.key, false))
+    } else if (change.type === "ready") {
+      tracker.onActive(change.key, getStorage(change.key))
+    } else if (change.type === "reset") {
+      tracker.onReset(change.key, getStorage(change.key))
+    } else if (change.type === "removed") {
+      const storage = getStorage(change.key)
+      tracker.onRemoved(change.key, storage)
+      storage.remove()
+    }
+  })
+}

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

@@ -1,5 +1,10 @@
-import { Subscription, distinctUntilChanged, map, of } from "rxjs"
-import { NestedMap, createStateNode, getInternals, mapRecord } from "./internal"
+import { distinctUntilChanged, map, of } from "rxjs"
+import {
+  createStateNode,
+  getInternals,
+  mapRecord,
+  trackParentChanges,
+} from "./internal"
 import { GetValueFn, KeysBaseType, StateNode } from "./types"
 
 export class InvalidRouteError extends Error {
@@ -45,11 +50,9 @@ export const routeState = <
     )
   })
 
-  const subscriptions = new NestedMap<keyof K, Subscription>()
-
   const watchInstanceRoutes = (instanceKey: K) => {
     let previousNode: any = null
-    const sub = keyState
+    return keyState
       .getInstance(instanceKey)
       .getState$()
       .subscribe({
@@ -63,29 +66,23 @@ export const routeState = <
           previousNode = node
         },
       })
-    const key = internalParent.keysOrder.map((k) => instanceKey[k])
-    subscriptions.set(key, sub)
   }
   const removeInstanceRoutes = (instanceKey: K) => {
-    const key = internalParent.keysOrder.map((k) => instanceKey[k])
-    const sub = subscriptions.get(key)
-    subscriptions.delete(key)
     Object.values(routedState).forEach((node) => {
       node.removeInstance(instanceKey)
     })
-    sub?.unsubscribe()
   }
 
-  for (let instance of internalParent.getInstances()) {
-    watchInstanceRoutes(instance.key)
-  }
-
-  internalParent.instanceChange$.subscribe((change) => {
-    if (change.type === "added") {
-      watchInstanceRoutes(change.key)
-    } else if (change.type === "removed") {
-      removeInstanceRoutes(change.key)
-    }
+  trackParentChanges(parent, {
+    onAdded(key) {
+      return watchInstanceRoutes(key)
+    },
+    onActive() {},
+    onReset() {},
+    onRemoved(key, storage) {
+      removeInstanceRoutes(key)
+      storage.value.unsubscribe()
+    },
   })
 
   return [keyState.public, mapRecord(routedState, (v) => v.public) as OT]
@@ -111,29 +108,19 @@ const createKeyState = <T, K extends Record<string, any>>(
       ),
   )
 
-  const addInstance = (instanceKey: K) => {
-    // TODO duplicate ?
-    keyNode.addInstance(instanceKey)
-  }
-  const removeInstance = (instanceKey: K) => {
-    keyNode.removeInstance(instanceKey)
-  }
-
-  for (let instance of internalParent.getInstances()) {
-    addInstance(instance.key)
-  }
-  for (let instance of internalParent.getInstances()) {
-    keyNode.activateInstance(instance.key)
-  }
-
-  internalParent.instanceChange$.subscribe((change) => {
-    if (change.type === "added") {
-      addInstance(change.key)
-    } else if (change.type === "ready") {
-      keyNode.activateInstance(change.key)
-    } else if (change.type === "removed") {
-      removeInstance(change.key)
-    }
+  trackParentChanges(parent, {
+    onAdded(key) {
+      keyNode.addInstance(key)
+    },
+    onActive(key) {
+      keyNode.activateInstance(key)
+    },
+    onReset(key) {
+      keyNode.resetInstance(key)
+    },
+    onRemoved(key) {
+      keyNode.removeInstance(key)
+    },
   })
 
   return keyNode

packages/context-state2/src/subinstance.ts

@@ -1,5 +1,5 @@
 import { Observable, map, scan, startWith } from "rxjs"
-import { NestedMap, createStateNode, getInternals } from "./internal"
+import { createStateNode, getInternals, trackParentChanges } from "./internal"
 import { substate } from "./substate"
 import {
   CtxFn,
@@ -74,9 +74,8 @@ export function subinstance<K extends KeysBaseType, KN extends string, KV, R>(
       ),
   )
 
-  const parentInstanceWatches = new NestedMap<K[keyof K], () => void>()
   function watchParentInstance(key: K) {
-    const sub = instanceKeys
+    return instanceKeys
       .getState$(key)
       .pipe(map((v, i) => [v, i] as const))
       .subscribe(([v, i]) => {
@@ -111,30 +110,17 @@ export function subinstance<K extends KeysBaseType, KN extends string, KV, R>(
           }
         }
       })
-
-    parentInstanceWatches.set(
-      parentInternals.keysOrder.map((k) => key[k]),
-      () => sub.unsubscribe(),
-    )
-  }
-  function stopWatchParentInstance(key: K) {
-    const orderedKey = parentInternals.keysOrder.map((k) => key[k])
-    const teardown = parentInstanceWatches.get(orderedKey)
-    parentInstanceWatches.delete(orderedKey)
-    teardown?.()
   }
 
-  // TODO this pattern is common on all operators... maybe it can be abstracted away? watch parent instances, do something on them, tear down their subscriptions.
-  for (let instance of parentInternals.getInstances()) {
-    watchParentInstance(instance.key)
-  }
-
-  parentInternals.instanceChange$.subscribe((change) => {
-    if (change.type === "added") {
-      watchParentInstance(change.key)
-    } else if (change.type === "removed") {
-      stopWatchParentInstance(change.key)
-    }
+  trackParentChanges(parent, {
+    onAdded(key) {
+      return watchParentInstance(key)
+    },
+    onActive() {},
+    onReset() {},
+    onRemoved(_, storage) {
+      storage.value.unsubscribe()
+    },
   })
 
   return [result.public, instanceKeys]

packages/context-state2/src/substate.ts

@@ -1,5 +1,5 @@
-import { EMPTY, Subscription, defer, distinctUntilChanged, skip } from "rxjs"
-import { NestedMap, createStateNode, getInternals } from "./internal"
+import { EMPTY, defer, distinctUntilChanged, skip } from "rxjs"
+import { createStateNode, getInternals, trackParentChanges } from "./internal"
 import {
   isSignal,
   type CtxFn,
@@ -25,12 +25,10 @@ export const substate = <T, K extends KeysBaseType>(
       ),
   )
 
-  const subscriptions = new NestedMap<K[keyof K], Subscription>()
-
   const addInstance = (instanceKey: K) => {
     // TODO duplicate ?
     stateNode.addInstance(instanceKey)
-    const sub = defer(() => {
+    return defer(() => {
       try {
         return parent.getState$(instanceKey)
       } catch (ex) {
@@ -51,36 +49,22 @@ export const substate = <T, K extends KeysBaseType>(
           // ?
         },
       })
-    subscriptions.set(
-      internalParent.keysOrder.map((k) => instanceKey[k]),
-      sub,
-    )
-  }
-  const removeInstance = (instanceKey: K) => {
-    const key = internalParent.keysOrder.map((k) => instanceKey[k])
-    const sub = subscriptions.get(key)
-    subscriptions.delete(key)
-    sub?.unsubscribe()
-    stateNode.removeInstance(instanceKey)
   }
 
-  for (let instance of internalParent.getInstances()) {
-    addInstance(instance.key)
-  }
-  for (let instance of internalParent.getInstances()) {
-    stateNode.activateInstance(instance.key)
-  }
-
-  internalParent.instanceChange$.subscribe((change) => {
-    if (change.type === "added") {
-      addInstance(change.key)
-    } else if (change.type === "ready") {
-      stateNode.activateInstance(change.key)
-    } else if (change.type === "removed") {
-      removeInstance(change.key)
-    } else if (change.type === "reset") {
-      stateNode.resetInstance(change.key)
-    }
+  trackParentChanges(parent, {
+    onAdded(key) {
+      return addInstance(key)
+    },
+    onActive(key) {
+      stateNode.activateInstance(key)
+    },
+    onReset(key) {
+      stateNode.resetInstance(key)
+    },
+    onRemoved(key, storage) {
+      stateNode.removeInstance(key)
+      storage.value.unsubscribe()
+    },
   })
 
   return stateNode.public