stateful.h

#pragma once
 
#include <condition_variable>
#include <cstddef>
#include <cstdint>
#include <exception>
#include <functional>
#include <map>
#include <memory>
#include <mutex>
#include <optional>
#include <stdexcept>
#include <string>
#include <thread>
#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 Invariant;

    inline constexpr const char* anonymous_group = "<anonymous>";

    template <typename M> class ConcurrentRule {
      public:
        using Function = std::function<void(TestCase&, M&)>;

        ConcurrentRule(std::string name, Function function)
            : ConcurrentRule(std::move(name), anonymous_group,
                             std::move(function)) {}

        ConcurrentRule(std::string name, std::string group, Function function)
            : name_(std::move(name)), group_(std::move(group)),
              function_(std::move(function)) {}

        const std::string& name() const { return name_; }

        const std::string& group() const { return group_; }

        const Function& function() const { return function_; }
 
      private:
        std::string name_;
        std::string group_;
        Function function_;
    };

    template <typename Derived> class ConcurrentStateMachine {
      public:
        std::vector<Invariant<Derived>> invariants() const { return {}; }
    };

    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 ConcurrentVariablesGenerator;

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

        void add(const TestCase& tc, T value) {
            std::lock_guard<std::mutex> lock(mutex_);
            int64_t variable = pool_handle_.add(tc);
            values_.emplace(variable, std::move(value));
        }

        bool empty() const {
            std::lock_guard<std::mutex> lock(mutex_);
            return values_.empty();
        }

        std::size_t size() const {
            std::lock_guard<std::mutex> lock(mutex_);
            return values_.size();
        }
 
        ConcurrentPool(const ConcurrentPool&) = delete;
        ConcurrentPool& operator=(const ConcurrentPool&) = delete;
 
      private:
        internal::PoolHandle pool_handle_;
        mutable std::mutex mutex_;
        std::map<int64_t, T> values_;
 
        friend class ConcurrentVariablesGenerator<T>;
    };

    template <typename T>
    class ConcurrentVariablesGenerator : public IGenerator<T> {
      public:
        ConcurrentVariablesGenerator(ConcurrentPool<T>& pool, bool consume)
            : pool_(pool), consume_(consume) {}
 
        T do_draw(const TestCase& tc) const override {
            std::lock_guard<std::mutex> lock(pool_.mutex_);
            int64_t variable = pool_.pool_handle_.draw_variable(tc, consume_);
            if (consume_) {
                T result = std::move(pool_.values_.at(variable));
                pool_.values_.erase(variable);
                return result;
            } else {
                return pool_.values_.at(variable);
            }
        }
 
      private:
        ConcurrentPool<T>& pool_;
        bool consume_;
    };

    template <typename T>
    Generator<T> values_consumed(ConcurrentPool<T>& pool) {
        return Generator<T>(new ConcurrentVariablesGenerator<T>(pool, true));
    }

    template <typename T>
    Generator<T> values_reusable(ConcurrentPool<T>& pool) {
        return Generator<T>(new ConcurrentVariablesGenerator<T>(pool, 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));
        }
    }
 
    template <typename T>
    void check_invariants(
        TestCase& tc, const std::string& origin, const T& state,
        const std::vector<Invariant<T>>& invariants,
        std::optional<std::reference_wrapper<internal::StateMachineHandle>>
            machine_handle = std::nullopt) {
        for (std::size_t index = 0; index < invariants.size(); ++index) {
            if (machine_handle.has_value() &&
                !machine_handle->get().should_check_invariant(
                    tc, static_cast<int64_t>(index))) {
                continue;
            }
            const Invariant<T>& inv = invariants[index];
            try {
                inv.invariant()(state);
            } catch (...) {
                tc.note("Invariant " + inv.name() + " violated " + origin);
                throw;
            }
        }
    }

    template <typename M>
    void run_concurrent(M& machine, TestCase& tc, int64_t min_concurrency,
                        int64_t max_concurrency) {
        static_assert(std::is_base_of<ConcurrentStateMachine<M>, M>::value,
                      "run_concurrent() requires a machine deriving from "
                      "ConcurrentStateMachine<M>.");
        if (min_concurrency < 1 || min_concurrency > max_concurrency) {
            throw std::invalid_argument(
                "concurrency bounds must satisfy 1 <= min <= max");
        }
 
        std::vector<ConcurrentRule<M>> rules = machine.rules();
        std::vector<Invariant<M>> invariants = machine.invariants();
        if (rules.empty()) {
            throw std::invalid_argument(
                "Cannot run a concurrent state machine with no rules.");
        }
 
        std::vector<std::string> rule_names;
        std::vector<std::string> invariant_names;
        std::vector<std::string> group_names;     // maps group ID to name
        std::map<std::string, int64_t> group_ids; // maps group name to ID
        std::vector<int64_t>
            rule_groups; // ID of group that rule_names[i] belongs to is the ID
                         // at rule_groups[i].
        rule_names.reserve(rules.size());
        rule_groups.reserve(rules.size());
        for (const ConcurrentRule<M>& rule : rules) {
            rule_names.push_back(rule.name());
            auto [group, inserted] = group_ids.emplace(
                rule.group(), static_cast<int64_t>(group_ids.size()));
            if (inserted) {
                group_names.push_back(rule.group());
            }
            rule_groups.push_back(group->second);
        }
        invariant_names.reserve(invariants.size());
        for (const Invariant<M>& invariant : invariants) {
            invariant_names.push_back(invariant.name());
        }
 
        internal::StateMachineHandle machine_handle(
            tc, rule_names, rule_groups, invariant_names, min_concurrency,
            max_concurrency);
        tc.note("Concurrency level: " +
                std::to_string(machine_handle.concurrency()));
        tc.note("Initial invariant check.");
        check_invariants(tc, "in the initial state", machine, invariants);
 
        enum class EventKind { RoundDone, Invalid, Overrun, Control, Panicked };
        struct Event {
            EventKind kind = EventKind::RoundDone;
            std::optional<internal::CapturedException> exception;
        };
        struct RoundState {
            std::mutex mutex;
            std::condition_variable start;
            std::condition_variable done;
            uint64_t round_number = 0;
            bool stop = false;
            std::size_t num_workers_completed = 0;
            std::vector<Event> events;
        };
 
        auto classify = [] {
            Event event;
            event.exception = internal::capture_current_exception();
            try {
                std::rethrow_exception(event.exception->exception);
            } catch (const internal::HegelStopTest&) {
                event.kind = EventKind::Overrun;
            } catch (const internal::HegelReject&) {
                // not the same as rule rejection - entire test case is rejected
                // requires some worker race on cloned test cases with same
                // parent
                event.kind = EventKind::Invalid; // GCOVR_EXCL_LINE
            } catch (const std::invalid_argument&) {
                event.kind = EventKind::Control;
            } catch (...) {
                event.kind = EventKind::Panicked;
            }
            return event;
        };
 
        auto run_round = [&](std::size_t worker, TestCase& worker_tc) {
            try {
                while (true) {
                    int64_t rule_index = machine_handle.next_rule(
                        worker_tc, static_cast<int64_t>(worker));
                    if (rule_index == internal::state_machine_done) {
                        return Event{};
                    }
                    if (rule_index < 0 ||
                        static_cast<std::size_t>(rule_index) >= rules.size()) {
                        // GCOVR_EXCL_START
                        throw std::runtime_error(
                            "state_machine_next_rule returned out-of-range "
                            "rule index. Please report this as a bug.");
                        // GCOVR_EXCL_STOP
                    }
 
                    const ConcurrentRule<M>& rule =
                        rules[static_cast<std::size_t>(rule_index)];
                    worker_tc.note("Rule: " + rule.name());
                    try {
                        internal::NoteIndentScope indent(worker_tc);
                        rule.function()(worker_tc, machine);
                    } catch (const internal::HegelReject&) {
                        machine_handle.rule_rejected(
                            worker_tc, static_cast<int64_t>(worker));
                        worker_tc.note(
                            "Rule stopped early due to violated assumption.");
                    }
                }
            } catch (...) {
                return classify();
            }
        };
 
        std::size_t worker_count =
            static_cast<std::size_t>(machine_handle.concurrency());
        RoundState state;
        state.events.resize(worker_count);
        std::vector<TestCase> worker_cases;
        worker_cases.reserve(worker_count);
        for (std::size_t worker = 0; worker < worker_count; ++worker) {
            worker_cases.push_back(tc.clone());
            internal::set_concurrent_worker(worker_cases.back(), worker);
        }
        std::vector<std::thread> workers;
        workers.reserve(worker_count);
        for (std::size_t worker = 0; worker < worker_count; ++worker) {
            workers.emplace_back(
                [&, worker,
                 worker_tc = std::move(worker_cases[worker])]() mutable {
                    uint64_t observed_round_number = 0;
                    while (true) {
                        {
                            std::unique_lock<std::mutex> lock(state.mutex);
                            state.start.wait(lock, [&] {
                                return state.stop ||
                                       state.round_number !=
                                           observed_round_number; // main thread
                                                                  // started a
                                                                  // new round
                            });
                            if (state.stop) {
                                return;
                            }
                            observed_round_number = state.round_number;
                        }
 
                        {
                            Event event = run_round(worker, worker_tc);
                            std::lock_guard<std::mutex> lock(state.mutex);
                            state.events[worker] = std::move(event);
                            ++state.num_workers_completed;
                        }
                        state.done.notify_one();
                    }
                });
        }
 
        auto stop_workers = [&] {
            {
                std::lock_guard<std::mutex> lock(state.mutex);
                state.stop = true;
            }
            state.start.notify_all();
            for (std::thread& worker : workers) {
                worker.join();
            }
        };
 
        try {
            while (true) {
                int64_t group = machine_handle.next_group(tc);
                if (group == internal::state_machine_done) {
                    break;
                }
                if (group < 0 ||
                    static_cast<std::size_t>(group) >= group_names.size()) {
                    // GCOVR_EXCL_START
                    throw std::runtime_error(
                        "state_machine_next_group returned an unknown group "
                        "index. Please report this as a bug.");
                    // GCOVR_EXCL_STOP
                }
                {
                    std::lock_guard<std::mutex> lock(state.mutex);
                    state.num_workers_completed = 0;
                    ++state.round_number;
                    tc.note("---------------- Round " +
                            std::to_string(state.round_number) + ": group \"" +
                            group_names[static_cast<std::size_t>(group)] +
                            "\" ----------------");
                }
                state.start.notify_all();
                {
                    std::unique_lock<std::mutex> lock(state.mutex);
                    state.done.wait(lock, [&] {
                        return state.num_workers_completed == worker_count;
                    });
                }
 
                std::optional<std::size_t> control;
                bool saw_overrun = false;
                bool saw_invalid = false;
                std::optional<std::size_t> panic;
                // need to collect every event because the precedence is
                // control > overrun > invalid > panic > round complete
                for (std::size_t worker = 0; worker < worker_count; ++worker) {
                    switch (state.events[worker].kind) {
                    case EventKind::RoundDone:
                        break;
                        // GCOVR_EXCL_START
                    case EventKind::Invalid:
                        saw_invalid = true;
                        break;
                        // GCOVR_EXCL_STOP
                    case EventKind::Control:
                        if (!control.has_value()) {
                            control = worker;
                        }
                        break;
                    case EventKind::Overrun:
                        saw_overrun = true;
                        break;
                    case EventKind::Panicked:
                        if (!panic.has_value()) {
                            panic = worker;
                        }
                        break;
                    }
                }
                if (control.has_value()) {
                    state.events[*control].exception->rethrow();
                }
                if (saw_overrun) {
                    throw internal::HegelStopTest();
                }
                if (saw_invalid) {
                    throw internal::HegelReject(); // GCOVR_EXCL_LINE
                }
                if (panic.has_value()) {
                    state.events[*panic].exception->rethrow();
                }
                check_invariants(
                    tc, "after round " + std::to_string(state.round_number),
                    machine, invariants, std::ref(machine_handle));
            }
 
            tc.note("Final invariant check.");
            check_invariants(tc, "in the final state", machine, invariants);
        } catch (...) {
            stop_workers();
            throw;
        }
        stop_workers();
    }

    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());
 
        std::vector<int64_t> rule_groups(rule_names.size(), 0);
        internal::StateMachineHandle machine_handle(tc, rule_names, rule_groups,
                                                    invariant_names, 1, 1);
        tc.note("Initial invariant check.");
        print_state(tc, machine, params);
        check_invariants(tc, "in the initial state", machine, invariants);
        int64_t steps_run = 0;
 
        // Sequential stateful testing is modeled as a concurrent stateful test
        // with one group and one worker.
        while (machine_handle.next_group(tc) != internal::state_machine_done) {
            while (true) {
                int64_t next_rule_idx = machine_handle.next_rule(tc, 0);
                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
                internal::start_span(tc, internal::SpanLabel::StatefulRule);
                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);
                    internal::stop_span(tc);
                } catch (const internal::HegelReject&) {
                    tc.note("Rule stopped early due to violated assumption.");
                    machine_handle.rule_rejected(tc, 0);
                    internal::stop_span(tc, true);
                } catch (...) {
                    internal::stop_span(tc);
                    throw;
                }
            }
            check_invariants(tc, "after step " + std::to_string(steps_run),
                             machine, invariants, std::ref(machine_handle));
        }
 
        tc.note("Final invariant check.");
        check_invariants(tc, "in the final state", machine, invariants);
    }

 
} // namespace hegel::stateful