state_machine.hpp

Go to the documentation of this file.

#ifndef STATE_MACHINE_HPP
#define STATE_MACHINE_HPP
 
#include <functional>
#include <unordered_map>
#include <vector>
#include <string>
#include <optional>

template <typename StateType> class PureStateMachine {
  public:
    using Condition = std::function<bool()>;

    struct Transition {
        StateType to;        
        Condition condition; 
    };

    enum class EventType {
        ENTER_STATE,  
        UPDATE_STATE, 
        EXIT_STATE,   
        TRANSITION    
    };

    struct Event {
        EventType type;       
        StateType state;      
        StateType next_state; 
    };

    enum class Mode {
        MULTI_EVENT, 
        SINGLE_EVENT 
    };

    explicit PureStateMachine(StateType initial_state, Mode mode = Mode::MULTI_EVENT)
        : current_state(initial_state), mode(mode) {}

    void add_state(StateType state) { states[state] = {}; }

    void add_transition(StateType from, StateType to, Condition condition) { states[from].push_back({to, condition}); }

    [[nodiscard]] std::vector<Event> update_pure() {
        switch (mode) {
        case Mode::MULTI_EVENT:
            return update_multi_event();
        case Mode::SINGLE_EVENT:
            return update_single_event();
        default:
            return {};
        }
    }

    [[nodiscard]] StateType get_state() const { return current_state; }

    void set_mode(Mode new_mode) { mode = new_mode; }

    [[nodiscard]] Mode get_mode() const { return mode; }
 
  private:
    [[nodiscard]] std::vector<Event> update_multi_event() {
        std::vector<Event> events;
 
        bool state_changed = current_state != previous_state;
        if (state_changed) {
            if (there_is_a_previous_state)
                events.push_back({EventType::EXIT_STATE, previous_state, current_state});
 
            events.push_back({EventType::ENTER_STATE, current_state, current_state});
            previous_state = current_state;
            there_is_a_previous_state = true;
        }
 
        events.push_back({EventType::UPDATE_STATE, current_state, current_state});
 
        for (auto &t : states[current_state]) {
            if (t.condition && t.condition()) {
                events.push_back({EventType::TRANSITION, current_state, t.to});
                current_state = t.to;
                break;
            }
        }
 
        return events;
    }

    [[nodiscard]] std::vector<Event> update_single_event() {
        bool state_changed = current_state != previous_state;
 
        if (state_changed) {
            if (there_is_a_previous_state) {
                previous_state = current_state;
                there_is_a_previous_state = true;
                return {{EventType::EXIT_STATE, previous_state, current_state}};
            } else {
                previous_state = current_state;
                there_is_a_previous_state = true;
                return {{EventType::ENTER_STATE, current_state, current_state}};
            }
        }
 
        // Evaluate transitions first
        for (auto &t : states[current_state]) {
            if (t.condition && t.condition()) {
                current_state = t.to;
                return {{EventType::TRANSITION, previous_state, current_state}};
            }
        }
 
        // Default: update event
        return {{EventType::UPDATE_STATE, current_state, current_state}};
    }
 
  private:
    StateType current_state;                
    StateType previous_state;               
    bool there_is_a_previous_state = false; 
    Mode mode;                              
 
    std::unordered_map<StateType, std::vector<Transition>> states; 
};

template <typename StateType> class StateMachine {
  public:
    using Action = std::function<void()>;
    using Condition = std::function<bool()>;
    using Pure = PureStateMachine<StateType>;
 
    struct State {
        std::optional<Action> on_enter;
        std::optional<Action> on_update;
        std::optional<Action> on_exit;
    };
 
    explicit StateMachine(StateType initial_state) : pure_fsm(initial_state) {}
 
    void add_state(StateType state, std::optional<Action> on_enter = std::nullopt,
                   std::optional<Action> on_update = std::nullopt, std::optional<Action> on_exit = std::nullopt) {
        pure_fsm.add_state(state);
        states[state] = {on_enter, on_update, on_exit};
    }
 
    void add_transition(StateType from, StateType to, Condition condition) {
        pure_fsm.add_transition(from, to, condition);
    }

    void update() {
        auto events = pure_fsm.update_pure();
 
        for (auto &event : events) {
 
            // NOTE: the order in this switch defines that exit is run first, then enter, and then update
            // which is how we would want it to be in terms of linear time
            switch (event.type) {
 
            case Pure::EventType::EXIT_STATE:
                if (states[event.state].on_exit)
                    (*states[event.state].on_exit)();
                break;
 
            case Pure::EventType::ENTER_STATE:
                if (states[event.state].on_enter)
                    (*states[event.state].on_enter)();
                break;
 
            case Pure::EventType::UPDATE_STATE:
                if (states[event.state].on_update)
                    (*states[event.state].on_update)();
                break;
 
            default:
                break;
            }
        }
    }
 
    [[nodiscard]] StateType get_state() const { return pure_fsm.get_state(); }
 
    [[nodiscard]] const Pure &get_pure_fsm() const { return pure_fsm; }
 
  private:
    Pure pure_fsm;
    std::unordered_map<StateType, State> states;
};
 
#endif // STATE_MACHINE_HPP