From 030c578f1076be42cfed53d6bd1a73b4b7695f34 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?V=C3=ADctor=20Oliva?= <olivarra1@gmail.com>
Date: Tue, 11 Jul 2023 20:10:55 +0200
Subject: [PATCH] wip

---
 packages/context-state/src/combineStates.ts   | 17 +++++++---
 .../src/internal/state-instance.ts            | 32 +++++++++++++------
 packages/context-state/src/objectBased.ts     |  6 +++-
 3 files changed, 40 insertions(+), 15 deletions(-)

diff --git a/packages/context-state/src/combineStates.ts b/packages/context-state/src/combineStates.ts
index d763f7a..98384d1 100644
--- a/packages/context-state/src/combineStates.ts
+++ b/packages/context-state/src/combineStates.ts
@@ -7,7 +7,7 @@ import {
   trackParentChanges,
 } from "./internal"
 import { NestedMap, Wildcard } from "./internal/nested-map"
-import { StateNode } from "./types"
+import { KeysBaseType, StateNode } from "./types"
 
 export type StringRecordNodeToStringRecord<
   States extends Record<string, StateNode<any, any>>,
@@ -17,7 +17,10 @@ export type StringRecordNodeToStringRecord<
 
 type StringRecordNodeToNodeStringRecord<
   States extends Record<string, StateNode<any, any>>,
-> = StateNode<StringRecordNodeToStringRecord<States>, any>
+> = StateNode<
+  StringRecordNodeToStringRecord<States>,
+  CombineStateKeys<MapKeys<States>>
+>
 
 export const combineStates = <
   States extends Record<string, StateNode<any, any>>,
@@ -102,11 +105,11 @@ export const combineStates = <
   })
   addInstances()
 
-  return result.public as StringRecordNodeToNodeStringRecord<States>
+  return result.public as any
 }
 
 type UnionToIntersection<U> = (U extends any ? (k: U) => void : never) extends (
-  k: infer I,
+  k: infer I extends KeysBaseType,
 ) => void
   ? I
   : never
@@ -138,7 +141,11 @@ type IsCompatible<KeysRecord, KeysIntersection> =
     ? true
     : false
 
+export type CombineStateKeys<KeysRecord> = UnionToIntersection<
+  KeysRecord[keyof KeysRecord]
+>
+
 export type KeysAreCompatible<KeysRecord> = IsCompatible<
   KeysRecord,
-  UnionToIntersection<KeysRecord[keyof KeysRecord]>
+  CombineStateKeys<KeysRecord>
 >
diff --git a/packages/context-state/src/internal/state-instance.ts b/packages/context-state/src/internal/state-instance.ts
index 0ac5229..f4ec610 100644
--- a/packages/context-state/src/internal/state-instance.ts
+++ b/packages/context-state/src/internal/state-instance.ts
@@ -1,4 +1,11 @@
-import { Observable, Subject, Subscription, filter, startWith } from "rxjs"
+import {
+  Observable,
+  Subject,
+  Subscription,
+  filter,
+  startWith,
+  defer,
+} from "rxjs"
 import type { KeysBaseType } from "../types"
 import { EMPTY_VALUE } from "./empty-value"
 import { StatePromise, createDeferredPromise } from "./promises"
@@ -21,19 +28,26 @@ export function createInstance<T, K extends KeysBaseType>(
   let currentValue: T | EMPTY_VALUE = EMPTY_VALUE
   let previousContextValue: T | EMPTY_VALUE = currentValue
   let subject = new Subject<T | EMPTY_VALUE>()
-  let state$ = subject.pipe(
-    startWith(currentValue),
-    filter((v) => v !== EMPTY_VALUE),
-  ) as Observable<T>
+  // Needs to be deferred for the `startWith(currentValue)`
+  let state$ = defer(
+    () =>
+      subject.pipe(
+        startWith(currentValue),
+        filter((v) => v !== EMPTY_VALUE),
+      ) as Observable<T>,
+  )
 
   const resetSubject = () => {
     previousContextValue = EMPTY_VALUE
     const oldSubject = subject
     subject = new Subject<T | EMPTY_VALUE>()
-    state$ = subject.pipe(
-      startWith(currentValue),
-      filter((v) => v !== EMPTY_VALUE),
-    ) as Observable<T>
+    state$ = defer(
+      () =>
+        subject.pipe(
+          startWith(currentValue),
+          filter((v) => v !== EMPTY_VALUE),
+        ) as Observable<T>,
+    )
     oldSubject.complete()
   }
 
diff --git a/packages/context-state/src/objectBased.ts b/packages/context-state/src/objectBased.ts
index 03106bc..5bdbbdd 100644
--- a/packages/context-state/src/objectBased.ts
+++ b/packages/context-state/src/objectBased.ts
@@ -3,6 +3,7 @@ import {
   MapKeys,
   StringRecordNodeToStringRecord,
   combineStates,
+  CombineStateKeys,
 } from "./combineStates"
 import { RootNodeKey, createRoot, RootNode as rootNode } from "./create-root"
 import { createSignal } from "./create-signal"
@@ -139,6 +140,9 @@ export function combineStateNodes<
   States extends Record<string, types.StateNode<any, any>>,
 >(
   states: KeysAreCompatible<MapKeys<States>> extends true ? States : never,
-): StateNode<StringRecordNodeToStringRecord<States>, any> {
+): StateNode<
+  StringRecordNodeToStringRecord<States>,
+  CombineStateKeys<MapKeys<States>> & types.KeysBaseType
+> {
   return new StateNode(combineStates(states))
 }