From c0c8027e3418f916f6ea7fc83dcddd7a14c7a61e Mon Sep 17 00:00:00 2001
From: John Archer <jahglow@gmail.com>
Date: Thu, 7 Jul 2022 11:13:39 +0300
Subject: [PATCH] fixes #1000

Bug: atomWithMachine typescript error after updating xstate to version 4.29.0
---
 package.json                  |   2 +-
 src/xstate/atomWithMachine.ts | 120 +++++++++++++++++++++++-----------
 yarn.lock                     |   8 +--
 3 files changed, 88 insertions(+), 42 deletions(-)

diff --git a/package.json b/package.json
index b3042fbedf..4aa6f41b6e 100644
--- a/package.json
+++ b/package.json
@@ -255,7 +255,7 @@
     "typescript": "^4.7.4",
     "valtio": "^1.6.1",
     "wonka": "^4.0.15",
-    "xstate": "4.28.1",
+    "xstate": "^4.32.1",
     "zustand": "^3.7.2"
   },
   "resolutions": {
diff --git a/src/xstate/atomWithMachine.ts b/src/xstate/atomWithMachine.ts
index 9c87ef1b61..02b7df1bbe 100644
--- a/src/xstate/atomWithMachine.ts
+++ b/src/xstate/atomWithMachine.ts
@@ -1,37 +1,68 @@
 import { interpret } from 'xstate'
 import type {
+  AnyInterpreter,
+  AnyStateMachine,
+  AreAllImplementationsAssumedToBeProvided,
   EventObject,
-  Interpreter,
+  InternalMachineOptions,
+  InterpreterFrom,
   InterpreterOptions,
-  MachineOptions,
-  State,
-  StateMachine,
-  Typestate,
+  Prop,
+  StateConfig,
+  StateFrom,
 } from 'xstate'
+import type { Atom, Getter, WritableAtom } from 'jotai'
 import { atom } from 'jotai'
-import type { Atom, Getter } from 'jotai'
+
+export interface MachineAtomOptions<TContext, TEvent extends EventObject> {
+  /**
+   * If provided, will be merged with machine's `context`.
+   */
+  context?: Partial<TContext>
+  /**
+   * The state to rehydrate the machine to. The machine will
+   * start at this state instead of its `initialState`.
+   */
+  state?: StateConfig<TContext, TEvent>
+}
+
+type Options<TMachine extends AnyStateMachine> =
+  AreAllImplementationsAssumedToBeProvided<
+    TMachine['__TResolvedTypesMeta']
+  > extends false
+    ? InterpreterOptions &
+        MachineAtomOptions<TMachine['__TContext'], TMachine['__TEvent']> &
+        InternalMachineOptions<
+          TMachine['__TContext'],
+          TMachine['__TEvent'],
+          TMachine['__TResolvedTypesMeta'],
+          true
+        >
+    : InterpreterOptions &
+        MachineAtomOptions<TMachine['__TContext'], TMachine['__TEvent']> &
+        InternalMachineOptions<
+          TMachine['__TContext'],
+          TMachine['__TEvent'],
+          TMachine['__TResolvedTypesMeta']
+        >
+
+type MaybeParam<T> = T extends (v: infer V) => unknown ? V : never
 
 export function atomWithMachine<
-  TContext,
-  TEvent extends EventObject,
-  TTypestate extends Typestate<TContext> = { value: any; context: TContext }
+  TMachine extends AnyStateMachine,
+  TInterpreter = InterpreterFrom<TMachine>
 >(
-  getMachine:
-    | StateMachine<TContext, any, TEvent, TTypestate>
-    | ((get: Getter) => StateMachine<TContext, any, TEvent, TTypestate>),
-  getOptions?:
-    | (Partial<InterpreterOptions> & Partial<MachineOptions<TContext, TEvent>>)
-    | ((
-        get: Getter
-      ) => Partial<InterpreterOptions> &
-        Partial<MachineOptions<TContext, TEvent>>)
-) {
-  type Machine = StateMachine<TContext, any, TEvent, TTypestate>
-  type Service = Interpreter<TContext, any, TEvent, TTypestate>
-  type MachineState = State<TContext, TEvent, any, TTypestate>
-  const cachedMachineAtom = atom<{ machine: Machine; service: Service } | null>(
-    null
-  )
+  getMachine: TMachine | ((get: Getter) => TMachine),
+  getOptions?: Options<TMachine> | ((get: Getter) => Options<TMachine>)
+): WritableAtom<
+  StateFrom<TMachine>,
+  MaybeParam<Prop<TInterpreter, 'send'>>,
+  void
+> {
+  const cachedMachineAtom = atom<{
+    machine: AnyStateMachine
+    service: AnyInterpreter
+  } | null>(null)
   const machineAtom = atom(
     (get) => {
       const cachedMachine = get(cachedMachineAtom)
@@ -45,30 +76,33 @@ export function atomWithMachine<
         }
         throw new Error('get not allowed after initialization')
       }
-      const machine =
-        typeof getMachine === 'function' ? getMachine(safeGet) : getMachine
-      const options =
-        typeof getOptions === 'function' ? getOptions(safeGet) : getOptions
+      const machine = isGetter(getMachine) ? getMachine(safeGet) : getMachine
+      const options = isGetter(getOptions) ? getOptions(safeGet) : getOptions
       initializing = false
       const {
         guards,
         actions,
-        activities,
         services,
         delays,
+        context,
         ...interpreterOptions
       } = options || {}
+
       const machineConfig = {
         ...(guards && { guards }),
         ...(actions && { actions }),
-        ...(activities && { activities }),
         ...(services && { services }),
         ...(delays && { delays }),
       }
+
       const machineWithConfig = machine.withConfig(
-        machineConfig,
-        machine.context
+        machineConfig as any,
+        () => ({
+          ...machine.context,
+          ...context,
+        })
       )
+
       const service = interpret(machineWithConfig, interpreterOptions)
       return { machine: machineWithConfig, service }
     },
@@ -76,16 +110,20 @@ export function atomWithMachine<
       set(cachedMachineAtom, get(machineAtom))
     }
   )
+
   machineAtom.onMount = (commit) => {
     commit()
   }
-  const cachedMachineStateAtom = atom<MachineState | null>(null)
+
+  const cachedMachineStateAtom = atom<StateFrom<TMachine> | null>(null)
+
   const machineStateAtom = atom(
     (get) =>
-      get(cachedMachineStateAtom) ?? get(machineAtom).machine.initialState,
+      get(cachedMachineStateAtom) ??
+      (get(machineAtom).machine.initialState as StateFrom<TMachine>),
     (get, set, registerCleanup: (cleanup: () => void) => void) => {
       const { service } = get(machineAtom)
-      service.onTransition((nextState) => {
+      service.onTransition((nextState: any) => {
         set(cachedMachineStateAtom, nextState)
       })
       service.start()
@@ -94,8 +132,10 @@ export function atomWithMachine<
       })
     }
   )
+
   machineStateAtom.onMount = (initialize) => {
     let unsub: (() => void) | undefined | false
+
     initialize((cleanup) => {
       if (unsub === false) {
         cleanup()
@@ -103,6 +143,7 @@ export function atomWithMachine<
         unsub = cleanup
       }
     })
+
     return () => {
       if (unsub) {
         unsub()
@@ -110,12 +151,17 @@ export function atomWithMachine<
       unsub = false
     }
   }
+
   const machineStateWithServiceAtom = atom(
     (get) => get(machineStateAtom),
-    (get, _set, event: Parameters<Service['send']>[0]) => {
+    (get, _set, event: Parameters<AnyInterpreter['send']>[0]) => {
       const { service } = get(machineAtom)
       service.send(event)
     }
   )
+
   return machineStateWithServiceAtom
 }
+
+const isGetter = <T>(v: T | ((get: Getter) => T)): v is (get: Getter) => T =>
+  typeof v === 'function'
diff --git a/yarn.lock b/yarn.lock
index 8517f2348a..3eb89e0520 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -5494,10 +5494,10 @@ xmlchars@^2.2.0:
   resolved "https://registry.yarnpkg.com/xmlchars/-/xmlchars-2.2.0.tgz#060fe1bcb7f9c76fe2a17db86a9bc3ab894210cb"
   integrity sha512-JZnDKK8B0RCDw84FNdDAIpZK+JuJw+s7Lz8nksI7SIuU3UXJJslUthsi+uWBUYOwPFwW7W7PRLRfUKpxjtjFCw==
 
-xstate@4.28.1:
-  version "4.28.1"
-  resolved "https://registry.yarnpkg.com/xstate/-/xstate-4.28.1.tgz#5cd6b5976aaf5943538462f149339cd67802f104"
-  integrity sha512-0xvaegeZNeHJAJvpjznyNr91qB1xy1PqeYMDyknh5S23TBPQJUHS81hk8W3UcvnB9uNs0YmOU2daoqb3WegzYQ==
+xstate@^4.32.1:
+  version "4.32.1"
+  resolved "https://registry.yarnpkg.com/xstate/-/xstate-4.32.1.tgz#1a09c808a66072938861a3b4acc5b38460244b70"
+  integrity sha512-QYUd+3GkXZ8i6qdixnOn28bL3EvA++LONYL/EMWwKlFSh/hiLndJ8YTnz77FDs+JUXcwU7NZJg7qoezoRHc4GQ==
 
 y18n@^5.0.5:
   version "5.0.8"