/// @file dpf/prg.hpp
/// @brief PRG aliases, the share expander, and the call counter.
/// @author Ryan Henry <ryan.henry@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_PRG_HPP__
#define LIBDPF_INCLUDE_DPF_PRG_HPP__

#include "hedley/hedley.h"

#include <cstring>
#include <stdexcept>
#include <type_traits>

#include "dpf/prg_count.hpp"
#include "dpf/prg_aes.hpp"
#include "dpf/prg_aes_ccr.hpp"
#include "dpf/prg_chacha.hpp"
#include "dpf/prg_dummy.hpp"
#include "dpf/prg_lowmc.hpp"
#include "dpf/secret_share.hpp"

namespace dpf
{

namespace prg
{

namespace detail
{

/// @brief Fill `T` from consecutive PRG blocks starting at `pos` (low bytes first).
/// @tparam PRG pseudorandom generator
/// @tparam T value type
/// @param seed the PRG seed
/// @param pos the 0-based index
/// @return the returned `T`
template <typename PRG, typename T>
HEDLEY_NO_THROW
HEDLEY_ALWAYS_INLINE
T expand_raw(typename PRG::block_type seed, psnip_uint32_t pos) noexcept
{
    static_assert(std::is_trivially_copyable_v<T>,
        "prg expand requires a trivially copyable value type");
    constexpr std::size_t nbytes = sizeof(T);
    constexpr std::size_t nblocks =
        (nbytes + sizeof(typename PRG::block_type) - 1)
            / sizeof(typename PRG::block_type);
    alignas(typename PRG::block_type) typename PRG::block_type
        blocks[nblocks ? nblocks : 1];
    PRG::eval(seed, blocks, static_cast<psnip_uint32_t>(nblocks ? nblocks : 1),
        pos);
    T value{};
    std::memcpy(&value, blocks, nbytes);
    return value;
}

template <typename PRG, typename T, std::size_t Party>
HEDLEY_NO_THROW
HEDLEY_ALWAYS_INLINE
auto expand_as_share(typename PRG::block_type seed,
    psnip_uint32_t pos) noexcept
{
    return subtractive_share<T, Party>::from_raw(expand_raw<PRG, T>(seed, pos));
}

}  // namespace detail

/// \complexity `ceil(sizeof(T) / sizeof(block_type))` PRG evaluations (`expand_raw`), then a copy.
template <typename AesKey>
template <typename T, std::size_t Party>
HEDLEY_NO_THROW
auto aes<AesKey>::expand(block_type seed, psnip_uint32_t pos) noexcept
{
    return detail::expand_as_share<aes<AesKey>, T, Party>(seed, pos);
}

/// \complexity `ceil(sizeof(T) / sizeof(block_type))` PRG evaluations (`expand_raw`), then a copy.
template <typename T, std::size_t Party>
HEDLEY_NO_THROW
auto dummy::expand(block_type seed, psnip_uint32_t pos) noexcept
{
    return detail::expand_as_share<dummy, T, Party>(seed, pos);
}

/// \complexity `ceil(sizeof(T) / sizeof(block_type))` PRG evaluations (`expand_raw`), then a copy.
template <typename T, std::size_t Party>
HEDLEY_NO_THROW
auto lowmc128::expand(block_type seed, psnip_uint32_t pos) noexcept
{
    return detail::expand_as_share<lowmc128, T, Party>(seed, pos);
}

/// \complexity `ceil(sizeof(T) / sizeof(block_type))` PRG evaluations (`expand_raw`), then a copy.
template <typename T, std::size_t Party>
HEDLEY_NO_THROW
auto aes128_ccr::expand(block_type seed, psnip_uint32_t pos) noexcept
{
    return detail::expand_as_share<aes128_ccr, T, Party>(seed, pos);
}

/// \complexity `ceil(sizeof(T) / sizeof(block_type))` PRG evaluations (`expand_raw`), then a copy.
template <unsigned Rounds>
template <typename T, std::size_t Party>
HEDLEY_NO_THROW
auto chacha<Rounds>::expand(block_type seed, psnip_uint32_t pos) noexcept
{
    return detail::expand_as_share<chacha<Rounds>, T, Party>(seed, pos);
}

/// @brief PRG adapter that exposes the shared TLS eval counter.
/// @details Counting lives in the concrete `PRG::eval` paths via
///          `note_eval`. This wrapper forwards and reports `eval_count()`.
template <typename PRG>
struct counter_wrapper final
{
    using block_type = typename PRG::block_type;

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    static block_type eval(block_type seed, psnip_uint32_t pos) noexcept
    {
        return PRG::eval(seed, pos);
    }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    static auto eval01(block_type seed) noexcept
    {
        return PRG::eval01(seed);
    }

    HEDLEY_ALWAYS_INLINE
    static void eval(block_type seed, block_type * HEDLEY_RESTRICT output,
        psnip_uint32_t count, psnip_uint32_t pos = 0)
    {
        PRG::eval(seed, output, count, pos);
    }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    HEDLEY_NON_NULL(1, 2, 3)
    static void eval01_x4(const block_type * HEDLEY_RESTRICT seeds,
        block_type * HEDLEY_RESTRICT left,
        block_type * HEDLEY_RESTRICT right) noexcept
    {
        PRG::eval01_x4(seeds, left, right);
    }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    HEDLEY_NON_NULL(1, 2)
    static void eval_x4(const block_type * HEDLEY_RESTRICT seeds,
        block_type * HEDLEY_RESTRICT output, psnip_uint32_t pos = 0) noexcept
    {
        PRG::eval_x4(seeds, output, pos);
    }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    HEDLEY_NON_NULL(1, 2)
    static void eval_x8(const block_type * HEDLEY_RESTRICT seeds,
        block_type * HEDLEY_RESTRICT output, psnip_uint32_t pos = 0) noexcept
    {
        PRG::eval_x8(seeds, output, pos);
    }

    template <typename T, std::size_t Party>
    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    static auto expand(block_type seed, psnip_uint32_t pos = 0) noexcept
    {
        return detail::expand_as_share<PRG, T, Party>(seed, pos);
    }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    static std::size_t count() noexcept
    {
        return static_cast<std::size_t>(eval_count());
    }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    static void reset() noexcept
    {
        reset_eval_count();
    }
};  // struct counter_wrapper

}  // namespace prg

}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_PRG_HPP__