A generic finite state machine implementation for managing stateful transitions.

FSM.java

import java.util.EnumMap;
import java.util.Map;
import java.util.Objects;
import java.util.function.UnaryOperator;

/**
 * A generic finite state machine implementation for managing state transitions.
 *
 * @param <S> the type of the state, which must be an enum.
 * @param <D> the type of the data associated with the states.
 * @author <a href="mailto:[email protected]">Diego Rocha</a>
 */
final class FSM<S extends Enum<S>, D> {

    private final Map<S, Map<S, Transition<S, D>>> transitions;
    private State<S, D> state;

    /**
     * Constructs a finite state machine with the given initial state.
     *
     * @param initial the initial state of the state machine.
     */
    private FSM(final State<S, D> initial) {
        this.transitions = new EnumMap<>(initial.type.getDeclaringClass());
        this.state = initial;
        for (final var v : initial.type.getDeclaringClass().getEnumConstants()) {
            this.transitions.put(v, new EnumMap<>(initial.type.getDeclaringClass()));
        }
    }

    /**
     * Factory method for creating an {@link FSM} instance.
     *
     * @param initial the initial state of the state machine.
     * @param state   the data associated with the initial state.
     * @param <S>     the type of the state, which must be an enum.
     * @param <D>     the type of the data associated with the states.
     * @return a new FSM instance initialized with the specified state.
     */
    static <S extends Enum<S>, D> FSM<S, D> on(final S initial, final D state) {
        return new FSM<>(new State<>(initial, state));
    }

    /**
     * Adds a transition from one state to another.
     *
     * @param in    the initial state from which the transition occurs.
     * @param out   the resulting state to which the transition leads.
     * @param trans the transition to be applied for this state change.
     */
    void add(final S in, final S out, final Transition<S, D> trans) {
        this.transitions.get(in).put(out, trans);
    }

    /**
     * Moves the state machine to the next state.
     *
     * @param next the next state to transition to.
     * @throws IllegalArgumentException if there is no transition defined from the current state to the next state.
     */
    void moveTo(final S next) {
        final var trans = this.transitions.get(this.state.type).get(next);
        if (trans == null) {
            throw new IllegalArgumentException("No transition from " + this.state.type + " to " + next);
        }
        this.state = trans.apply(this.state);
    }

    /**
     * A functional interface representing a transition between states.
     *
     * @param <S> the type of the state, which must be an enum.
     * @param <D> the type of the data associated with the states.
     * @author <a href="mailto:[email protected]">Diego Rocha</a>
     */
    interface Transition<S extends Enum<S>, D> extends UnaryOperator<State<S, D>> {
    }

    /**
     * A record representing a state with a type and associated data.
     *
     * @param type the type of the state.
     * @param data the data associated with the state.
     * @param <S>  the type of the state, which must be an enum.
     * @param <D>  the type of the data associated with the states.
     * @author <a href="mailto:[email protected]">Diego Rocha</a>
     */
    record State<S extends Enum<S>, D>(S type, D data) {

        public State {
            Objects.requireNonNull(type, "State type cannot be null");
            Objects.requireNonNull(data, "State data cannot be null");
        }
    }
}