Typescript state machine that prevents dis-allowed transitions with compile errors

Question

I am trying to create a state machine type where the transitions are enforced by the compiler.

I want to define the state machine like this:

type States = {
    state1: {state: 'state1'},
    state2: {state: 'state2'}
}

const Transitions: TransitionsType<States> = {
    init() {
        return {state: 'state1'}
    },
    state1: {
        state2({...anyOtherItems}) {
            // processing with anyOtherItems
            // maybe return a load of extra stuff here
            return {...anyOtherItems, state: 'state2'}
        }
    },
    state2: {
        state1({...others}) {
            // same deal here
            return {state: 'state1'}
        }
    }
};

transition(transitions, {state: 'state1'}, 'state2');

States is a type that describes the states of the machine, and a datatype for that state.

Transitions is an object that has the from state as the first level of the map, and the to is the second level of the map. The functions will be called for each transition, and they will take the type related to the from state, and return the type related to the to state. The init function is to kick it all off.

transition is a function that takes a state machine definition, a state object and a to state, and returns a new state object.

I want to get compile errors if the transition functions return objects with the wrong states, and if the transition function is called with a state object and a to state that are disallowed.

So far I have (and feel free to completely ignore this):

type StatesType<S> = {
    [T in keyof S]: {state: T};
};

type TransitionTosType<S extends StatesType<S>, F extends keyof StatesType<S>> = {
    [T in keyof S]?: (arg: S[F]) => S[T]
}

type TransitionsType<S extends StatesType<S>> = {
    [F in keyof S]: TransitionTosType<S, F>
} & {
    init: () => S[keyof S]
}

function transition<S extends StatesType<S>>(transitions: TransitionsType<S>,
                                            state: S[keyof S],
                                            to: keyof S & string): {state: string} {
    const fromStateTransitions = transitions[state.state as keyof TransitionsType<S>];
    if (fromStateTransitions) {
        const toTransition = fromStateTransitions[to];
        if (toTransition)
            return toTransition(state);
    }
    throw new TypeError(`no transition from ${state.state} to ${String(to)}`);
}

This gives me an error on the line return toTransition(state);, and it does not give me an error if I try to transition to a disallowed state.

I'd be happy to accept any answer that can take the states and transitions given (States and Transitions in the example), and provides a function which takes a set of transitions (e.g. Transitions), a state context of the correct type (as given in a type like States), and string for the transition to, and where there will be compile errors if the state context is the wrong type, or the to state isn't a valid transition from the from state.

I am fine with the transition function accepting type parameters.

I will also accept an explanation of why it is impossible as an answer.

Answer 1

The main problem I see with your code example is that you have a conflict between manually specifying that transitions is of type TransitionsType<States>, and allowing the compiler to infer something more specific from the initializing object literal. The type TransitionsType<States> has a fairly explicit representation of the States type, which you need for transition(transitions, ⋯) to function correctly. But if you use that type then the compiler has no idea about the specifics of the state graph and which nodes are connected to which. You really need to use a narrower type, such as shown here:

const transitions = {
  init() {
    return { state: 'state1' }
  },
  state1: {
    state2({ ...anyOtherItems }) {
      return { ...anyOtherItems, state: 'state2' }
    }
  },
  state2: {
    state1({ ...others }) {
      return { state: 'state1' }
    }
  }
} satisfies TransitionsType<States>;

where I've used the satisfies operator to both verify the assignability to that type and to provide a context in which to infer the type of transitions. That inferred type is:

const transitions: {
    init(): {
        state: "state1";
    };
    state1: {
        state2({ ...anyOtherItems }: {
            state: 'state1';
        }): {
            state: "state2";
        };
    };
    state2: {
        state1({ ...others }: {
            state: 'state2';
        }): {
            state: "state1";
        };
    };
}

But then the type of transitions doesn't explicitly contain the States type. So you'd need transition(transitions, ⋯) to be able to tease apart the type of transitions to recover the original States type, which isn't necessarily fully present (imagine a graph where one of the states is disconnected).

There might be clever ways to get the best of both worlds here, but in my attempts the typing for transition() became so complex that it doesn't seem worth it to me.

---

Instead, my approach would be to accept that you need to explicitly pass two pieces of information to transition(): both the States type at the type level, and the specific shape of transitions. Ideally then the call would look like

transition<States>(transitions, state, to);

But unfortunately, that type-level type passing would require what's known as "partial type argument inference" as requested in microsoft/TypeScript#26272, which isn't currently supported in TypeScript. See Typescript: infer type of generic after optional first generic, for example. The workarounds involve some refactoring. The most common one is a type-level currying so that instead of transition<States>(transitions, state, to), the call transition<States>() returns a function which you call with (transitions, state, to) as arguments.

It could look like this:

function transition<S extends StatesType<S>>() {

  function f<T extends TransitionsType<S>, F extends S[keyof S]>(
    transitions: T,
    state: F,
    to: keyof T[F["state"] & keyof T]
  ): void;

  function f(transitions: any,
    state: any,
    to: any): { state: string } {
    const fromStateTransitions = transitions[state.state];
    if (fromStateTransitions) {
      const toTransition = fromStateTransitions[to];
      if (toTransition)
        return toTransition(state);
    }
    throw new TypeError(`no transition from ${state.state} to ${String(to)}`);
  }
  return f;
}

So now when you call transition<States>(), the type of S inside the function is set in stone, and can be used to constrain T (the generic type corresponding to transitions, and F (the generic type corresponding to state), and consequently the type of to is an easily-written function of F and T.

That is, if T is the type of transitions, and F is the type of state, then to must be of type keyof T[F["state"]] (one of the keys of the property of transitions corresponding to the state F). Note that I had to write T[F["state"] & keyof T] to convince the compiler that F["state"] is acceptable as a key of T. It's not directly written that this is true, at least not for generics, and the compiler can't deduce it. An intersection frees up the compiler from worrying about it.

---

Okay, let's try it out:

const transitionFunc = transition<States>();
transitionFunc(transitions, { state: 'state1' }, 'state2'); // okay
transitionFunc(transitions, { state: 'state1' }, 'state1'); // error
transitionFunc(transitions, { state: 'state2' }, 'state2'); // error
transitionFunc(transitions, { state: 'state2' }, 'state1'); // okay
transitionFunc(transitions, { state: 'state3' }, 'state2'); // error

Looks good. The compiler only allows you to pass in arguments of the appropriate types, and will warn you if your to state is not directly accessible from state.

Playground link to code