@xstate.fsm: State value added to transition type (#4043)

* feat: state value added to transition type

* fix: types changed

* change `state.matches` signature in `@xstate/fsm` to mimick the one in `xstate`

* add runtime context to what was intended to be type-only test

* add default to the new type param in `Transition`

* Update .changeset/modern-spies-cry.md

---------

Co-authored-by: Сергей Краснов <krasnov.s@releaseband.com>
Co-authored-by: Mateusz Burzyński <mateuszburzynski@gmail.com>

.changeset/modern-spies-cry.md

@@ -0,0 +1,5 @@
+---
+'@xstate/fsm': patch
+---
+
+Pass around `TState['value']` type to `Transition` and `.initial` property of the machine configuration.

packages/xstate-fsm/src/index.ts

@@ -130,7 +130,7 @@ function handleActions<
 export function createMachine<
   TContext extends object,
   TEvent extends EventObject = EventObject,
-  TState extends Typestate<TContext> = { value: any; context: TContext }
+  TState extends Typestate<TContext> = Typestate<TContext>
 >(
   fsmConfig: StateMachine.Config<TContext, TEvent, TState>,
   implementations: {
@@ -182,18 +182,21 @@ export function createMachine<
 
       if (stateConfig.on) {
         const transitions: Array<
-          StateMachine.Transition<TContext, TEvent>
+          StateMachine.Transition<TContext, TEvent, TState['value']>
         > = toArray(stateConfig.on[eventObject.type]);
 
         for (const transition of transitions) {
           if (transition === undefined) {
             return createUnchangedState(value, context);
           }
 
-          const { target, actions = [], cond = () => true } =
-            typeof transition === 'string'
-              ? { target: transition }
-              : transition;
+          const {
+            target,
+            actions = [],
+            cond = () => true
+          } = typeof transition === 'string'
+            ? { target: transition }
+            : transition;
 
           const isTargetless = target === undefined;
 
@@ -209,11 +212,12 @@ export function createMachine<
           }
 
           if (cond(context, eventObject)) {
-            const allActions = (isTargetless
-              ? toArray(actions)
-              : ([] as any[])
-                  .concat(stateConfig.exit, actions, nextStateConfig.entry)
-                  .filter((a) => a)
+            const allActions = (
+              isTargetless
+                ? toArray(actions)
+                : ([] as any[])
+                    .concat(stateConfig.exit, actions, nextStateConfig.entry)
+                    .filter((a) => a)
             ).map<StateMachine.ActionObject<TContext, TEvent>>((action) =>
               toActionObject(action, (machine as any)._options.actions)
             );

packages/xstate-fsm/src/types.ts

@@ -86,10 +86,14 @@ export namespace StateMachine {
     assignment: Assigner<TContext, TEvent> | PropertyAssigner<TContext, TEvent>;
   }
 
-  export type Transition<TContext extends object, TEvent extends EventObject> =
-    | string
+  export type Transition<
+    TContext extends object,
+    TEvent extends EventObject,
+    TStateValue extends string = string
+  > =
+    | TStateValue
     | {
-        target?: string;
+        target?: TStateValue;
         actions?: SingleOrArray<Action<TContext, TEvent>>;
         cond?: (context: TContext, event: TEvent) => boolean;
       };
@@ -104,9 +108,15 @@ export namespace StateMachine {
     changed?: boolean | undefined;
     matches: <TSV extends TState['value']>(
       value: TSV
-    ) => this is TState extends { value: TSV }
-      ? TState & { value: TSV }
-      : never;
+    ) => this is State<
+      (TState extends any
+        ? { value: TSV; context: any } extends TState
+          ? TState
+          : never
+        : never)['context'],
+      TEvent,
+      TState
+    > & { value: TSV };
   }
 
   export type AnyMachine = StateMachine.Machine<any, any, any>;
@@ -118,16 +128,20 @@ export namespace StateMachine {
   export interface Config<
     TContext extends object,
     TEvent extends EventObject,
-    TState extends Typestate<TContext> = { value: any; context: TContext }
+    TState extends Typestate<TContext> = Typestate<TContext>
   > {
     id?: string;
-    initial: string;
+    initial: TState['value'];
     context?: TContext;
     states: {
       [key in TState['value']]: {
         on?: {
           [K in TEvent['type']]?: SingleOrArray<
-            Transition<TContext, TEvent extends { type: K } ? TEvent : never>
+            Transition<
+              TContext,
+              TEvent extends { type: K } ? TEvent : never,
+              TState['value']
+            >
           >;
         };
         exit?: SingleOrArray<Action<TContext, TEvent>>;
@@ -157,9 +171,7 @@ export namespace StateMachine {
     TState extends Typestate<TContext> = { value: any; context: TContext }
   > {
     send: (event: TEvent | TEvent['type']) => void;
-    subscribe: (
-      listener: StateListener<State<TContext, TEvent, TState>>
-    ) => {
+    subscribe: (listener: StateListener<State<TContext, TEvent, TState>>) => {
       unsubscribe: () => void;
     };
     start: (

packages/xstate-fsm/test/types.test.ts

@@ -18,4 +18,44 @@ describe('matches', () => {
       ((_accept: string) => {})(state.context.count);
     }
   });
+
+  it('should narrow context using typestates', () => {
+    type MyContext = { user: { name: string } | null; count: number };
+    type Typestates =
+      | {
+          value: 'idle';
+          context: { user: null; count: number };
+        }
+      | {
+          value: 'fetched';
+          context: { user: { name: string }; count: number };
+        };
+    const machine = createMachine<MyContext, { type: string }, Typestates>({
+      context: {
+        user: null,
+        count: 0
+      },
+      initial: 'idle',
+      states: { idle: {}, fetched: {} }
+    });
+    const state = machine.initialState;
+
+    if (state.matches('idle')) {
+      ((_accept: null) => {})(state.context.user);
+      // @ts-expect-error
+      ((_accept: { name: string }) => {})(state.context.user);
+
+      ((_accept: number) => {})(state.context.count);
+      // @ts-expect-error
+      ((_accept: string) => {})(state.context.count);
+    } else if (state.matches('fetched')) {
+      // @ts-expect-error
+      ((_accept: null) => {})(state.context.user);
+      ((_accept: { name: string }) => {})(state.context.user);
+
+      ((_accept: number) => {})(state.context.count);
+      // @ts-expect-error
+      ((_accept: string) => {})(state.context.count);
+    }
+  });
 });