ImplVariant.hpp

Go to the documentation of this file.

// -------------------------------------------------------//
//
// SHAMROCK code for hydrodynamics
// Copyright (c) 2021-2026 Timothée David--Cléris <tim.shamrock@proton.me>
// SPDX-License-Identifier: CeCILL Free Software License Agreement v2.1
// Shamrock is licensed under the CeCILL 2.1 License, see LICENSE for more information
//
// -------------------------------------------------------//
 
#pragma once

 
#include "shambase/exception.hpp"
#include "sham/format/format.hpp"
#include "shambackends/DeviceScheduler.hpp"
#include <fmt/ranges.h>
#include <nlohmann/json.hpp>
#include <string_view>
#include <concepts>
#include <functional>
#include <optional>
#include <string>
#include <utility>
#include <variant>
#include <vector>
 
namespace shamalgs {

    template<class Alt>
    struct ImplVariantParams {
        static inline nlohmann::json to_json(const Alt &) { return nlohmann::json::object(); }
        static inline Alt from_json(const nlohmann::json &) { return Alt{}; }
    };

    template<class Alt>
    concept HasCustomDefaults = requires {
        { Alt::variant_custom_defaults() } -> std::convertible_to<std::vector<Alt>>;
    };
 
    // Forward declaration: defined below, needed by impl_variant_alts::default_config_list()
    template<class Variant>
    inline std::string variant_to_config_string(const Variant &v);
 
    namespace details {
        template<class Alt>
        inline std::vector<Alt> alt_default_list() {
            if constexpr (HasCustomDefaults<Alt>) {
                return Alt::variant_custom_defaults();
            } else {
                return {Alt{}};
            }
        }

        template<class Variant>
        struct impl_variant_alts;
 
        template<class... Alts>
        struct impl_variant_alts<std::variant<Alts...>> {
            static inline std::vector<std::string> default_type_names() {
                return {std::string(Alts::variant_type_name)...};
            }

            static inline std::vector<std::string> default_config_list() {
                std::vector<std::string> out;
                auto add_alt = [&]<class Alt>() {
                    for (auto &alt : alt_default_list<Alt>()) {
                        out.push_back(
                            variant_to_config_string<std::variant<Alts...>>(
                                std::variant<Alts...>{std::move(alt)}));
                    }
                };
                (add_alt.template operator()<Alts>(), ...);
                return out;
            }
 
            static inline std::variant<Alts...> from_config_string(std::string_view s) {
                nlohmann::json j      = nlohmann::json::parse(s);
                std::string name      = j.at("implementation").get<std::string>();
                nlohmann::json params = j.value("parameters", nlohmann::json::object());
 
                std::optional<std::variant<Alts...>> result;
                (void) ((name == Alts::variant_type_name
                             ? (result
                                = std::variant<Alts...>{ImplVariantParams<Alts>::from_json(params)},
                                true)
                             : false)
                        || ...);
 
                if (!result) {
                    throw shambase::make_except_with_loc<std::invalid_argument>(sham::format(
                        "invalid implementation : {}, possible implementations : {}",
                        name,
                        default_type_names()));
                }
                return *result;
            }
        };
    } // namespace details

    template<class Variant>
    inline std::string variant_to_config_string(const Variant &v) {
        return std::visit(
            [](const auto &alt) {
                using Alt = std::decay_t<decltype(alt)>;
                nlohmann::json j;
                j["implementation"] = std::string(Alt::variant_type_name);
                j["parameters"]     = ImplVariantParams<Alt>::to_json(alt);
                return j.dump();
            },
            v);
    }

    template<class Variant>
    inline Variant variant_from_config_string(std::string_view s) {
        return details::impl_variant_alts<Variant>::from_config_string(s);
    }

    template<class Variant>
    inline std::vector<std::string> variant_default_type_names() {
        return details::impl_variant_alts<Variant>::default_type_names();
    }

    class IImplVariant {
        public:
        virtual ~IImplVariant() = default;

        virtual std::string get_current_config() const = 0;

        virtual std::vector<std::string> get_default_config_list() const = 0;

        virtual void set(std::string_view config_json) = 0;

        virtual bool is_set() const = 0;

        virtual void autoselect(const sham::DeviceScheduler_ptr &sched) = 0;
    };

    template<class... Alts>
    class ImplVariantGlobal : public IImplVariant {
        public:
        using Variant = std::variant<Alts...>;

        using AutoselectFn
            = std::function<void(const sham::DeviceScheduler_ptr &, ImplVariantGlobal &)>;

        explicit ImplVariantGlobal(AutoselectFn fn) : autoselect_fn(std::move(fn)) {}

        inline bool is_set() const override { return current.has_value(); }

        /// Select the algorithm's default implementation for the device behind `sched`
        inline void autoselect(const sham::DeviceScheduler_ptr &sched) override {
            autoselect_fn(sched, *this);
        }

        inline const Variant &get() const { return *current; }

        /// null if no implementation has been selected yet (see is_set())
        inline std::string get_current_config() const override {
            if (!is_set()) {
                return nlohmann::json(nullptr).dump();
            }
            return variant_to_config_string(get());
        }

        /// contribute one entry per instance in their variant_custom_defaults() list.
        inline std::vector<std::string> get_default_config_list() const override {
            return details::impl_variant_alts<Variant>::default_config_list();
        }

        inline void set(Variant v) { current = std::move(v); }

        /// Select an implementation from a {"implementation": ..., "parameters": ...} json string
        inline void set(std::string_view config_json) override {
            current = variant_from_config_string<Variant>(config_json);
        }
 
        private:
        std::optional<Variant> current;
        AutoselectFn autoselect_fn;
    };
 
} // namespace shamalgs