state machine: add handler to state in a type safe way

useStateWithAction.ts

type AnyState = { step: string };
type AnyActionMap = Record<string, (...args: any[]) => unknown>;
type AnyAction<State extends AnyState> = {
	[S in State['step']]?: (
		state: { step: S } & State,
		setState: (newState: SetStateAction<State>) => void
	) => AnyActionMap;
};

type StateWithActions<
	State extends AnyState,
	Action extends AnyAction<State>
> = {
	[S in State['step']]: { step: S } & State &
		(Action[S] extends (...args: any[]) => infer ActionMap
			? {
					action: ActionMap;
			  }
			: { action: Record<string, never> });
}[State['step']];

export function useStateWithAction<
	State extends AnyState,
	Action extends AnyAction<State>
>(initialState: State, action: Action): StateWithActions<State, Action> {
	const [state, setState] = useState<AnyState>(initialState);
	const createAction = (action as unknown as AnyAction<AnyState>)[state.step];
	return {
		...state,
		action: createAction ? createAction(state, setState) : {},
	} as unknown as StateWithActions<State, AnyAction<State>>;
}

type MyState =
    { step: '1', data1: number }
    | { step: '2', data2: string }

let initialState: MyState;
const current = useStateWithAction(initialState, {
    '1': (state, setState) => ({
        'to2': () => setState({ step: '2', data2: state.data1.toString() }),
        'add': (val: number) => setState({ step: '1', data1: state.data1 + val }),
    }),
    '2': (state, setState) => ({
        'to1': () => setState({ step: '1', data1: Number(state.data2) }),
    })
});

if (current.step === '1') {
    current.data1
    current.action.to2()
    current.action.add(33)
}
if (current.step === '2') {
    current.data2
    current.action.to1()
}