fix types + subinstance test case

packages/context-state/src/combineStates.ts

@@ -108,9 +108,11 @@ export const combineStates = <
   return result.public as any
 }
 
-type UnionToIntersection<U> = (U extends any ? (k: U) => void : never) extends (
-  k: infer I extends KeysBaseType,
-) => void
+type UnionToIntersection<U> = U extends never
+  ? never
+  : (U extends any ? (k: U) => void : never) extends (
+      k: infer I extends KeysBaseType,
+    ) => void
   ? I
   : never
 

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

@@ -2,13 +2,20 @@ import { of } from "rxjs"
 import { createStateNode } from "./internal"
 import { StateNode } from "./types"
 
-export type RootNodeKey<CtxValue, KeyName extends string> = KeyName extends ""
+export type RootNodeKey<KeyName extends string, KeyValue> = KeyName extends ""
   ? {}
-  : Record<KeyName, CtxValue>
+  : Record<KeyName, KeyValue>
 
 type TeardownFn = () => void
+export type RunFn<CtxValue, KeyName, KeyValue> = KeyName extends ""
+  ? CtxValue extends null
+    ? () => TeardownFn
+    : (key: unknown, ctxValue: CtxValue) => TeardownFn
+  : CtxValue extends null
+  ? (key: KeyValue) => TeardownFn
+  : (key: KeyValue, ctxValue: CtxValue) => TeardownFn
 export interface RootNode<CtxValue, KeyName extends string, KeyValue>
-  extends StateNode<CtxValue, RootNodeKey<KeyValue, KeyName>> {
+  extends StateNode<CtxValue, RootNodeKey<KeyName, KeyValue>> {
   run: KeyName extends ""
     ? never extends CtxValue
       ? () => TeardownFn
@@ -32,7 +39,7 @@ export function createRoot<CtxValue, KeyName extends string, KeyValue>(
 ): RootNode<CtxValue, KeyName, KeyValue> {
   const contextValues = new Map<KeyValue, CtxValue>()
   const internalNode = createStateNode<
-    RootNodeKey<KeyValue, KeyName>,
+    RootNodeKey<KeyName, KeyValue>,
     CtxValue
   >(
     keyName ? [keyName] : [],
@@ -58,7 +65,7 @@ export function createRoot<CtxValue, KeyName extends string, KeyValue>(
                 [keyName]: root,
               }
             : {}
-        ) as RootNodeKey<KeyValue, KeyName>
+        ) as RootNodeKey<KeyName, KeyValue>
 
         // TODO throw if instance already exists?
         const contextValueKey = root ?? ("" as KeyValue)

packages/context-state/src/index.tsx

@@ -1,4 +1,5 @@
 export * from "./create-root"
+export * from "./combineStates"
 export * from "./route-state"
 export * from "./subinstance"
 export * from "./substate"

packages/context-state/src/objectBased.ts

@@ -1,3 +1,4 @@
+import { subtree } from "./subtree"
 import {
   KeysAreCompatible,
   MapKeys,
@@ -5,7 +6,12 @@ import {
   combineStates,
   CombineStateKeys,
 } from "./combineStates"
-import { RootNodeKey, createRoot, RootNode as rootNode } from "./create-root"
+import {
+  RootNodeKey,
+  RunFn,
+  createRoot,
+  RootNode as rootNode,
+} from "./create-root"
 import { createSignal } from "./create-signal"
 import { getInternals, mapRecord, setInternals } from "./internal"
 import { routeState } from "./route-state"
@@ -57,17 +63,26 @@ class StateNode<T, K extends types.KeysBaseType>
   }
 }
 
-export class RootNode<K extends string, V> extends StateNode<
-  never,
-  RootNodeKey<K, V>
-> {
+export class RootNode<
+  CtxValue,
+  K extends string = "",
+  KeyValue = unknown,
+> extends StateNode<CtxValue, RootNodeKey<K, KeyValue>> {
   constructor(keyName?: K) {
-    super(keyName ? createRoot(keyName) : createRoot())
+    super(keyName ? createRoot(keyName) : (createRoot() as any))
   }
 
-  run(key?: V) {
-    ;(this.node as rootNode<V, K>).run(key!)
+  withTypes<NewCtxValue, NewKeyValue = KeyValue>() {
+    return this as unknown as RootNode<NewCtxValue, K, NewKeyValue>
   }
+
+  run: RunFn<CtxValue, K, KeyValue> = function (
+    this: RootNode<CtxValue, K, KeyValue>,
+    root?: KeyValue,
+    ctxValue?: CtxValue,
+  ) {
+    return (this.node as rootNode<CtxValue, K, KeyValue>).run(root!, ctxValue!)
+  } as any
 }
 
 type ResultingRoutes<O, T, K extends types.KeysBaseType> = {
@@ -142,7 +157,7 @@ export function combineStateNodes<
   states: KeysAreCompatible<MapKeys<States>> extends true ? States : never,
 ): StateNode<
   StringRecordNodeToStringRecord<States>,
-  CombineStateKeys<MapKeys<States>> & types.KeysBaseType
+  CombineStateKeys<MapKeys<States>>
 > {
   return new StateNode(combineStates(states))
 }

packages/context-state/src/react-bindings.tsx

@@ -8,6 +8,7 @@ import React, {
   useState,
   useSyncExternalStore,
 } from "react"
+import { Subscription } from "rxjs"
 import { Instance, StatePromise, getInternals } from "./internal"
 import { KeysBaseType, StateNode } from "./types"
 
@@ -77,18 +78,21 @@ export const useStateNode = <O, K extends KeysBaseType>(
   if (ref.source$ !== instance) {
     ref.source$ = instance
     ref.args[0] = (next: () => void) => {
-      let subscription = instance.getState$().subscribe({
-        next,
-        error: (e) => {
-          setError(() => {
-            throw e
-          })
-        },
-        complete() {
-          next()
-          subscription.add(ref.args[0](next))
-        },
-      })
+      const subscription = new Subscription()
+      subscription.add(
+        instance.getState$().subscribe({
+          next,
+          error: (e) => {
+            setError(() => {
+              throw e
+            })
+          },
+          complete() {
+            next()
+            subscription.add(ref.args[0](next))
+          },
+        }),
+      )
       return () => {
         subscription.unsubscribe()
       }

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

@@ -1,4 +1,13 @@
-import { NEVER, Subject, filter, map, of, startWith } from "rxjs"
+import {
+  EMPTY,
+  NEVER,
+  Subject,
+  filter,
+  map,
+  of,
+  startWith,
+  switchMap,
+} from "rxjs"
 import { describe, expect, it, vi } from "vitest"
 import { InstanceUpdate, createRoot, subinstance, substate } from "./"
 
@@ -160,6 +169,71 @@ describe("subinstance", () => {
       expect(() => instances.getValue({ keyName: "b" })).toThrow()
       expect(() => [...(keys.getValue() as Set<string>)]).toThrow()
     })
+
+    it("can access instances as soon as they are announced", () => {
+      const root = createRoot()
+      const instance$ = new Subject<InstanceUpdate<string>>()
+      const [instanceNode, keys] = subinstance(
+        root,
+        "keyName",
+        () => instance$,
+        (id) => of(id),
+      )
+      root.run()
+
+      let error = null
+      let lastActive = null
+      keys
+        .getState$()
+        .pipe(
+          switchMap((v) =>
+            v.size
+              ? instanceNode.getState$({
+                  keyName: Array.from(v.values()).at(-1)!,
+                })
+              : EMPTY,
+          ),
+        )
+        .subscribe({
+          next: (v) => (lastActive = v),
+          error: (e) => (error = e),
+        })
+
+      instance$.next({
+        add: ["a"],
+      })
+      expect(lastActive).toEqual("a")
+      instance$.next({
+        add: ["b"],
+      })
+      expect(lastActive).toEqual("b")
+      expect(error).toBe(null)
+    })
+
+    it("continues working after the parent changes value", () => {
+      const root = createRoot()
+      const subnode$ = new Subject<string>()
+      const subnode = substate(root, () => subnode$.pipe(startWith("a")))
+      const [instanceNode, keys] = subinstance(
+        subnode,
+        "keyName",
+        (ctx) =>
+          ctx(subnode) === "a"
+            ? EMPTY
+            : of({
+                add: ["a", "b"],
+              }),
+        (id) => of(id),
+      )
+      root.run()
+
+      expect(Array.from(keys.getValue() as Set<string>)).toEqual([])
+      subnode$.next("b")
+      expect(Array.from(keys.getValue() as Set<string>)).toEqual(["a", "b"])
+
+      expect(instanceNode.getValue({ keyName: "a" })).toEqual("a")
+      expect(instanceNode.getValue({ keyName: "b" })).toEqual("b")
+    })
   })
 
   describe("key selector", () => {

packages/context-state/src/subinstance.ts

@@ -113,6 +113,10 @@ export function subinstance<K extends KeysBaseType, KN extends string, KV, R>(
     },
     onActive() {},
     onReset() {},
+    onAfterChange(key, storage) {
+      storage.value.unsubscribe()
+      storage.setValue(watchParentInstance(key))
+    },
     onRemoved(key, storage) {
       storage.value.unsubscribe()
       const orderedKey = parentInternals.keysOrder.map((k) => key[k])

packages/context-state/src/test-utils/finalizationRegistry.ts

@@ -1,4 +1,5 @@
 import "expose-gc"
+import { expect } from "vitest"
 
 export function testFinalizationRegistry() {
   const promises = new Map<