/// @file dpf/prg.hpp
/// @brief
/// @details
/// @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 <atomic>
#include <cstring>
#include <type_traits>

#include "dpf/prg_aes.hpp"
#include "dpf/prg_dummy.hpp"
#include "dpf/prg_lowmc.hpp"
#include "dpf/secret_share.hpp"

namespace dpf
{

namespace prg
{

namespace detail
{

/// Fill `T` from consecutive PRG blocks starting at `pos` (low bytes first).
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

template <typename AesKey>
template <typename T, std::size_t Party>
auto aes<AesKey>::expand(block_type seed, psnip_uint32_t pos) noexcept
{
    return detail::expand_as_share<aes<AesKey>, T, Party>(seed, pos);
}

template <typename T, std::size_t Party>
auto dummy::expand(block_type seed, psnip_uint32_t pos) noexcept
{
    return detail::expand_as_share<dummy, T, Party>(seed, pos);
}

template <typename T, std::size_t Party>
auto lowmc128::expand(block_type seed, psnip_uint32_t pos) noexcept
{
    return detail::expand_as_share<lowmc128, T, Party>(seed, pos);
}

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
    {
        count_.fetch_add(1, std::memory_order::memory_order_relaxed);
        return PRG::eval(seed, pos);
    }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    static auto eval01(block_type seed) noexcept
    {
        count_.fetch_add(2, std::memory_order::memory_order_relaxed);
        return PRG::eval01(seed);
    }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    static void eval(block_type seed, block_type * HEDLEY_RESTRICT output,
        psnip_uint32_t count, psnip_uint32_t pos = 0) noexcept
    {
        count_.fetch_add(count, std::memory_order::memory_order_relaxed);
        PRG::eval(seed, output, count, pos);
    }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    static void eval01_x4(const block_type * HEDLEY_RESTRICT seeds,
        block_type * HEDLEY_RESTRICT left,
        block_type * HEDLEY_RESTRICT right) noexcept
    {
        count_.fetch_add(8, std::memory_order::memory_order_relaxed);
        PRG::eval01_x4(seeds, left, right);
    }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    static void eval_x4(const block_type * HEDLEY_RESTRICT seeds,
        block_type * HEDLEY_RESTRICT output, psnip_uint32_t pos = 0) noexcept
    {
        count_.fetch_add(4, std::memory_order::memory_order_relaxed);
        PRG::eval_x4(seeds, output, pos);
    }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    static void eval_x8(const block_type * HEDLEY_RESTRICT seeds,
        block_type * HEDLEY_RESTRICT output, psnip_uint32_t pos = 0) noexcept
    {
        count_.fetch_add(8, std::memory_order::memory_order_relaxed);
        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
    {
        count_.fetch_add(
            static_cast<std::size_t>(
                (sizeof(T) + sizeof(block_type) - 1) / sizeof(block_type)),
            std::memory_order::memory_order_relaxed);
        return detail::expand_as_share<PRG, T, Party>(seed, pos);
    }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    static std::size_t count() noexcept
    {
        return count_;
    }

  private:
    inline static std::atomic_size_t count_{0};
};  // struct counter_wrapper

}  // namespace prg

}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_PRG_HPP__