/// @file dpf/caller_fold.hpp
/// @brief One-pass caller fold for the full-domain / point / sequence walks.
/// @details `sketch_ref` folds every extractable share into three `fp61`
///          moments. A protocol whose audit lives in a different group (the
///          Express multiplication proof, Pika's Schwartz–Zippel check, a
///          keyword-PIR record XOR) wants the same one-pass hook without a
///          second expansion and without the `fp61` conversion. These helpers
///          take an inlined `fold(index, share)` callable, invoked once per
///          written output in the same loop that writes the buffer. The fold is
///          a template argument, not a virtual call, so it inlines exactly the
///          way `sketch_ref::absorb` does today. `dpf::sketch` stays the
///          weight-1 `fp61` fold; existing call sites are untouched.
/// @copyright Copyright (c) 2019-2026 Ryan Henry and [others](@ref authors)
/// @license Released under a GNU General Public v2.0 (GPLv2) license.

#ifndef LIBDPF_INCLUDE_DPF_CALLER_FOLD_HPP__
#define LIBDPF_INCLUDE_DPF_CALLER_FOLD_HPP__

#include <cstddef>
#include <iterator>
#include <type_traits>
#include <utility>

#include "hedley/hedley.h"

#include "dpf/dpf_key.hpp"
#include "dpf/eval_common.hpp"
#include "dpf/eval_full.hpp"
#include "dpf/eval_point.hpp"
#include "dpf/eval_target.hpp"
#include "dpf/eval_walk.hpp"

namespace dpf
{

namespace detail_fold
{

/// @brief True when `Fold` is callable as `fold(std::size_t, share)` for the
///        group element written by output `I` of `DpfKey`.
template <typename Fold, typename DpfKey, std::size_t I, typename = void>
struct is_index_fold : std::false_type {};

template <typename Fold, typename DpfKey, std::size_t I>
struct is_index_fold<Fold, DpfKey, I,
    std::void_t<decltype(std::declval<Fold &>()(
        std::declval<std::size_t>(),
        std::declval<typename DpfKey::template concrete_output_type<I>>()))>>
    : std::true_type {};

template <typename Fold, typename DpfKey, std::size_t I>
inline constexpr bool is_index_fold_v =
    is_index_fold<std::decay_t<Fold>, DpfKey, I>::value;

}  // namespace detail_fold

// N.B.: the one-pass `eval_full_add_into(buf, key, fold)` hook lives in
// `dpf/eval_walk.hpp` alongside the `rotate` / `sketch_ref` overloads. This
// header adds the remaining fold-carrying walks (allocate-and-fold full eval,
// point eval, and the full-domain keyword-PIR XOR) so every walk that accepts
// a `sketch_ref` also accepts a generic caller fold.

/// @brief Full-domain expansion of output `I`, folding each written share.
/// @details Allocates a fresh buffer (like `eval_full(key)`), folds every point
///          into `fold`, and returns the `(buffer, iterable)` pair so the caller
///          keeps the expanded shares as well as the audit.
/// \complexity One full-domain expansion, `Θ(2^n)`, plus one fold per point.
template <std::size_t I = 0, typename DpfKey, typename Fold,
          std::enable_if_t<looks_like_dpf_key_v<DpfKey>
              && !is_multilevel_key_v<DpfKey>
              && detail_fold::is_index_fold_v<Fold, DpfKey, I>, bool> = true>
HEDLEY_WARN_UNUSED_RESULT
auto eval_full_fold(const DpfKey & dpf, Fold fold)
{
    auto result = eval_full<I>(dpf);
    auto & iter = result.second;
    std::size_t i = 0;
    for (auto it = std::begin(iter); it != std::end(iter); ++it, ++i)
        fold(i, detail_walk::group_value(*it));
    return result;
}

/// @brief Evaluate output `I` at `x` and fold the single written share.
/// @details The fold is called once, with index `0`, matching the one output
///          `eval_point` writes.
/// \complexity O(n) time; one interior traversal per level.
template <std::size_t I = 0, typename DpfKey, typename InputT, typename Fold,
          std::enable_if_t<looks_like_dpf_key_v<DpfKey>
              && !is_multilevel_key_v<DpfKey>
              && detail_fold::is_index_fold_v<Fold, DpfKey, I>, bool> = true>
auto eval_point_fold(const DpfKey & dpf, InputT && x, Fold fold)
{
    auto out = eval_point<I>(dpf, std::forward<InputT>(x));
    fold(std::size_t{0}, detail_walk::group_value(*out));
    return out;
}

/// @brief Keyword-PIR XOR fold: `⊕ records[i]` over `i` where `DPF_I(i)` is set.
/// @details Runs one full-domain bit expansion of output `I` and, in that same
///          loop, XORs `records[i]` into a running accumulator whenever the
///          party's bit share at `i` is 1 — without ever materializing the bit
///          vector. Each server returns its share of the XOR; the two shares
///          reconstruct (XOR) to `⊕ records[i]` over the 1-set of the DPF, i.e.
///          the matched record for a point key. `records` is indexed by domain
///          point and must cover the domain (or at least the evaluated prefix).
/// \complexity One full-domain expansion, `Θ(2^n)`, plus one XOR per set bit.
template <std::size_t I = 0, typename DpfKey, typename Records,
          std::enable_if_t<looks_like_dpf_key_v<DpfKey>
              && !is_multilevel_key_v<DpfKey>, bool> = true>
HEDLEY_WARN_UNUSED_RESULT
auto eval_sequence_xor(const DpfKey & key, const Records & records)
{
    using record_type = std::decay_t<decltype(records[std::size_t{0}])>;
    auto result = eval_full<I>(key);
    auto & iter = result.second;
    record_type acc{};
    std::size_t i = 0;
    const std::size_t n = static_cast<std::size_t>(std::size(records));
    for (auto it = std::begin(iter); it != std::end(iter) && i < n; ++it, ++i)
    {
        if (static_cast<bool>(*it))
            acc = static_cast<record_type>(acc ^ records[i]);
    }
    return acc;
}

}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_CALLER_FOLD_HPP__