/// @file dpf/prg_count.hpp
/// @brief Thread-local counters of symmetric-key block invocations.
/// @details Each concrete block primitive (fixed-key AES-128 or AES-256 in
///          Matyas–Meyer–Oseas mode, ChaCha, LowMC) calls
///          `note_eval(n, primitive)` once per batch of `n` blocks. The count
///          is filed under the thread's current `purpose`:
///
///          - `expand`: PRG work, such as tree expansion, leaf conversion, and
///            label or column generation. This is the default.
///          - `hash`: correlation-robust hashing, such as IKNP row hashes,
///            garbled-gate hashes, and verification hashes.
///          - `harness`: blocks the measurement machinery computes for itself,
///            such as the `experiment` seed stream.
///
///          `eval_count()` is the protocol total: `expand` plus `hash` over
///          every primitive. Harness blocks are never part of it.
///          `count(purpose, primitive)` reads one cell and `snapshot()` reads
///          all of them. Call sites set the purpose with `purpose_scope`.
///          The counters are per thread, so work on another thread (a compute
///          pool, another party) is not included. Reset around a timed region
///          with `reset_eval_count`. The increment is one TLS add.
#ifndef LIBDPF_INCLUDE_DPF_PRG_COUNT_HPP__
#define LIBDPF_INCLUDE_DPF_PRG_COUNT_HPP__

#include <array>
#include <cstddef>
#include <cstdint>

#include "hedley/hedley.h"

namespace dpf
{
namespace prg
{

/// @brief Which block primitive produced a counted block.
enum class primitive : unsigned char
{
    aes128 = 0,
    aes256 = 1,
    chacha = 2,
    lowmc = 3
};

inline constexpr std::size_t primitive_count = 4;

/// @brief Why the block was computed.
enum class purpose : unsigned char
{
    expand = 0,
    hash = 1,
    harness = 2
};

inline constexpr std::size_t purpose_count = 3;

inline const char * primitive_name(primitive p) noexcept
{
    switch (p)
    {
    case primitive::aes128:
        return "aes128";
    case primitive::aes256:
        return "aes256";
    case primitive::chacha:
        return "chacha";
    case primitive::lowmc:
        return "lowmc";
    }
    return "aes128";
}

inline const char * purpose_name(purpose u) noexcept
{
    switch (u)
    {
    case purpose::expand:
        return "expand";
    case purpose::hash:
        return "hash";
    case purpose::harness:
        return "harness";
    }
    return "expand";
}

/// @brief Every cell, indexed `purpose * primitive_count + primitive`.
using counts = std::array<std::uint64_t, purpose_count * primitive_count>;

namespace detail
{

inline thread_local std::uint64_t eval_counts_tls[purpose_count * primitive_count] = {};
/// First cell of the current purpose's row (`purpose * primitive_count`).
inline thread_local std::size_t purpose_row_tls = 0;

}  // namespace detail

/// @brief Record `n` block invocations of `p` on this thread.
HEDLEY_ALWAYS_INLINE
void note_eval(std::uint64_t n = 1, primitive p = primitive::aes128) noexcept
{
    detail::eval_counts_tls[detail::purpose_row_tls + static_cast<std::size_t>(p)] += n;
}

/// @brief Zero every counter on this thread.
HEDLEY_ALWAYS_INLINE
void reset_eval_count() noexcept
{
    for (auto & c : detail::eval_counts_tls)
        c = 0;
}

/// @brief Blocks of `p` computed for `u` since the last reset on this thread.
HEDLEY_ALWAYS_INLINE
std::uint64_t count(purpose u, primitive p) noexcept
{
    return detail::eval_counts_tls[static_cast<std::size_t>(u) * primitive_count
        + static_cast<std::size_t>(p)];
}

/// @brief Protocol blocks (`expand` + `hash`, every primitive) since the last
///        reset on this thread.
HEDLEY_ALWAYS_INLINE
std::uint64_t eval_count() noexcept
{
    std::uint64_t total = 0;
    for (std::size_t i = 0; i < 2 * primitive_count; ++i)
        total += detail::eval_counts_tls[i];
    return total;
}

/// @brief Every cell on this thread.
inline counts snapshot() noexcept
{
    counts out{};
    for (std::size_t i = 0; i < out.size(); ++i)
        out[i] = detail::eval_counts_tls[i];
    return out;
}

/// @brief Add cells counted on another thread (a compute-pool task this
///        thread waited for) to this thread's cells.
inline void absorb(const counts & delta) noexcept
{
    for (std::size_t i = 0; i < delta.size(); ++i)
        detail::eval_counts_tls[i] += delta[i];
}

/// @brief File this thread's block invocations under `u` until the scope ends.
class purpose_scope
{
public:
    explicit purpose_scope(purpose u) noexcept : prev_(detail::purpose_row_tls)
    {
        detail::purpose_row_tls = static_cast<std::size_t>(u) * primitive_count;
    }
    ~purpose_scope() { detail::purpose_row_tls = prev_; }
    purpose_scope(const purpose_scope &) = delete;
    purpose_scope & operator=(const purpose_scope &) = delete;

private:
    std::size_t prev_;
};

}  // namespace prg
}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_PRG_COUNT_HPP__