control.hh

#pragma once
 
#include <clingo/backend.hh>
#include <clingo/base.hh>
#include <clingo/config.hh>
#include <clingo/core.hh>
#include <clingo/ground.hh>
#include <clingo/observe.hh>
#include <clingo/profile.hh>
#include <clingo/propagate.hh>
#include <clingo/solve.hh>
#include <clingo/stats.hh>
#include <clingo/symbol.hh>
 
#include <clingo/control.h>
 
#include <cassert>
#include <optional>
#include <span>
 
namespace Clingo {
 
namespace AST {
 
class Program;
auto c_cast(Program const &x) -> clingo_program_t *;
 
} // namespace AST
 
namespace Detail {
 
void join(clingo_control_t *ctx, clingo_program_t const *prg);
 
} // namespace Detail
 
 
class ConstMap {
  public:
    using key_type = std::string_view;
    using mapped_type = Symbol;
    using value_type = std::pair<key_type, mapped_type>;
    using size_type = std::size_t;
    using difference_type = std::ptrdiff_t;
    using reference = value_type;
    using pointer = Detail::ArrowProxy<value_type>;
    using iterator = Detail::RandomAccessIterator<ConstMap>;
 
    explicit ConstMap(clingo_const_map_t const *map) : map_{map} {}
 
    [[nodiscard]] auto contains(key_type name) const -> bool {
        return Detail::call<clingo_const_map_find>(map_, name.data(), name.size(), nullptr);
    }
 
    [[nodiscard]] auto operator[](key_type name) const -> mapped_type {
        clingo_symbol_t sym = 0;
        bool found = false;
        Detail::handle_error(clingo_const_map_find(map_, name.data(), name.size(), &sym, &found));
        return found ? Symbol{sym, true} : throw std::out_of_range{"key not found"};
    }
 
    [[nodiscard]] auto at(size_t index) const -> value_type {
        clingo_string_t name;
        clingo_symbol_t sym = 0;
        Detail::handle_error(clingo_const_map_at(map_, index, &name, &sym));
        return {{name.data, name.size}, Symbol{sym, true}};
    }
 
    [[nodiscard]] auto size() const -> size_type { return Detail::call<clingo_const_map_size>(map_); }
 
    [[nodiscard]] auto begin() const -> iterator { return iterator{*this, 0}; }
 
    [[nodiscard]] auto end() const -> iterator { return iterator{*this, size()}; }
 
  private:
    clingo_const_map_t const *map_;
};
 
enum class ControlMode : clingo_mode_t {
    parse = clingo_mode_parse,     
    rewrite = clingo_mode_rewrite, 
    ground = clingo_mode_ground,   
    solve = clingo_mode_solve,     
};
 
enum class WriteAspifFlags : clingo_write_aspif_mode_t {
    none = 0,                                              
    preamble = clingo_write_aspif_mode_preamble,           
    preamble_auto = clingo_write_aspif_mode_preamble_auto, 
    append = clingo_write_aspif_mode_append,               
    preprocess = clingo_write_aspif_mode_preprocess,       
    symbols = clingo_write_aspif_mode_symbols,             
};
CLINGO_ENABLE_BITSET_ENUM(WriteAspifFlags);
 
enum class SolveFlags : clingo_solve_mode_bitset_t {
    empty = 0,                       
    yield = clingo_solve_mode_yield, 
    async = clingo_solve_mode_async, 
};
CLINGO_ENABLE_BITSET_ENUM(SolveFlags);
 
enum class DiscardType {
    minimize = clingo_discard_type_e::minimize, 
    project = clingo_discard_type_e::project,   
};
 
class Control {
  public:
    using Context = std::function<SymbolVector(std::string_view, SymbolSpan)>;
 
    explicit Control(Library const &lib, StringList arguments) : Clingo::Control{lib, StringSpan{arguments}} {}
 
    explicit Control(Library const &lib, StringSpan arguments = {}) {
        auto cstrs = Detail::transform(arguments, [](auto const &x) { return clingo_string_t{x.data(), x.size()}; });
        ctl_.reset(Detail::call<clingo_control_new>(c_cast(lib), cstrs.data(), arguments.size()), false);
    }
 
    explicit Control(clingo_control_t *rep, bool acquire) : ctl_{rep, acquire} {}
 
    [[nodiscard]] friend auto c_cast(Control const &ctl) -> clingo_control_t * { return ctl.ctl_.get(); }
 
    [[nodiscard]] auto mode() const -> ControlMode {
        return static_cast<ControlMode>(Detail::call<clingo_control_mode>(ctl_.get()));
    }
 
    void write_aspif(std::string_view path, WriteAspifFlags flags = WriteAspifFlags::none) const {
        Detail::handle_error(clingo_control_write_aspif(ctl_.get(), path.data(), path.size(),
                                                        static_cast<clingo_write_aspif_mode_t>(flags)));
    }
 
    void parse_files(StringSpan files) const {
        auto cfiles = Detail::transform(files, [](auto const &x) { return clingo_string_t{x.data(), x.size()}; });
        Detail::handle_error(clingo_control_parse_files(ctl_.get(), cfiles.data(), cfiles.size()));
    }
 
    void parse_files(StringList files) const { parse_files(StringSpan{files.begin(), files.end()}); }
 
    void parse_string(std::string_view program) const {
        Detail::handle_error(clingo_control_parse_string(ctl_.get(), program.data(), program.size()));
    }
 
    void ground(std::optional<PartSpan> parts = std::nullopt, Context ctx = nullptr) const {
        auto [default_parts, c_parts] = c_parts_(parts);
        constexpr static clingo_ground_event_handler_t handler = clingo_ground_event_handler_t{
            []([[maybe_unused]] char const *name, [[maybe_unused]] size_t name_size,
               [[maybe_unused]] size_t arguments_size, [[maybe_unused]] void *data, bool *result) -> bool {
                CLINGO_TRY {
                    *result = true;
                }
                CLINGO_CATCH;
            },
            []([[maybe_unused]] clingo_lib_t *lib, [[maybe_unused]] clingo_location_t const *location, char const *name,
               size_t name_size, clingo_symbol_t const *arguments, size_t arguments_size, void *data,
               clingo_symbol_callback_t symbol_callback, void *symbol_callback_data) -> bool {
                CLINGO_TRY {
                    auto &cb = *static_cast<std::function<SymbolVector(std::string_view, SymbolSpan)> *>(data);
                    auto syms = cb({name, name_size}, {cpp_cast(arguments), arguments_size});
                    auto const *c_syms = c_cast(syms.data());
                    return symbol_callback(c_syms, syms.size(), symbol_callback_data);
                }
                CLINGO_CATCH;
            },
            nullptr,
            nullptr,
        };
 
        Detail::handle_error(
            clingo_control_ground(ctl_.get(), c_parts.data(), c_parts.size(), ctx ? &handler : nullptr, &ctx));
    }
 
    void ground(std::initializer_list<Part> parts, Context ctx = nullptr) const {
        ground(PartSpan{parts}, std::move(ctx));
    }
 
    template <Detail::UserData<GroundEventHandler> Handler = std::nullptr_t>
    [[nodiscard]] auto start_ground(std::optional<PartSpan> parts = std::nullopt, Handler &&handler = nullptr) const
        -> GroundHandle {
        auto [default_parts, c_parts] = c_parts_(parts);
        using UserData = Detail::UserDataTraits<Handler>;
        auto user_data = UserData::create(std::forward<Handler>(handler));
        clingo_ground_event_handler_t const *c_handler_ptr = nullptr;
        if constexpr (UserData::has_data) {
            if (UserData::has_value(user_data)) {
                static constexpr auto c_handler = clingo_ground_event_handler_t{
                    [](char const *name, size_t name_size, size_t arguments_size, void *data, bool *result) -> bool {
                        CLINGO_TRY {
                            *result = UserData::cast(data)->callable(std::string_view{name, name_size}, arguments_size);
                        }
                        CLINGO_CATCH;
                    },
                    []([[maybe_unused]] clingo_lib_t *lib, clingo_location_t const *location, char const *name,
                       size_t name_size, clingo_symbol_t const *arguments, size_t arguments_size, void *data,
                       clingo_symbol_callback_t symbol_callback, void *symbol_callback_data) -> bool {
                        CLINGO_TRY {
                            Location loc{location};
                            auto syms = UserData::cast(data)->call(loc, {name, name_size},
                                                                   {cpp_cast(arguments), arguments_size});
                            auto const *c_syms = c_cast(syms.data());
                            return symbol_callback(c_syms, syms.size(), symbol_callback_data);
                        }
                        CLINGO_CATCH;
                    },
                    [](clingo_ground_result_t result, void *data) {
                        UserData::cast(data)->finish(static_cast<GroundResult>(result));
                    },
                    [](void *data) { UserData::free(data); },
                };
                c_handler_ptr = &c_handler;
            }
        }
        clingo_ground_handle_t *handle = nullptr;
        Detail::handle_error(clingo_control_start_ground(ctl_.get(), c_parts.data(), c_parts.size(), c_handler_ptr,
                                                         UserData::release(user_data), &handle));
        return GroundHandle{handle};
    }
 
    template <Detail::UserData<GroundEventHandler> Handler = std::nullptr_t>
    [[nodiscard]] auto start_ground(std::initializer_list<Part> parts, Handler &&handler = nullptr) const
        -> GroundHandle {
        return start_ground(PartSpan{parts}, std::forward<Handler>(handler));
    }
 
    [[nodiscard]] auto base() const -> Base { return Base{Detail::call<clingo_control_base>(ctl_.get())}; }
 
    [[nodiscard]] auto stats() const -> ConstStats {
        auto const *stats = Detail::call<clingo_control_stats>(ctl_.get());
        auto key = Detail::call<clingo_stats_root>(stats);
        return ConstStats{stats, key};
    }
 
    auto profile() -> std::vector<ProfileNode> {
        struct Builder {
            static auto internal(size_t depth, char const *key, size_t key_size, bool nested, void *data) -> bool {
                CLINGO_TRY {
                    auto *self = static_cast<Builder *>(data);
                    assert(depth <= self->stack.size());
                    self->stack.resize(depth);
                    auto node = ProfileNodeInternal{std::string{key, key_size}, nested};
                    if (self->stack.empty()) {
                        self->roots.emplace_back(std::move(node));
                        self->stack.emplace_back(&self->roots.back());
                    } else {
                        auto &children = std::get<ProfileNodeInternal>(*self->stack.back()).children;
                        children.emplace_back(std::move(node));
                        self->stack.emplace_back(&children.back());
                    }
                }
                CLINGO_CATCH;
            }
            static auto leaf(size_t depth, clingo_profile_data_t *values, clingo_profile_type_t type, void *data)
                -> bool {
                CLINGO_TRY {
                    auto *self = static_cast<Builder *>(data);
                    assert(depth <= self->stack.size());
                    self->stack.resize(depth);
                    auto node = ProfileNodeLeaf{static_cast<ProfileType>(type), values->matches, values->instances,
                                                values->time_instantiate, values->time_propagate};
                    if (self->stack.empty()) {
                        self->roots.emplace_back(node);
                    } else {
                        std::get<ProfileNodeInternal>(*self->stack.back()).children.emplace_back(node);
                    }
                }
                CLINGO_CATCH;
            }
 
            std::vector<ProfileNode *> stack;
            std::vector<ProfileNode> roots;
        } builder;
 
        auto visitor = clingo_profile_visitor_t{&Builder::internal, &Builder::leaf};
        clingo_control_profile(ctl_.get(), &visitor, &builder);
        return std::move(builder.roots);
    }
 
    template <Detail::UserData<SolveEventHandler> Handler = std::nullptr_t>
    [[nodiscard]] auto solve(ProgramLiteralSpan assumptions = {}, Handler &&handler = nullptr) const -> SolveResult {
        return solve_(std::forward<Handler>(handler), assumptions, SolveFlags::empty).get();
    }
 
    template <Detail::UserData<SolveEventHandler> Handler = std::nullptr_t>
    [[nodiscard]] auto solve(std::initializer_list<ProgramLiteral> assumptions, Handler &&handler = nullptr) const
        -> SolveResult {
        return solve_(std::forward<Handler>(handler), ProgramLiteralSpan{assumptions}, SolveFlags::empty).get();
    }
 
    template <Detail::UserData<SolveEventHandler> Handler = std::nullptr_t>
    [[nodiscard]] auto start_solve(ProgramLiteralSpan assumptions = {}, SolveFlags flags = SolveFlags::yield,
                                   Handler &&handler = nullptr) const -> SolveHandle {
        return solve_(std::forward<Handler>(handler), assumptions, flags);
    }
 
    template <Detail::UserData<SolveEventHandler> Handler = std::nullptr_t>
    [[nodiscard]] auto start_solve(std::initializer_list<ProgramLiteral> assumptions,
                                   SolveFlags flags = SolveFlags::yield, Handler &&handler = nullptr) const
        -> SolveHandle {
        return solve_(std::forward<Handler>(handler), ProgramLiteralSpan{assumptions}, flags);
    }
 
    void main() const { Detail::handle_error(clingo_control_main(ctl_.get())); }
 
    void interrupt() const { clingo_control_interrupt(ctl_.get()); }
 
    void discard(DiscardType type) const {
        Detail::handle_error(clingo_control_discard(ctl_.get(), static_cast<clingo_discard_type_t>(type)));
    }
 
    [[nodiscard]] auto buffer() const -> std::string_view {
        auto [data, size] = Detail::call<clingo_control_buffer>(ctl_.get());
        return {data, size};
    }
 
    [[nodiscard]] auto parts() const -> std::optional<PartVector> {
        clingo_part_t const *parts = nullptr;
        size_t size = 0;
        bool has_value = false;
        Detail::handle_error(clingo_control_get_parts(ctl_.get(), &parts, &size, &has_value));
        if (has_value) {
            return Detail::transform(std::span{parts, size}, [](auto const &part) {
                auto params = std::span{cpp_cast(part.params), part.params_size};
                return Part{{part.name, part.name_size}, {params.begin(), params.end()}};
            });
        }
        return std::nullopt;
    }
 
    void parts(PartList parts) const { this->parts(std::span{parts.begin(), parts.end()}); }
 
    void parts(std::optional<PartSpan> parts) const {
        if (parts) {
            auto cparts = Detail::transform(*parts, [](auto const &part) {
                return clingo_part_t{part.name.data(), part.name.size(), c_cast(part.params.data()),
                                     part.params.size()};
            });
            Detail::handle_error(clingo_control_set_parts(ctl_.get(), cparts.data(), cparts.size(), true));
        } else {
            Detail::handle_error(clingo_control_set_parts(ctl_.get(), nullptr, 0, true));
        }
    }
 
    [[nodiscard]] auto config() const -> Config {
        auto *config = Detail::call<clingo_control_config>(ctl_.get());
        auto key = Detail::call<clingo_config_root>(config);
        return Config{config, key};
    }
 
    [[nodiscard]] auto const_map() const -> ConstMap {
        return ConstMap{Detail::call<clingo_control_const_map>(ctl_.get())};
    }
 
    void observe(Observer &obs, bool preprocess = true) const { obs.observe(ctl_.get(), preprocess); }
 
    [[nodiscard]] auto backend() const -> ProgramBackend {
        return ProgramBackend{Detail::call<clingo_control_backend>(ctl_.get())};
    }
 
    template <Detail::UserData<Propagator, false> P> auto register_propagator(P &&propagator) const -> decltype(auto) {
        using UserData = Detail::UserDataTraits<P>;
        auto user_data = UserData::create(std::forward<P>(propagator));
        if (!UserData::has_value(user_data)) {
            throw std::invalid_argument{"propagator cannot be null"};
        }
        auto &res = *UserData::get(user_data);
 
        static constexpr auto decide = []() {
            if constexpr (std::is_base_of_v<Heuristic, typename UserData::ValueType>) {
                return [](clingo_assignment_t const *assignment, clingo_literal_t fallback, void *data,
                          clingo_literal_t *decision) -> bool {
                    CLINGO_TRY {
                        *decision = UserData::cast(data)->decide(Assignment{assignment}, fallback);
                    }
                    CLINGO_CATCH;
                };
            } else {
                return nullptr;
            }
        };
 
        static constexpr auto c_propagator = clingo_propagator_t{
            [](clingo_assignment_t const *assignment, clingo_propagate_init_t *init, void *data) -> bool {
                CLINGO_TRY {
                    UserData::cast(data)->init(Assignment{assignment}, PropagateInit{init});
                }
                CLINGO_CATCH;
            },
            [](clingo_assignment_t const *assignment, clingo_propagate_control_t *control, void *data) {
                CLINGO_TRY {
                    UserData::cast(data)->attach(Assignment{assignment}, PropagateControl{control});
                }
                CLINGO_CATCH;
            },
            [](clingo_assignment_t const *assignment, clingo_propagate_control_t *control,
               clingo_literal_t const *changes, size_t size, void *data) -> bool {
                CLINGO_TRY {
                    UserData::cast(data)->propagate(Assignment{assignment}, PropagateControl{control},
                                                    SolverLiteralSpan{changes, size});
                }
                CLINGO_CATCH;
            },
            [](clingo_assignment_t const *assignment, clingo_literal_t const *changes, size_t size, void *data) {
                try {
                    UserData::cast(data)->undo(Assignment{assignment}, SolverLiteralSpan{changes, size});
                } catch (std::exception const &e) {
                    fprintf(stderr, "panic: %s\n", e.what());
                    std::terminate();
                }
            },
            [](clingo_assignment_t const *assignment, clingo_propagate_control_t *control, void *data) -> bool {
                CLINGO_TRY {
                    UserData::cast(data)->check(Assignment{assignment}, PropagateControl{control});
                }
                CLINGO_CATCH;
            },
            decide(),
            [](void *data) { UserData::free(data); },
        };
 
        Detail::handle_error(
            clingo_control_register_propagator(ctl_.get(), &c_propagator, UserData::release(user_data)));
        return res;
    }
 
    void join(AST::Program const &prg) const { Detail::join(ctl_.get(), c_cast(prg)); }
 
  private:
    friend class Detail::intrusive_handle<Control, clingo_control_t>;
 
    static auto acquire(clingo_control_t *ptr) { clingo_control_acquire(ptr); }
 
    static auto release(clingo_control_t *ptr) { clingo_control_release(ptr); }
 
    template <Detail::UserData<SolveEventHandler> Handler = std::nullptr_t>
    [[nodiscard]] auto solve_(Handler &&handler, ProgramLiteralSpan const &assumptions, SolveFlags flags) const
        -> SolveHandle {
        using UserData = Detail::UserDataTraits<Handler>;
        auto user_data = UserData::create(std::forward<Handler>(handler));
        clingo_solve_event_handler_t const *c_handler_ptr = nullptr;
        if constexpr (UserData::has_data) {
            if (UserData::has_value(user_data)) {
                static constexpr auto c_solve_event_handler = clingo_solve_event_handler_t{
                    [](clingo_model_t *model, void *data, bool *goon) -> bool {
                        CLINGO_TRY {
                            auto *hnd = UserData::cast(data);
                            assert(hnd != nullptr);
                            auto mdl = Model{model};
                            *goon = hnd->model(mdl);
                            return true;
                        }
                        CLINGO_CATCH;
                    },
                    [](int64_t const *values, size_t size, void *data) -> bool {
                        CLINGO_TRY {
                            auto *hnd = UserData::cast(data);
                            assert(hnd != nullptr);
                            hnd->unsat({values, size});
                        }
                        CLINGO_CATCH;
                    },
                    [](clingo_stats_t *stats, void *data) -> bool {
                        CLINGO_TRY {
                            auto *hnd = UserData::cast(data);
                            assert(hnd != nullptr);
                            std::string_view user_step = "user_step";
                            std::string_view user_accu = "user_accu";
                            uint64_t root = Detail::call<clingo_stats_root>(stats);
                            uint64_t step = Detail::call<clingo_stats_map_add_subkey>(
                                stats, root, user_step.data(), user_step.size(), clingo_stats_type_map);
                            uint64_t accu = Detail::call<clingo_stats_map_add_subkey>(
                                stats, root, user_accu.data(), user_accu.size(), clingo_stats_type_map);
                            hnd->stats(Stats{stats, step}, Stats{stats, accu});
                        }
                        CLINGO_CATCH;
                    },
                    [](clingo_solve_result_bitset_t result, void *data) -> void {
                        auto *hnd = UserData::cast(data);
                        assert(hnd != nullptr);
                        hnd->finish(static_cast<SolveResult>(result));
                    },
                    [](void *data) { UserData::free(data); },
                };
                c_handler_ptr = &c_solve_event_handler;
            }
        }
 
        clingo_solve_handle_t *res = nullptr;
        Detail::handle_error(clingo_control_solve(ctl_.get(), static_cast<clingo_solve_mode_bitset_t>(flags),
                                                  assumptions.data(), assumptions.size(), c_handler_ptr,
                                                  UserData::release(user_data), &res));
        return SolveHandle{res};
    }
 
    [[nodiscard]] auto c_parts_(std::optional<PartSpan> parts = std::nullopt) const
        -> std::pair<std::optional<PartVector>, std::vector<clingo_part_t>> {
        auto res = std::pair<std::optional<PartVector>, std::vector<clingo_part_t>>{};
        auto &[default_parts, c_parts] = res;
        if (!parts) {
            parts = default_parts = this->parts();
        }
        if (!parts) {
            static auto default_part = Part{"base", {}};
            parts = PartSpan{&default_part, 1};
        }
        assert(parts);
        c_parts.reserve(parts->size());
        for (auto const &part : *parts) {
            c_parts.emplace_back(part.name.data(), part.name.size(), c_cast(part.params.data()), part.params.size());
        }
        return res;
    }
 
    Detail::intrusive_handle<Control, clingo_control_t> ctl_;
};
 
 
} // namespace Clingo