From 6300214de0cb1b210f7967d3f5a877d27e61c8cc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?V=C3=ADctor=20Oliva?= <olivarra1@gmail.com>
Date: Wed, 19 Oct 2022 15:16:36 +0200
Subject: [PATCH] feat: add compatible keys restriction to combineStates

---
 .../context-state/src/combineStates.test.ts   | 10 +++--
 packages/context-state/src/combineStates.ts   | 40 ++++++++++++++++++-
 packages/context-state/src/route-state.ts     |  8 ++--
 3 files changed, 50 insertions(+), 8 deletions(-)

diff --git a/packages/context-state/src/combineStates.test.ts b/packages/context-state/src/combineStates.test.ts
index e12a7c9..3edd287 100644
--- a/packages/context-state/src/combineStates.test.ts
+++ b/packages/context-state/src/combineStates.test.ts
@@ -2,7 +2,7 @@ import { createRoot } from "./create-root"
 import { BehaviorSubject, of, Subject } from "rxjs"
 import { substate } from "./substate"
 import { combineStates } from "./combineStates"
-import { StateNode } from "./types"
+import { StateNode, StringRecord } from "./types"
 import { routeState } from "./route-state"
 
 describe("combineStates", () => {
@@ -43,7 +43,9 @@ describe("combineStates", () => {
   })
 
   it("only activates if all branches are active", () => {
-    function createActivableNode(root: StateNode<any, any>) {
+    function createActivableNode<K extends StringRecord<any>>(
+      root: StateNode<any, K>,
+    ) {
       const ctxSource = new BehaviorSubject(false)
       const ctxNode = substate(root, () => ctxSource)
       const [, { node }] = routeState(
@@ -84,7 +86,9 @@ describe("combineStates", () => {
   })
 
   it("doesn't emit a value until all branches have one", async () => {
-    function createSettableNode(root: StateNode<any, any>) {
+    function createSettableNode<K extends StringRecord<any>>(
+      root: StateNode<never, K>,
+    ) {
       const source = new Subject()
       const node = substate(root, () => source)
       return [node, (value: any) => source.next(value)] as const
diff --git a/packages/context-state/src/combineStates.ts b/packages/context-state/src/combineStates.ts
index c74877d..4b8c618 100644
--- a/packages/context-state/src/combineStates.ts
+++ b/packages/context-state/src/combineStates.ts
@@ -28,7 +28,7 @@ interface CombinedStateInstance {
 }
 
 export const combineStates = <States extends StringRecord<StateNode<any, any>>>(
-  states: States,
+  states: KeysAreCompatible<MapKeys<States>> extends true ? States : never,
 ): StringRecordNodeToNodeStringRecord<States> => {
   const instances = new NestedMap<any[], CombinedStateInstance>()
   const nKeys = Object.keys(states).length
@@ -86,3 +86,41 @@ export const combineStates = <States extends StringRecord<StateNode<any, any>>>(
 
   return result.public as StringRecordNodeToNodeStringRecord<States>
 }
+
+type UnionToIntersection<U> = (U extends any ? (k: U) => void : never) extends (
+  k: infer I,
+) => void
+  ? I
+  : never
+
+/**
+ * Converts Record<string, State<any, K>> to Record<string, K>
+ */
+type MapKeys<States> = {
+  [K in keyof States]: States[K] extends StateNode<any, infer K> ? K : never
+}
+
+/**
+ * For each of the keys, check if they are compatible with the intersection
+ */
+type IndividualIsCompatible<KeysRecord, KeysIntersection> = {
+  [K in keyof KeysRecord]: KeysRecord[K] extends KeysIntersection ? true : false
+}
+
+/**
+ * It will be compatible if one of the individual ones returns true.
+ * If all of them are false, true extends false => false
+ * if one of them is true, true extends boolean => true
+ */
+type IsCompatible<KeysRecord, KeysIntersection> =
+  true extends IndividualIsCompatible<
+    KeysRecord,
+    KeysIntersection
+  >[keyof KeysRecord]
+    ? true
+    : false
+
+type KeysAreCompatible<KeysRecord> = IsCompatible<
+  KeysRecord,
+  UnionToIntersection<KeysRecord[keyof KeysRecord]>
+>
diff --git a/packages/context-state/src/route-state.ts b/packages/context-state/src/route-state.ts
index 6adb2f7..17faa4b 100644
--- a/packages/context-state/src/route-state.ts
+++ b/packages/context-state/src/route-state.ts
@@ -24,10 +24,10 @@ export const routeState = <
   K extends StringRecord<any>,
   O extends StringRecord<((value: T) => any) | null>,
   OT extends {
-    [K in keyof O]: null extends O[K]
-      ? StateNode<T, any>
-      : O[K] extends (value: T) => infer V
-      ? StateNode<V, any>
+    [KOT in keyof O]: null extends O[KOT]
+      ? StateNode<T, K>
+      : O[KOT] extends (value: T) => infer V
+      ? StateNode<V, K>
       : unknown
   },
 >(