From a26a2cfaf725e6f6f62ea83c1a8aac0179ee251e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?V=C3=ADctor=20Oliva?= <olivarra1@gmail.com>
Date: Tue, 16 Jan 2024 14:53:45 +0100
Subject: [PATCH] wip

---
 packages/context-state/src/subtree.test.ts |  4 ++--
 packages/context-state/src/subtree.ts      | 18 +++++++++++++++---
 packages/utils/src/suspend.ts              |  1 +
 3 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/packages/context-state/src/subtree.test.ts b/packages/context-state/src/subtree.test.ts
index 4715002..62b847f 100644
--- a/packages/context-state/src/subtree.test.ts
+++ b/packages/context-state/src/subtree.test.ts
@@ -12,7 +12,7 @@ describe("subtree", () => {
 
     const otherRoot = createRoot().withTypes<{ value: string }>()
 
-    subtree(subnode, (ctx) => otherRoot.run(null, { value: ctx(subnode) }))
+    subtree(subnode, otherRoot, (ctx) => [null, { value: ctx(subnode) }])
 
     mainRoot.run()
 
@@ -31,7 +31,7 @@ describe("subtree", () => {
     const otherRootCopy = substate(subnode, () => activeObs$(otherRoot))
 
     const otherRoot = createRoot().withTypes<{ value: string }>()
-    subtree(subnode, (ctx) => otherRoot.run(null, { value: ctx(subnode) }))
+    subtree(subnode, otherRoot, (ctx) => [null, { value: ctx(subnode) }])
 
     mainRoot.run()
 
diff --git a/packages/context-state/src/subtree.ts b/packages/context-state/src/subtree.ts
index c3496db..fe73755 100644
--- a/packages/context-state/src/subtree.ts
+++ b/packages/context-state/src/subtree.ts
@@ -7,17 +7,29 @@ import {
   StateNode,
   isSignal,
 } from "./types"
+import { RootNode } from "./create-root"
 
-export function subtree<K extends KeysBaseType>(
+export function subtree<K extends KeysBaseType, CtxValue, KeyValue>(
   parent: StateNode<unknown, K>,
-  createInstance: (
+  node: RootNode<CtxValue, string, KeyValue>,
+  instanceSelector: (
     ctxValue: GetValueFn,
     ctxObservable: GetObservableFn<K>,
     key: K,
-  ) => () => void,
+  ) => [KeyValue] | [KeyValue, CtxValue],
 ): void {
   const internalParent = getInternals(parent)
 
+  const createInstance = (
+    ctxValue: GetValueFn,
+    ctxObservable: GetObservableFn<K>,
+    key: K,
+  ) => {
+    const [keyValue, context] = instanceSelector(ctxValue, ctxObservable, key)
+
+    return node.run(keyValue, context!)
+  }
+
   const teardowns = new NestedMap<K[keyof K], () => void>()
   const nestedMapKey = (key: K) => internalParent.keysOrder.map((k) => key[k])
   const teardown = (key: K) => teardowns.get(nestedMapKey(key))?.()
diff --git a/packages/utils/src/suspend.ts b/packages/utils/src/suspend.ts
index 1ba0077..832a7c4 100644
--- a/packages/utils/src/suspend.ts
+++ b/packages/utils/src/suspend.ts
@@ -2,6 +2,7 @@ import { ObservableInput, from, Observable } from "rxjs"
 import { SUSPENSE } from "@react-rxjs/core"
 import { defaultStart } from "./internal-utils"
 
+//
 /**
  * A RxJS creation operator that prepends a SUSPENSE on the source observable.
  *