core.h

#pragma once
 
#include <functional>
#include <memory>
#include <optional>
#include <stdexcept>
#include <type_traits>
 
#include "config.h"
#include "internal.h"
#include "repr.h"
#include "test_case.h"

namespace hegel::generators {
 
    template <typename T> class CompositeGenerator;
    template <typename T, typename U> class MappedGenerator;

    template <typename T> struct IGenerator {
        IGenerator() {}
        virtual ~IGenerator() = default;

        // Every generator produces its value by driving the engine's typed
        // draw primitives (hegel::internal::draw_*) against the test case.
        virtual T do_draw(const TestCase& tc) const = 0;
    };

    template <typename T> class Generator {
      public:
        using value_type = T;

        Generator(IGenerator<T>* p) : inner_(p) {}
        Generator(std::shared_ptr<IGenerator<T>> p) : inner_(std::move(p)) {}
 
        T do_draw(const TestCase& tc) const { return inner_->do_draw(tc); }

        template <typename F>
        Generator<std::invoke_result_t<F, T>> map(F&& f) const {
            using ResultType = std::invoke_result_t<F, T>;
            return Generator<ResultType>(
                new MappedGenerator<T, ResultType>(inner_, std::forward<F>(f)));
        }

        template <typename F> std::invoke_result_t<F, T> flat_map(F&& f) const {
            // F: T -> Generator<ResultType>; flat_map returns the same
            // Generator<ResultType>.
            using ResultType = typename std::invoke_result_t<F, T>::value_type;
            auto inner = inner_;
            return compose([inner, f = std::forward<F>(f)](
                               const TestCase& tc) -> ResultType {
                internal::start_span(tc, internal::SpanLabel::FlatMap);
                ResultType result = f(inner->do_draw(tc)).do_draw(tc);
                internal::stop_span(tc);
                return result;
            });
        }

        Generator<T> filter(std::function<bool(const T&)> pred) const {
            auto inner = inner_;
            return compose([inner, pred](const TestCase& tc) -> T {
                for (int i = 0; i < 3; ++i) {
                    internal::start_span(tc, internal::SpanLabel::Filter);
                    T value = inner->do_draw(tc);
                    if (pred(value)) {
                        internal::stop_span(tc);
                        return value;
                    }
                    // Discard the rejected span so the engine retries from
                    // before it opened.
                    internal::stop_span(tc, true);
                }
                tc.assume(false);
                // unreachable: assume(false) throws
                throw internal::HegelReject();
            });
        }
 
      private:
        std::shared_ptr<IGenerator<T>> inner_;
    };

    // Generator that produces values by invoking a user-provided function.
    // Users never construct or reference this type directly; it's produced
    // internally by compose() and by map()/flat_map()/filter().
    template <typename T> class CompositeGenerator : public IGenerator<T> {
      public:
        explicit CompositeGenerator(std::function<T(const TestCase&)> fn)
            : gen_fn_(std::move(fn)) {}
 
        T do_draw(const TestCase& tc) const override { return gen_fn_(tc); }
 
      private:
        std::function<T(const TestCase&)> gen_fn_;
    };

    // Generator that applies a client-side transformation to values drawn
    // from a source generator. Produced internally by Generator<T>::map().
    template <typename T, typename U>
    class MappedGenerator : public IGenerator<U> {
      public:
        MappedGenerator(std::shared_ptr<IGenerator<T>> source,
                        std::function<U(T)> f)
            : source_(std::move(source)), f_(std::move(f)) {}
 
        U do_draw(const TestCase& tc) const override {
            internal::start_span(tc, internal::SpanLabel::Mapped);
            U result = f_(source_->do_draw(tc));
            internal::stop_span(tc);
            return result;
        }
 
      private:
        std::shared_ptr<IGenerator<T>> source_;
        std::function<U(T)> f_;
    };

    template <typename F> auto compose(F&& fn) {
        using T = std::invoke_result_t<F, const TestCase&>;
        return Generator<T>(new CompositeGenerator<T>(
            std::function<T(const TestCase&)>(std::forward<F>(fn))));
    }
 
} // namespace hegel::generators
 
namespace hegel {
 
    template <typename T>
    T TestCase::draw(const generators::Generator<T>& gen) const {
        return draw(std::string_view(), gen);
    }
 
    template <typename T>
    T TestCase::draw(std::string_view name, const generators::Generator<T>& gen,
                     bool repeatable) const {
        internal::DrawLogScope scope(*this, name, repeatable);
        T value = gen.do_draw(*this);
        if (scope.should_log()) {
            scope.log(internal::repr(value));
        }
        return value;
    }
 
} // namespace hegel