Tags in state

statelyai · Aug 29, 2023 · 175fc14 · 175fc14
1 parent 79da3e6
commit 175fc14
Show file tree

Hide file tree

Showing 8 changed files with 96 additions and 44 deletions.
diff --git a/packages/core/src/State.ts b/packages/core/src/State.ts
@@ -70,6 +70,7 @@ export class State<
   TContext extends MachineContext,
   TEvent extends EventObject,
   TActor extends ProvidedActor,
+  TTag extends string,
   TOutput,
   TResolvedTypesMeta = TypegenDisabled
 > {
@@ -106,13 +107,29 @@ export class State<
     TContext extends MachineContext,
     TEvent extends EventObject = EventObject
   >(
-    stateValue: State<TContext, TEvent, TODO, any, any> | StateValue,
+    stateValue:
+      | State<
+          TContext,
+          TEvent,
+          TODO,
+          any, // tags
+          any, // output
+          any // typege
+        >
+      | StateValue,
     context: TContext = {} as TContext,
     machine: AnyStateMachine
-  ): State<TContext, TEvent, TODO, any, any> {
+  ): State<
+    TContext,
+    TEvent,
+    TODO,
+    any, // tags
+    any, // output
+    any // typegen
+  > {
     if (stateValue instanceof State) {
       if (stateValue.context !== context) {
-        return new State<TContext, TEvent, TODO, any>(
+        return new State<TContext, TEvent, TODO, any, any, any>(
           {
             value: stateValue.value,
             context,
@@ -131,7 +148,7 @@ export class State<
       getStateNodes(machine.root, stateValue)
     );
 
-    return new State<TContext, TEvent, TODO, any>(
+    return new State<TContext, TEvent, TODO, any, any, any>(
       {
         value: stateValue,
         context,
@@ -212,7 +229,7 @@ export class State<
   public hasTag(
     tag: TResolvedTypesMeta extends TypegenEnabled
       ? Prop<Prop<TResolvedTypesMeta, 'resolved'>, 'tags'>
-      : string
+      : TTag
   ): boolean {
     return this.tags.has(tag as string);
   }

diff --git a/packages/core/src/StateMachine.ts b/packages/core/src/StateMachine.ts
@@ -71,10 +71,10 @@ export class StateMachine<
 > implements
     ActorLogic<
       TEvent,
-      State<TContext, TEvent, TActor, TOutput, TResolvedTypesMeta>,
-      State<TContext, TEvent, TActor, TOutput, TResolvedTypesMeta>,
+      State<TContext, TEvent, TActor, TTag, TOutput, TResolvedTypesMeta>,
+      State<TContext, TEvent, TActor, TTag, TOutput, TResolvedTypesMeta>,
       PersistedMachineState<
-        State<TContext, TEvent, TActor, TOutput, TResolvedTypesMeta>
+        State<TContext, TEvent, TActor, TTag, TOutput, TResolvedTypesMeta>
       >,
       TODO,
       TInput,
@@ -204,7 +204,7 @@ export class StateMachine<
    * @param state The state to resolve
    */
   public resolveState(
-    state: State<TContext, TEvent, TActor, TOutput, TResolvedTypesMeta>
+    state: State<TContext, TEvent, TActor, TTag, TOutput, TResolvedTypesMeta>
   ): typeof state {
     const configurationSet = getConfiguration(
       getStateNodes(this.root, state.value)
@@ -223,7 +223,7 @@ export class StateMachine<
     ...[context]: Equals<TContext, MachineContext> extends true
       ? []
       : [TContext]
-  ): State<TContext, TEvent, TActor, TOutput, TResolvedTypesMeta> {
+  ): State<TContext, TEvent, TActor, TTag, TOutput, TResolvedTypesMeta> {
     const resolvedStateValue = resolveStateValue(this.root, stateValue);
 
     return this.resolveState(State.from(resolvedStateValue, context, this));
@@ -237,10 +237,10 @@ export class StateMachine<
    * @param event The received event
    */
   public transition(
-    state: State<TContext, TEvent, TActor, TOutput, TResolvedTypesMeta>,
+    state: State<TContext, TEvent, TActor, TTag, TOutput, TResolvedTypesMeta>,
     event: TEvent,
     actorCtx: ActorContext<TEvent, typeof state>
-  ): State<TContext, TEvent, TActor, TOutput, TResolvedTypesMeta> {
+  ): State<TContext, TEvent, TActor, TTag, TOutput, TResolvedTypesMeta> {
     // TODO: handle error events in a better way
     if (
       isErrorEvent(event) &&
@@ -264,15 +264,15 @@ export class StateMachine<
    * @param event The received event
    */
   public microstep(
-    state: State<TContext, TEvent, TActor, TOutput, TResolvedTypesMeta>,
+    state: State<TContext, TEvent, TActor, TTag, TOutput, TResolvedTypesMeta>,
     event: TEvent,
     actorCtx: AnyActorContext
-  ): Array<State<TContext, TEvent, TActor, TOutput, TResolvedTypesMeta>> {
+  ): Array<State<TContext, TEvent, TActor, TTag, TOutput, TResolvedTypesMeta>> {
     return macrostep(state, event, actorCtx).microstates;
   }
 
   public getTransitionData(
-    state: State<TContext, TEvent, TActor, TOutput, TResolvedTypesMeta>,
+    state: State<TContext, TEvent, TActor, TTag, TOutput, TResolvedTypesMeta>,
     event: TEvent
   ): Array<TransitionDefinition<TContext, TEvent>> {
     return transitionNode(this.root, state.value, state, event) || [];
@@ -285,7 +285,7 @@ export class StateMachine<
   private getPreInitialState(
     actorCtx: AnyActorContext,
     initEvent: any
-  ): State<TContext, TEvent, TActor, TOutput, TResolvedTypesMeta> {
+  ): State<TContext, TEvent, TActor, TTag, TOutput, TResolvedTypesMeta> {
     const { context } = this.config;
 
     const preInitial = this.resolveState(
@@ -319,10 +319,10 @@ export class StateMachine<
   public getInitialState(
     actorCtx: ActorContext<
       TEvent,
-      State<TContext, TEvent, TActor, TOutput, TResolvedTypesMeta>
+      State<TContext, TEvent, TActor, TTag, TOutput, TResolvedTypesMeta>
     >,
     input?: TInput
-  ): State<TContext, TEvent, TActor, TOutput, TResolvedTypesMeta> {
+  ): State<TContext, TEvent, TActor, TTag, TOutput, TResolvedTypesMeta> {
     const initEvent = createInitEvent(input) as unknown as TEvent; // TODO: fix;
 
     const preInitialState = this.getPreInitialState(actorCtx, initEvent);
@@ -353,7 +353,7 @@ export class StateMachine<
   }
 
   public start(
-    state: State<TContext, TEvent, TActor, TOutput, TResolvedTypesMeta>
+    state: State<TContext, TEvent, TActor, TTag, TOutput, TResolvedTypesMeta>
   ): void {
     Object.values(state.children).forEach((child: any) => {
       if (child.status === 0) {
@@ -387,25 +387,25 @@ export class StateMachine<
   }
 
   public getPersistedState(
-    state: State<TContext, TEvent, TActor, TOutput, TResolvedTypesMeta>
+    state: State<TContext, TEvent, TActor, TTag, TOutput, TResolvedTypesMeta>
   ): PersistedMachineState<
-    State<TContext, TEvent, TActor, TOutput, TResolvedTypesMeta>
+    State<TContext, TEvent, TActor, TTag, TOutput, TResolvedTypesMeta>
   > {
     return getPersistedState(state);
   }
 
   public createState(
     stateConfig:
-      | State<TContext, TEvent, TActor, TOutput, TResolvedTypesMeta>
+      | State<TContext, TEvent, TActor, TTag, TOutput, TResolvedTypesMeta>
       | StateConfig<TContext, TEvent>
-  ): State<TContext, TEvent, TActor, TOutput, TResolvedTypesMeta> {
+  ): State<TContext, TEvent, TActor, TTag, TOutput, TResolvedTypesMeta> {
     return stateConfig instanceof State
       ? stateConfig
       : new State(stateConfig, this);
   }
 
   public getStatus(
-    state: State<TContext, TEvent, TActor, TOutput, TResolvedTypesMeta>
+    state: State<TContext, TEvent, TActor, TTag, TOutput, TResolvedTypesMeta>
   ) {
     return state.error
       ? { status: 'error', data: state.error }
@@ -416,13 +416,13 @@ export class StateMachine<
 
   public restoreState(
     state: PersistedMachineState<
-      State<TContext, TEvent, TActor, TOutput, TResolvedTypesMeta>
+      State<TContext, TEvent, TActor, TTag, TOutput, TResolvedTypesMeta>
     >,
     _actorCtx: ActorContext<
       TEvent,
-      State<TContext, TEvent, TActor, TOutput, TResolvedTypesMeta>
+      State<TContext, TEvent, TActor, TTag, TOutput, TResolvedTypesMeta>
     >
-  ): State<TContext, TEvent, TActor, TOutput, TResolvedTypesMeta> {
+  ): State<TContext, TEvent, TActor, TTag, TOutput, TResolvedTypesMeta> {
     const children: Record<string, AnyActorRef> = {};
 
     Object.keys(state.children).forEach((actorId) => {
@@ -452,6 +452,7 @@ export class StateMachine<
       TContext,
       TEvent,
       TActor,
+      TTag,
       TOutput,
       TResolvedTypesMeta
     > = this.createState(new State({ ...state, children }, this));

diff --git a/packages/core/src/StateNode.ts b/packages/core/src/StateNode.ts
@@ -372,7 +372,7 @@ export class StateNode<
   }
 
   public next(
-    state: State<TContext, TEvent, TODO, TODO>,
+    state: State<TContext, TEvent, TODO, TODO, TODO, TODO>,
     event: TEvent
   ): TransitionDefinition<TContext, TEvent>[] | undefined {
     const eventType = event.type;

diff --git a/packages/core/src/stateUtils.ts b/packages/core/src/stateUtils.ts
@@ -670,7 +670,7 @@ export function getStateNodes<
   TEvent extends EventObject
 >(
   stateNode: AnyStateNode,
-  state: StateValue | State<TContext, TEvent, TODO, TODO>
+  state: StateValue | State<TContext, TEvent, TODO, TODO, TODO, TODO>
 ): Array<AnyStateNode> {
   const stateValue = state instanceof State ? state.value : toStateValue(state);
 
@@ -706,7 +706,7 @@ export function transitionAtomicNode<
 >(
   stateNode: AnyStateNode,
   stateValue: string,
-  state: State<TContext, TEvent, TODO, TODO>,
+  state: State<TContext, TEvent, TODO, TODO, TODO, TODO>,
   event: TEvent
 ): Array<TransitionDefinition<TContext, TEvent>> | undefined {
   const childStateNode = getStateNode(stateNode, stateValue);
@@ -725,7 +725,7 @@ export function transitionCompoundNode<
 >(
   stateNode: AnyStateNode,
   stateValue: StateValueMap,
-  state: State<TContext, TEvent, TODO, TODO>,
+  state: State<TContext, TEvent, TODO, TODO, TODO, TODO>,
   event: TEvent
 ): Array<TransitionDefinition<TContext, TEvent>> | undefined {
   const subStateKeys = Object.keys(stateValue);
@@ -751,7 +751,7 @@ export function transitionParallelNode<
 >(
   stateNode: AnyStateNode,
   stateValue: StateValueMap,
-  state: State<TContext, TEvent, TODO, TODO>,
+  state: State<TContext, TEvent, TODO, TODO, TODO, TODO>,
   event: TEvent
 ): Array<TransitionDefinition<TContext, TEvent>> | undefined {
   const allInnerTransitions: Array<TransitionDefinition<TContext, TEvent>> = [];
@@ -787,7 +787,14 @@ export function transitionNode<
 >(
   stateNode: AnyStateNode,
   stateValue: StateValue,
-  state: State<TContext, TEvent, TODO, TODO, TODO>,
+  state: State<
+    TContext,
+    TEvent,
+    TODO,
+    TODO,
+    TODO, // output
+    TODO // tags
+  >,
   event: TEvent
 ): Array<TransitionDefinition<TContext, TEvent>> | undefined {
   // leaf node

diff --git a/packages/core/src/types.ts b/packages/core/src/types.ts
@@ -784,7 +784,14 @@ export type AnyStateNode = StateNode<any, any>;
 
 export type AnyStateNodeDefinition = StateNodeDefinition<any, any>;
 
-export type AnyState = State<any, any, any, any, any>;
+export type AnyState = State<
+  any,
+  any,
+  any,
+  any,
+  any, // delays
+  any // tags
+>;
 
 export type AnyStateMachine = StateMachine<
   any,
@@ -1663,7 +1670,7 @@ export type ActorRefFrom<T> = ReturnTypeOrValue<T> extends infer R
       infer _TAction,
       infer _TGuard,
       infer _TDelay,
-      infer _TTag,
+      infer TTag,
       infer _TInput,
       infer TOutput,
       infer TResolvedTypesMeta
@@ -1674,6 +1681,7 @@ export type ActorRefFrom<T> = ReturnTypeOrValue<T> extends infer R
           TContext,
           TEvent,
           TActor,
+          TTag,
           TOutput,
           AreAllImplementationsAssumedToBeProvided<TResolvedTypesMeta> extends false
             ? MarkAllImplementationsAsProvided<TResolvedTypesMeta>
@@ -1709,18 +1717,18 @@ export type InterpreterFrom<
   infer _TAction,
   infer _TGuard,
   infer _TDelay,
-  infer _TTag,
+  infer TTag,
   infer TInput,
   infer TOutput,
   infer TResolvedTypesMeta
 >
   ? Actor<
       ActorLogic<
         TEvent,
-        State<TContext, TEvent, TActor, TOutput, TResolvedTypesMeta>,
-        State<TContext, TEvent, TActor, TOutput, TResolvedTypesMeta>,
+        State<TContext, TEvent, TActor, TTag, TOutput, TResolvedTypesMeta>,
+        State<TContext, TEvent, TActor, TTag, TOutput, TResolvedTypesMeta>,
         PersistedMachineState<
-          State<TContext, TEvent, TActor, TOutput, TResolvedTypesMeta>
+          State<TContext, TEvent, TActor, TTag, TOutput, TResolvedTypesMeta>
         >,
         ActorSystem<any>,
         TInput

diff --git a/packages/core/test/types.test.ts b/packages/core/test/types.test.ts
@@ -2972,4 +2972,19 @@ describe('tags', () => {
       }
     });
   });
+
+  it('typed tags should work in state', () => {
+    const machine = createMachine({
+      types: {} as {
+        tags: 'a' | 'b' | 'c';
+      }
+    });
+
+    const actor = createActor(machine).start();
+
+    actor.getSnapshot().hasTag('a');
+    actor.getSnapshot().hasTag('b');
+    // @ts-expect-error
+    actor.getSnapshot().hasTag('unknown');
+  });
 });
diff --git a/packages/xstate-solid/src/types.ts b/packages/xstate-solid/src/types.ts
@@ -15,10 +15,11 @@ type StateObject<
   TContext extends MachineContext,
   TEvent extends EventObject = EventObject,
   TActor extends ProvidedActor = ProvidedActor,
+  TTag extends string = string,
   TOutput = unknown,
   TResolvedTypesMeta = TypegenDisabled
 > = Pick<
-  State<TContext, TEvent, TActor, TOutput, TResolvedTypesMeta>,
+  State<TContext, TEvent, TActor, TTag, TOutput, TResolvedTypesMeta>,
   keyof AnyState
 >;
 
@@ -28,10 +29,11 @@ export type CheckSnapshot<Snapshot> = Snapshot extends State<
   infer TContext,
   infer TEvents,
   infer TActor,
+  infer TTag,
   infer TOutput,
   infer TResolvedTypesMeta
 >
-  ? StateObject<TContext, TEvents, TActor, TOutput, TResolvedTypesMeta>
+  ? StateObject<TContext, TEvents, TActor, TTag, TOutput, TResolvedTypesMeta>
   : Snapshot;
 
 type InternalMachineOpts<

diff --git a/packages/xstate-test/src/types.ts b/packages/xstate-test/src/types.ts
@@ -96,9 +96,11 @@ export type TestMachineOptions<
 export interface TestMeta<T, TContext extends MachineContext> {
   test?: (
     testContext: T,
-    state: State<TContext, any, any, any>
+    state: State<TContext, any, any, any, any, any>
   ) => Promise<void> | void;
-  description?: string | ((state: State<TContext, any, any, any>) => string);
+  description?:
+    | string
+    | ((state: State<TContext, any, any, any, any, any>) => string);
   skip?: boolean;
 }
 interface TestStateResult {
@@ -168,7 +170,7 @@ export interface TestTransitionConfig<
   TTestContext
 > extends TransitionConfig<TContext, TEvent, TEvent, TODO, TODO> {
   test?: (
-    state: State<TContext, TEvent, any, any>,
+    state: State<TContext, TEvent, any, any, any, any>,
     testContext: TTestContext
   ) => void;
 }