/// @file dpf/path_memoizer.hpp
/// @brief Workspaces that resume a root-to-leaf DPF walk.
/// @details `basic_path_memoizer` keeps one interior node per level. The next
///          `eval_point` recomputes the suffix after the common prefix with
///          the previous input. `nonmemoizing_path_memoizer` keeps a single
///          node and starts at the root on every call.
///
///          Pass a mutable lvalue to `eval_point`. The factories unwrap
///          `party_key`, so a memoizer built from either party's type
///          accepts both parties. A different root restarts the path.
/// @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"
#include "dpf/verifiable.hpp"

namespace dpf
{

/// @brief 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;

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

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

/// @brief One interior node per level. `assign_x` returns the first level that the
/// next walk must recompute. `filled_to` is the deepest level already
/// written for the current input. Callers pass this object to `eval_point`;
/// they do not call `assign_x` themselves.
/// @tparam DpfKey DPF key type
/// \complexity O(n) nodes, one per level (`arr_` has `depth + 1` slots). `assign_x` is O(1) after a common-prefix XOR.
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} { }
    HEDLEY_NO_THROW
    basic_path_memoizer(basic_path_memoizer &&) noexcept = default;
    basic_path_memoizer(const basic_path_memoizer &) = default;
    HEDLEY_NO_THROW
    basic_path_memoizer & operator=(basic_path_memoizer &&) noexcept = default;
    basic_path_memoizer & operator=(const basic_path_memoizer &) = default;
    ~basic_path_memoizer() = default;

    HEDLEY_NO_THROW
    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;
    }

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

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

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

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

    HEDLEY_NO_THROW
    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)
};

/// @brief A single interior node. Every `assign_x` restarts at the root.
/// @tparam DpfKey DPF key type
/// \complexity O(1) space (one node). Every `assign_x` restarts at the root, so `eval_point` still walks n levels.
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} { }
    HEDLEY_NO_THROW
    nonmemoizing_path_memoizer(nonmemoizing_path_memoizer &&) noexcept = default;
    nonmemoizing_path_memoizer(const nonmemoizing_path_memoizer &) = default;
    HEDLEY_NO_THROW
    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);
    }

    HEDLEY_NO_THROW
    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);
}

/// @brief Walk interior nodes so `path[0..to_level]` is valid for `x`.
/// @details When `pi` is non-null and the key is verifiable, each newly
///          traversed node is folded into the proof token (LightShark-style
///          path authentication on aggregate walks). When `fold_cached` is
///          true, nodes already held by the memoizer are hashed into `pi`
///          without re-expanding the PRG — required after `init_proof` on a
///          warm path. Leave `fold_cached` false when continuing one
///          accumulator that already contains those nodes.
/// @tparam DpfKey DPF key type
/// @tparam PathMemoizer path memoizer type
/// @param dpf the DPF key
/// @param x the `x`
/// @param path the root-to-leaf path
/// @param to_level the `to_level`
/// @param pi optional VDPF proof accumulator
/// @param fold_cached whether to fold memoized prefix nodes into a fresh token
template <typename DpfKey, typename PathMemoizer>
void ensure_level(const DpfKey & dpf, typename DpfKey::input_type x,
    PathMemoizer & path, std::size_t to_level, proof_token * pi = nullptr,
    bool fold_cached = false)
{
    auto level_index = path_resume_for_level(path, dpf, x, to_level);
    if constexpr (DpfKey::is_verifiable)
    {
        if (pi != nullptr && fold_cached && level_index > 1)
        {
            for (std::size_t li = 1; li < level_index && li <= to_level; ++li)
            {
                const auto x_bits = static_cast<psnip_uint64_t>(
                    utils::to_integral_type<typename DpfKey::input_type>{}(x)
                    >> (utils::bitlength_of_v<typename DpfKey::input_type>
                        - li));
                detail::vdpf::fold_node(*pi, li - 1, x_bits, path[li],
                    dpf.correction_seeds()[li - 1]);
            }
        }
    }
    // Resume past `to_level` means the spine is already valid. Do not form
    // `msb_mask >> (level_index - 1)` in that case: for a full-width domain
    // (`to_level == input_bits`) that shift is exactly the type width (UB).
    if (level_index <= 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);
            const bool is_last = DpfKey::tree::is_last_level(level_index - 1,
                dpf.depth);
            path[level_index] =
                DpfKey::traverse_interior(path[level_index - 1], cw, bit,
                    is_last);
            if constexpr (DpfKey::is_verifiable)
            {
                if (pi != nullptr)
                {
                    const auto x_bits = static_cast<psnip_uint64_t>(
                        utils::to_integral_type<typename DpfKey::input_type>{}(x)
                        >> (utils::bitlength_of_v<typename DpfKey::input_type>
                            - level_index));
                    detail::vdpf::fold_node(*pi, level_index - 1, x_bits,
                        path[level_index],
                        dpf.correction_seeds()[level_index - 1]);
                }
            }
        }
    }
    path_note_filled_to(path, to_level);
}

}  // namespace detail

/// @brief Path workspace for `DpfKey`. A `party_key` argument is unwrapped, and the
/// result accepts both parties.
/// @snippet evaluation/memoizers.cpp path-memoizer
/// @tparam DpfKey DPF key type
/// @return Path workspace for `DpfKey`
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>();
}

/// @brief Single-node path workspace. Suitable for one query.
/// @tparam DpfKey DPF key type
/// @return Single-node path workspace
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__