combinators.h

#pragma once
 
#include <cstdint>
#include <initializer_list>
#include <memory>
#include <optional>
#include <stdexcept>
#include <string>
#include <tuple>
#include <variant>
#include <vector>
 
#include "hegel/core.h"
 
namespace hegel::generators {

    // Concrete IGenerator for sampled_from(). Draws an index into the
    // captured `elements_` vector and returns that element.
    template <typename T> class SampledFromGenerator : public IGenerator<T> {
      public:
        explicit SampledFromGenerator(std::vector<T> elements)
            : elements_(std::move(elements)) {
            if (elements_.empty()) {
                throw std::invalid_argument(
                    "sampled_from requires a non-empty vector");
            }
        }
 
        T do_draw(const TestCase& tc) const override {
            int64_t index = hegel::internal::draw_integer(
                tc, 0, static_cast<int64_t>(elements_.size() - 1));
            return elements_[static_cast<size_t>(index)];
        }
 
      private:
        std::vector<T> elements_;
    };
 
    // Concrete IGenerator for one_of(). Draws a branch index inside a
    // one_of span, then delegates to that branch.
    template <typename T> class OneOfGenerator : public IGenerator<T> {
      public:
        explicit OneOfGenerator(std::vector<Generator<T>> gens)
            : gens_(std::move(gens)) {
            if (gens_.empty()) {
                throw std::invalid_argument(
                    "one_of requires a non-empty vector of generators");
            }
        }
 
        T do_draw(const TestCase& tc) const override {
            namespace hi = hegel::internal;
            hi::start_span(tc, hi::SpanLabel::OneOf);
            int64_t index =
                hi::draw_integer(tc, 0, static_cast<int64_t>(gens_.size() - 1));
            T result = gens_[static_cast<size_t>(index)].do_draw(tc);
            hi::stop_span(tc);
            return result;
        }
 
      private:
        std::vector<Generator<T>> gens_;
    };


    template <typename T>
    Generator<T> sampled_from(const std::vector<T>& elements) {
        return Generator<T>(new SampledFromGenerator<T>(elements));
    }

    template <typename T>
    Generator<T> sampled_from(std::initializer_list<T> elements) {
        return sampled_from(std::vector<T>(elements));
    }

    inline Generator<std::string>
    sampled_from(std::initializer_list<const char*> elements) {
        std::vector<std::string> strings;
        strings.reserve(elements.size());
        for (const char* s : elements) {
            strings.push_back(s);
        }
        return sampled_from(strings);
    }



    template <typename T> Generator<T> one_of(std::vector<Generator<T>> gens) {
        return Generator<T>(new OneOfGenerator<T>(std::move(gens)));
    }

    template <typename T>
    Generator<T> one_of(std::initializer_list<Generator<T>> gens) {
        return one_of(std::vector<Generator<T>>(gens));
    }

    namespace detail {
 
        template <typename Variant, typename GenTuple, size_t I = 0>
        Variant draw_variant_impl(const GenTuple& gens, size_t idx,
                                  const TestCase& tc) {
            if constexpr (I < std::tuple_size_v<GenTuple>) {
                if (idx == I) {
                    return Variant{std::in_place_index<I>,
                                   std::get<I>(gens).do_draw(tc)};
                }
                return draw_variant_impl<Variant, GenTuple, I + 1>(gens, idx,
                                                                   tc);
            } else {
                // Unreachable: idx is always in [0, N), so an earlier branch
                // matches before the recursion bottoms out.
                return Variant{}; // GCOVR_EXCL_LINE
            }
        }
 
    } // namespace detail
 
    // Concrete IGenerator for variant(). Draws a branch index inside a
    // one_of span; branches can have heterogeneous types.
    template <typename... Ts>
    class VariantGenerator : public IGenerator<std::variant<Ts...>> {
      public:
        using Result = std::variant<Ts...>;
 
        explicit VariantGenerator(Generator<Ts>... gens)
            : gens_(std::move(gens)...) {}
 
        Result do_draw(const TestCase& tc) const override {
            namespace hi = hegel::internal;
            constexpr size_t N = sizeof...(Ts);
            hi::start_span(tc, hi::SpanLabel::OneOf);
            int64_t index =
                hi::draw_integer(tc, 0, static_cast<int64_t>(N - 1));
            Result result = detail::draw_variant_impl<Result, decltype(gens_)>(
                gens_, static_cast<size_t>(index), tc);
            hi::stop_span(tc);
            return result;
        }
 
      private:
        std::tuple<Generator<Ts>...> gens_;
    };
 
    // Concrete IGenerator for optional(). A boolean draw inside an
    // optional span gates presence; false shrinks toward nullopt.
    template <typename T>
    class OptionalGenerator : public IGenerator<std::optional<T>> {
      public:
        explicit OptionalGenerator(Generator<T> gen) : gen_(std::move(gen)) {}
 
        std::optional<T> do_draw(const TestCase& tc) const override {
            namespace hi = hegel::internal;
            hi::start_span(tc, hi::SpanLabel::Optional);
            std::optional<T> result;
            if (hi::draw_boolean(tc, 0.5)) {
                result = gen_.do_draw(tc);
            }
            hi::stop_span(tc);
            return result;
        }
 
      private:
        Generator<T> gen_;
    };

    template <typename... Ts>
    Generator<std::variant<Ts...>> variant(Generator<Ts>... gens) {
        return Generator<std::variant<Ts...>>(
            new VariantGenerator<Ts...>(std::move(gens)...));
    }

    template <typename T>
    Generator<std::optional<T>> optional(Generator<T> gen) {
        return Generator<std::optional<T>>(
            new OptionalGenerator<T>(std::move(gen)));
    }


    // Generator that delegates to an implementation supplied later through a
    // shared slot. Handed out by DeferredGeneratorDefinition::generator().
    template <typename T> class DeferredGenerator : public IGenerator<T> {
      public:
        explicit DeferredGenerator(
            std::shared_ptr<std::optional<Generator<T>>> slot)
            : slot_(std::move(slot)) {}
 
        T do_draw(const TestCase& tc) const override {
            if (!*slot_) {
                throw std::runtime_error(
                    "deferred generator drawn before set() was called");
            }
            return (*slot_)->do_draw(tc);
        }
 
      private:
        std::shared_ptr<std::optional<Generator<T>>> slot_;
    };

    template <typename T> class DeferredGeneratorDefinition {
      public:
        DeferredGeneratorDefinition()
            : slot_(std::make_shared<std::optional<Generator<T>>>()) {}

        Generator<T> generator() const {
            return Generator<T>(new DeferredGenerator<T>(slot_));
        }

        void set(Generator<T> gen) {
            if (*slot_) {
                throw std::runtime_error(
                    "deferred generator set() called more than once");
            }
            *slot_ = std::move(gen);
        }
 
      private:
        std::shared_ptr<std::optional<Generator<T>>> slot_;
    };

    template <typename T> DeferredGeneratorDefinition<T> deferred() {
        return DeferredGeneratorDefinition<T>();
    }
 
} // namespace hegel::generators