stateful.h

#pragma once
 
#include <cstddef>
#include <cstdint>
#include <functional>
#include <map>
#include <stdexcept>
#include <string>
#include <type_traits>
#include <utility>
#include <vector>
 
#include "hegel/core.h"

namespace hegel::stateful {
 
    using hegel::generators::Generator;
    using hegel::generators::IGenerator;
 
    template <typename T> class VariablesGenerator;

    template <typename T> class Pool {
      public:
        explicit Pool(const TestCase& tc) : tc_(tc), pool_handle_(tc) {}

        void add(const T& element) {
            int64_t var_id = pool_handle_.add(tc_);
            // GCOVR_EXCL_START
            if (pool_.find(var_id) != pool_.end()) {
                throw std::runtime_error("unexpected variable id in map");
            }
            // GCOVR_EXCL_STOP
            pool_.emplace(var_id, element);
        }

        void add(T&& element) {
            int64_t var_id = pool_handle_.add(tc_);
            // GCOVR_EXCL_START
            if (pool_.find(var_id) != pool_.end()) {
                throw std::runtime_error("unexpected variable id in map");
            }
            // GCOVR_EXCL_STOP
            pool_.emplace(var_id, std::move(element));
        }

        std::size_t size() const { return pool_.size(); }
 
        Pool(const Pool&) = delete;
 
      private:
        const TestCase& tc_;
        hegel::internal::PoolHandle pool_handle_;
        std::map<int64_t, T> pool_;
 
        friend class VariablesGenerator<T>;
    };

    // Concrete IGenerator for variables.
    template <typename T> class VariablesGenerator : public IGenerator<T> {
      public:
        explicit VariablesGenerator(Pool<T>& p, bool consume)
            : p_(p), consume_(consume) {}
 
        T do_draw(const TestCase& tc) const override {
            int64_t variable = p_.pool_handle_.draw_variable(tc, consume_);
 
            auto it = p_.pool_.find(variable);
            // GCOVR_EXCL_START
            if (it == p_.pool_.end()) {
                throw std::runtime_error(
                    "Pool state diverged between the engine and the "
                    "client, or a bug in the pool bookkeeping.");
            }
            // GCOVR_EXCL_STOP
            if (consume_) {
                auto val = std::move(it->second);
                p_.pool_.erase(it);
                return val;
            }
            return it->second;
        }
 
      private:
        Pool<T>& p_;
        bool consume_;
    };


    template <typename T> Generator<T> values_consumed(Pool<T>& p) {
        return Generator<T>(new VariablesGenerator<T>(p, true));
    }

    template <typename T> Generator<T> values_reusable(Pool<T>& p) {
        return Generator<T>(new VariablesGenerator<T>(p, false));
    }

    template <typename T> class Rule {
      public:
        explicit Rule(std::string name, std::function<void(TestCase&, T&)> step)
            : name_(std::move(name)), step_(std::move(step)) {}

        const std::string& name() const { return name_; }
        const std::function<void(TestCase&, T&)>& step() const { return step_; }
 
      private:
        std::string name_;
        std::function<void(TestCase&, T&)> step_;
    };

    template <typename T> class Invariant {
      public:
        explicit Invariant(std::string name,
                           std::function<void(const T&)> invariant)
            : name_(std::move(name)), invariant_(std::move(invariant)) {}

        const std::string& name() const { return name_; }
        const std::function<void(const T&)>& invariant() const {
            return invariant_;
        }
 
      private:
        std::string name_;
        std::function<void(const T&)> invariant_;
    };

    template <typename State> struct StateMachineParams {
        State initial_state;
    };

    template <typename Derived, typename State> class StateMachine {
      public:
        using state_type = State;
        explicit StateMachine(StateMachineParams<State> params)
            : state(std::move(params.initial_state)) {}

        State state;

        std::vector<Invariant<Derived>> invariants() { return {}; }
 
      protected:
        ~StateMachine() = default;
    };

    struct RunParams {
        bool print_state = true;
    };

    // prints the machine's state.
    template <typename M>
    void print_state(TestCase& tc, const M& machine, const RunParams& params) {
        if (params.print_state) {
            tc.note("state = " + internal::repr(machine.state));
        }
    }
 
    // check if invariants hold on a given state
    template <typename T>
    void check_invariants(TestCase& tc, const std::string& origin,
                          const T& state,
                          const std::vector<Invariant<T>>& invariants) {
        for (const auto& inv : invariants) {
            try {
                (inv.invariant())(state);
            } catch (...) {
                tc.note("Invariant " + inv.name() + " violated " + origin);
                throw;
            }
        }
    }

    template <typename M>
    void run(M& machine, TestCase& tc, const RunParams& params = {}) {
        static_assert(
            std::is_base_of<StateMachine<M, typename M::state_type>, M>::value,
            "run() requires a machine deriving from "
            "StateMachine<M, State>.");
        std::vector<Rule<M>> rules = machine.rules();
        std::vector<Invariant<M>> invariants = machine.invariants();
        if (rules.empty()) {
            throw std::invalid_argument(
                "Cannot run a state machine with no rules.");
        }
        std::vector<std::string> rule_names;
        rule_names.reserve(rules.size());
        for (const Rule<M>& rule : rules)
            rule_names.push_back(rule.name());
 
        std::vector<std::string> invariant_names;
        invariant_names.reserve(invariants.size());
        for (const Invariant<M>& invariant : invariants)
            invariant_names.push_back(invariant.name());
 
        print_state(tc, machine, params);
        check_invariants(tc, "in the initial state", machine, invariants);
 
        internal::StateMachineHandle machine_handle(tc, rule_names,
                                                    invariant_names);
        int64_t steps_run = 0;
 
        while (true) {
            internal::start_span(tc, internal::SpanLabel::StatefulRule);
            int64_t next_rule_idx = machine_handle.next_rule(tc);
            if (next_rule_idx == internal::state_machine_done) {
                break;
            }
            // GCOVR_EXCL_START
            if (next_rule_idx < 0 ||
                static_cast<size_t>(next_rule_idx) >= rules.size()) {
                throw std::runtime_error(
                    "state_machine_next_rule returned out-of-range "
                    "rule index. Please report this as a bug.");
            }
            // GCOVR_EXCL_STOP
            steps_run++;
            const Rule<M>& rule = rules[static_cast<size_t>(next_rule_idx)];
            tc.note("Step " + std::to_string(steps_run) + ": " + rule.name());
 
            try {
                // nest the draws the step makes under its "Step N" header.
                {
                    internal::NoteIndentScope indent(tc);
                    rule.step()(tc, machine);
                }
                print_state(tc, machine, params);
                check_invariants(tc, "after step " + std::to_string(steps_run),
                                 machine, invariants);
                internal::stop_span(tc);
            } catch (const internal::HegelReject&) {
                tc.note("Rule stopped early due to violated assumption.");
                machine_handle.rule_rejected(tc);
                internal::stop_span(tc, true);
            } catch (...) {
                internal::stop_span(tc);
                throw;
            }
        }
    }

 
} // namespace hegel::stateful