/// @file dpf/path_memoizer.hpp
/// @brief
/// @details
/// @author Ryan Henry <ryan.henry@ucalgary.ca>
/// @author Christopher Jiang <christopher.jiang@ucalgary.ca>
/// @copyright Copyright (c) 2019-2024 Ryan Henry and [others](@ref authors)
/// @license Released under a GNU General Public v2.0 (GPLv2) license;
///          see [LICENSE.md](@ref license) for details.

#ifndef LIBDPF_INCLUDE_DPF_PATH_MEMOIZER_HPP__
#define LIBDPF_INCLUDE_DPF_PATH_MEMOIZER_HPP__

#include "hedley/hedley.h"

#include <cstddef>
#include <cstring>
#include <algorithm>
#include <type_traits>
#include <functional>
#include <memory>
#include <array>
#include <optional>

#include "dpf/utils.hpp"
#include "dpf/secret_share.hpp"

namespace dpf
{

/// Path memoizers key on the underlying DPF key type. `party_key` wrappers
/// share the same tree layout, so a memoizer built for party 0 also accepts
/// party 1 (and bare keys).
template <typename T>
using path_memoizer_key_t = unwrap_party_key_t<T>;

template <typename DpfKey,
          typename ReturnT = const typename DpfKey::interior_node *>
struct path_memoizer_base
{
  public:
    using dpf_type = DpfKey;
    using input_type = typename DpfKey::input_type;
    using node_type = typename DpfKey::interior_node;
    using return_type = ReturnT;
    using iterator_type = return_type;

    virtual std::size_t assign_x(const dpf_type &, input_type) noexcept = 0;
    virtual node_type & operator[](std::size_t) noexcept = 0;

    virtual return_type begin() const noexcept = 0;
    virtual return_type end() const noexcept = 0;
};

template <typename DpfKey>
struct alignas(alignof(typename path_memoizer_key_t<DpfKey>::interior_node))
basic_path_memoizer final
    : public path_memoizer_base<path_memoizer_key_t<DpfKey>>
{
  public:
    using dpf_type = path_memoizer_key_t<DpfKey>;
    using input_type = typename dpf_type::input_type;
    using node_type = typename dpf_type::interior_node;
HEDLEY_PRAGMA(GCC diagnostic push)
HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
    using return_type = std::add_pointer_t<std::add_const_t<node_type>>;
HEDLEY_PRAGMA(GCC diagnostic pop)
    using iterator_type = return_type;
    static constexpr auto depth = dpf_type::depth;

    basic_path_memoizer()
      : dpf_{std::nullopt}, x_{std::nullopt}, filled_to_{0} { }
    basic_path_memoizer(basic_path_memoizer &&) noexcept = default;
    basic_path_memoizer(const basic_path_memoizer &) = default;
    basic_path_memoizer & operator=(basic_path_memoizer &&) noexcept = default;
    basic_path_memoizer & operator=(const basic_path_memoizer &) = default;
    ~basic_path_memoizer() = default;

    std::size_t assign_x(const dpf_type & dpf, input_type new_x) noexcept override
    {
        static constexpr auto clz_xor = utils::countl_zero_symmetric_difference<input_type>{};
        if (dpf_.has_value() == true && std::memcmp(&dpf_root_, &dpf.root(), sizeof(node_type)) == 0
            && std::memcmp(&dpf_common_part_hash_, &dpf.common_part_hash(), sizeof(digest_type)) == 0)
        {
            static constexpr auto complement_of = std::bit_not{};
            input_type old_x = x_.value_or(complement_of(new_x));
            x_ = new_x;
            const std::size_t resume = clz_xor(old_x, new_x) + 1;
            // Path slots below `resume` remain valid; invalidate the rest.
            if (filled_to_ >= resume)
                filled_to_ = resume - 1;
            return resume;
        }

        this->operator[](0) = dpf.root();
        dpf_ = std::cref(dpf);
        dpf_root_ = dpf.root();
        dpf_common_part_hash_ = dpf.common_part_hash();
        x_ = new_x;
        filled_to_ = 0;
        return 1;
    }

    node_type & operator[](std::size_t i) noexcept override
    {
        return arr_[i];
    }

    return_type begin() const noexcept override
    {
        if (x_.has_value() == true)
        {
            return std::addressof(arr_[depth]);
        }
        else
        {
            return end();
        }
    }

    return_type end() const noexcept override
    {
        return std::addressof(arr_[depth+1]);
    }

    /// Inclusive high-water: `arr_[0..filled_to_]` are valid for the current x.
    std::size_t filled_to() const noexcept { return filled_to_; }

    void note_filled(std::size_t level) noexcept
    {
        if (level > filled_to_)
            filled_to_ = level;
    }

  private:
    std::optional<std::reference_wrapper<const dpf_type>> dpf_;
    node_type dpf_root_;
    digest_type dpf_common_part_hash_;
    std::optional<input_type> x_;
    std::size_t filled_to_;
HEDLEY_PRAGMA(GCC diagnostic push)
HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
    std::array<node_type, depth+1> arr_;
HEDLEY_PRAGMA(GCC diagnostic pop)
};

template <typename DpfKey>
struct nonmemoizing_path_memoizer final
    : public path_memoizer_base<path_memoizer_key_t<DpfKey>>
{
  public:
    using dpf_type = path_memoizer_key_t<DpfKey>;
    using input_type = typename dpf_type::input_type;
    using node_type = typename dpf_type::interior_node;
HEDLEY_PRAGMA(GCC diagnostic push)
HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
    using return_type = std::add_pointer_t<std::add_const_t<node_type>>;
HEDLEY_PRAGMA(GCC diagnostic pop)
    using iterator_type = return_type;

    nonmemoizing_path_memoizer()
      : dpf_{std::nullopt} { }
    nonmemoizing_path_memoizer(nonmemoizing_path_memoizer &&) noexcept = default;
    nonmemoizing_path_memoizer(const nonmemoizing_path_memoizer &) = default;
    nonmemoizing_path_memoizer & operator=(nonmemoizing_path_memoizer &&) noexcept = default;
    nonmemoizing_path_memoizer & operator=(const nonmemoizing_path_memoizer &) = default;
    ~nonmemoizing_path_memoizer() = default;

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    std::size_t assign_x(const dpf_type & dpf, input_type) noexcept override
    {
        dpf_ = dpf;
        v = dpf.root();
        return 1;
    }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    node_type & operator[](std::size_t) noexcept override
    {
        return v;
    }

    return_type begin() const noexcept override
    {
        return std::addressof(v);
    }

    return_type end() const noexcept override
    {
        return std::addressof(v) + 1;
    }

  private:
    std::optional<std::reference_wrapper<const dpf_type>> dpf_;
    node_type v;
};

namespace detail
{

template <typename MemoizerT>
HEDLEY_ALWAYS_INLINE
auto make_path_memoizer()
{
    return MemoizerT();
}

template <typename T, typename = void>
struct has_path_high_water : std::false_type {};

template <typename T>
struct has_path_high_water<T,
    std::void_t<decltype(std::declval<T &>().filled_to()),
                decltype(std::declval<T &>().note_filled(std::size_t{}))>>
    : std::true_type {};

template <typename DpfKey, typename PathMemoizer>
std::size_t path_resume_for_level(PathMemoizer & path, const DpfKey & dpf,
    typename DpfKey::input_type x, std::size_t to_level)
{
    auto resume = path.assign_x(dpf, x);
    if constexpr (has_path_high_water<PathMemoizer>::value)
    {
        resume = std::min(resume, path.filled_to() + 1);
        if (resume == 0)
            resume = 1;
    }
    else if (resume > to_level)
    {
        resume = 1;
    }
    return resume;
}

template <typename PathMemoizer>
void path_note_filled_to(PathMemoizer & path, std::size_t level)
{
    if constexpr (has_path_high_water<PathMemoizer>::value)
        path.note_filled(level);
}

/// Walk interior nodes so `path[0..to_level]` is valid for `x`.
template <typename DpfKey, typename PathMemoizer>
void ensure_level(const DpfKey & dpf, typename DpfKey::input_type x,
    PathMemoizer & path, std::size_t to_level)
{
    auto level_index = path_resume_for_level(path, dpf, x, to_level);
    DPF_UNROLL_LOOP
    for (auto mask = dpf.msb_mask >> (level_index - 1);
         level_index <= to_level; ++level_index, mask >>= 1)
    {
        bool bit = !!(mask & x);
        auto cw = dpf.correction_word(level_index - 1, bit);
        path[level_index] =
            DpfKey::traverse_interior(path[level_index - 1], cw, bit);
    }
    path_note_filled_to(path, to_level);
}

}  // namespace detail

template <typename DpfKey>
auto make_basic_path_memoizer()
{
    return detail::make_path_memoizer<basic_path_memoizer<DpfKey>>();
}

template <typename DpfKey>
auto make_basic_path_memoizer(const DpfKey &)
{
    return make_basic_path_memoizer<DpfKey>();
}

template <typename DpfKey>
auto make_nonmemoizing_path_memoizer()
{
    return detail::make_path_memoizer<nonmemoizing_path_memoizer<DpfKey>>();
}

template <typename DpfKey>
auto make_nonmemoizing_path_memoizer(const DpfKey &)
{
    return make_nonmemoizing_path_memoizer<DpfKey>();
}

}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_PATH_MEMOIZER_HPP__