libdpf/include/dpf/buffered_prg.hpp

/// @file dpf/buffered_prg.hpp
/// @brief Buffered, seekable lanes over a libdpf PRG.
/// @details The default PRG is `dpf::prg::aes128`. Any PRG with
///          `block_type`, `eval(seed, pos)`, and `eval(seed, out, count, pos)`
///          can be substituted. A master block is the recorded seed. Role `r`
///          owns two streams: values are `PRG::eval(master, r)` and share
///          masks are `PRG::eval(tweaked_master, r)`. Element `index` is a
///          contiguous run of blocks on that stream, so a forward scan is one
///          multi-block `eval`.
/// @copyright Copyright (c) 2019-2026 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_BUFFERED_PRG_HPP__
#define LIBDPF_INCLUDE_DPF_BUFFERED_PRG_HPP__

#include <cstddef>
#include <cstdint>
#include <cstring>
#include <map>
#include <stdexcept>
#include <tuple>
#include <type_traits>
#include <utility>
#include <vector>

#include "hedley/hedley.h"

#include "dpf/aligned_allocator.hpp"
#include "dpf/prg.hpp"
#include "dpf/random.hpp"

namespace dpf
{
namespace randomness
{

namespace detail
{

template <typename PRG>
typename PRG::block_type mask_master(typename PRG::block_type master) noexcept
{
    unsigned char raw[sizeof(master)];
    std::memcpy(raw, &master, sizeof(master));
    raw[sizeof(master) - 1] ^= 0x01u;
    typename PRG::block_type out{};
    std::memcpy(&out, raw, sizeof(out));
    return out;
}

template <typename PRG, typename T>
struct lane_codec
{
    static_assert(std::is_trivially_copyable_v<T>,
        "prg lanes require a trivially copyable value type");

    using block_type = typename PRG::block_type;

    static T at(block_type seed, std::uint64_t index)
    {
        T out{};
        fill(seed, index, &out, 1u);
        return out;
    }

    /// Element `index` is the packed byte range `[index * sizeof(T), ...)`.
    static void fill(block_type seed, std::uint64_t index, T * out, std::size_t count)
    {
        if (count == 0)
            return;
        constexpr std::uint64_t block_bytes = sizeof(block_type);
        std::uint64_t byte_off = index * static_cast<std::uint64_t>(sizeof(T));
        std::uint64_t nbytes = static_cast<std::uint64_t>(count) * sizeof(T);
        std::uint64_t start = byte_off / block_bytes;
        std::uint64_t end = byte_off + nbytes;
        std::uint64_t nblocks = (end + block_bytes - 1u) / block_bytes - start;
        if (start > static_cast<std::uint64_t>(UINT32_MAX)
            || nblocks > static_cast<std::uint64_t>(UINT32_MAX)
            || start > static_cast<std::uint64_t>(UINT32_MAX) - nblocks)
        {
            throw std::invalid_argument("prg lane index is out of range");
        }
        HEDLEY_PRAGMA(GCC diagnostic push)
        HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
        dpf::aligned_allocator<block_type> alloc;
        auto blocks = alloc.allocate_unique_ptr(static_cast<std::size_t>(nblocks));
        HEDLEY_PRAGMA(GCC diagnostic pop)
        PRG::eval(seed, blocks.get(), static_cast<psnip_uint32_t>(nblocks),
            static_cast<psnip_uint32_t>(start));
        auto * bytes = reinterpret_cast<const unsigned char *>(blocks.get());
        std::memcpy(out, bytes + static_cast<std::size_t>(byte_off % block_bytes),
            static_cast<std::size_t>(nbytes));
    }
};

template <typename PRG, typename T>
struct buffered_slot
{
    using block_type = typename PRG::block_type;

    explicit buffered_slot(block_type seed, std::size_t buffer_elems)
      : seed_(seed),
        buffer_(buffer_elems),
        absolute_pos_(0u),
        filled_(0u),
        next_(0u)
    {
        refill(0u);
    }

    T get()
    {
        if (next_ >= filled_)
            refill(absolute_pos_);
        T v = buffer_[next_];
        ++next_;
        ++absolute_pos_;
        return v;
    }

    void fill(T * out, std::size_t count)
    {
        std::size_t written = 0u;
        while (written < count)
        {
            if (next_ >= filled_)
                refill(absolute_pos_);
            std::size_t available = filled_ - next_;
            std::size_t take = available < count - written ? available : count - written;
            std::memcpy(out + written, buffer_.data() + next_, take * sizeof(T));
            written += take;
            next_ += take;
            absolute_pos_ += static_cast<std::uint64_t>(take);
        }
    }

    T at(std::uint64_t index) const
    {
        return lane_codec<PRG, T>::at(seed_, index);
    }

    std::uint64_t sampled() const noexcept { return absolute_pos_; }

private:
    void refill(std::uint64_t at_elem)
    {
        next_ = 0u;
        filled_ = buffer_.size();
        lane_codec<PRG, T>::fill(seed_, at_elem, buffer_.data(), buffer_.size());
    }

    block_type seed_;
    std::vector<T> buffer_;
    std::uint64_t absolute_pos_;
    std::size_t filled_;
    std::size_t next_;
};

}  // namespace detail

template <typename PRG = dpf::prg::aes128>
typename PRG::block_type sample_master_seed()
{
    return dpf::uniform_sample<typename PRG::block_type>();
}

/// Fixed lanes. Lane `I` is `PRG::eval(master, I)`.
template <typename PRG, typename... Ts>
class buffered_prg
{
public:
    using prg_type = PRG;
    using seed_type = typename PRG::block_type;
    static constexpr std::size_t stream_count = sizeof...(Ts);

    explicit buffered_prg(std::size_t per_stream_buffer_elems = 1024u)
      : seed_(sample_master_seed<PRG>()),
        buffers_(make_buffers(per_stream_buffer_elems))
    { }

    explicit buffered_prg(seed_type seed, std::size_t per_stream_buffer_elems = 1024u)
      : seed_(seed),
        buffers_(make_buffers(per_stream_buffer_elems))
    { }

    const seed_type & seed() const noexcept { return seed_; }

    template <std::size_t I>
    auto get()
    {
        static_assert(I < stream_count, "stream index out of range");
        return std::get<I>(buffers_).get();
    }

    template <std::size_t I>
    void fill(std::tuple_element_t<I, std::tuple<Ts...>> * out, std::size_t count)
    {
        static_assert(I < stream_count, "stream index out of range");
        std::get<I>(buffers_).fill(out, count);
    }

    template <std::size_t I>
    auto at(std::uint64_t index) const
    {
        static_assert(I < stream_count, "stream index out of range");
        return std::get<I>(buffers_).at(index);
    }

    template <std::size_t I>
    std::uint64_t sampled() const noexcept
    {
        static_assert(I < stream_count, "stream index out of range");
        return std::get<I>(buffers_).sampled();
    }

private:
    template <std::size_t... I>
    std::tuple<detail::buffered_slot<PRG, Ts>...>
    make_buffers_impl(std::size_t n, std::index_sequence<I...>) const
    {
        return std::tuple<detail::buffered_slot<PRG, Ts>...>(
            detail::buffered_slot<PRG, Ts>(
                PRG::eval(seed_, static_cast<psnip_uint32_t>(I)), n)...);
    }

    std::tuple<detail::buffered_slot<PRG, Ts>...> make_buffers(std::size_t n) const
    {
        if (n == 0)
            throw std::invalid_argument("per_stream_buffer_elems must be positive");
        return make_buffers_impl(n, std::make_index_sequence<stream_count>{});
    }

    seed_type seed_;
    std::tuple<detail::buffered_slot<PRG, Ts>...> buffers_;
};

template <typename... Ts>
using aes_buffered_prg = buffered_prg<dpf::prg::aes128, Ts...>;

/// Dynamic lanes of one value type. `value_at(role, index)` and
/// `mask_at(role, index)` are independent of call order. A window cache
/// refills from the requested index.
template <typename T, typename PRG = dpf::prg::aes128>
class lane_table
{
public:
    using prg_type = PRG;
    using seed_type = typename PRG::block_type;
    using value_type = T;

    explicit lane_table(std::size_t window = 256u)
      : lane_table(sample_master_seed<PRG>(), window)
    { }

    explicit lane_table(seed_type seed, std::size_t window = 256u)
      : seed_(seed),
        mask_seed_(detail::mask_master<PRG>(seed)),
        window_(window)
    {
        if (window_ == 0)
            throw std::invalid_argument("prg lane window must be positive");
    }

    lane_table(const lane_table &) = delete;
    lane_table & operator=(const lane_table &) = delete;
    lane_table(lane_table &&) = default;
    lane_table & operator=(lane_table &&) = default;

    const seed_type & seed() const noexcept { return seed_; }

    T value_at(std::uint32_t role, std::uint64_t index) const
    {
        return cached(role, true, index);
    }

    T mask_at(std::uint32_t role, std::uint64_t index) const
    {
        return cached(role, false, index);
    }

    void fill_values(std::uint32_t role, std::uint64_t index, T * out, std::size_t n) const
    {
        if (n == 0)
            return;
        detail::lane_codec<PRG, T>::fill(slot(role).value_seed, index, out, n);
    }

    void fill_masks(std::uint32_t role, std::uint64_t index, T * out, std::size_t n) const
    {
        if (n == 0)
            return;
        detail::lane_codec<PRG, T>::fill(slot(role).mask_seed, index, out, n);
    }

private:
    using block_type = typename PRG::block_type;

    struct slot_pair
    {
        block_type value_seed{};
        block_type mask_seed{};
        std::vector<T> value_cache;
        std::vector<T> mask_cache;
        std::uint64_t value_base = 0;
        std::uint64_t mask_base = 0;
        bool value_hot = false;
        bool mask_hot = false;
    };

    slot_pair & slot(std::uint32_t role) const
    {
        auto it = slots_.find(role);
        if (it != slots_.end())
            return it->second;
        slot_pair created;
        created.value_seed = PRG::eval(seed_, role);
        created.mask_seed = PRG::eval(mask_seed_, role);
        created.value_cache.resize(window_);
        created.mask_cache.resize(window_);
        auto inserted = slots_.emplace(role, std::move(created));
        return inserted.first->second;
    }

    T cached(std::uint32_t role, bool values, std::uint64_t index) const
    {
        slot_pair & s = slot(role);
        std::vector<T> & buf = values ? s.value_cache : s.mask_cache;
        std::uint64_t & base = values ? s.value_base : s.mask_base;
        bool & hot = values ? s.value_hot : s.mask_hot;
        block_type stream = values ? s.value_seed : s.mask_seed;
        if (hot && index >= base && index < base + buf.size())
            return buf[static_cast<std::size_t>(index - base)];
        base = index;
        hot = true;
        detail::lane_codec<PRG, T>::fill(stream, index, buf.data(), buf.size());
        return buf[0];
    }

    seed_type seed_;
    seed_type mask_seed_;
    std::size_t window_;
    mutable std::map<std::uint32_t, slot_pair> slots_;
};

}  // namespace randomness
}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_BUFFERED_PRG_HPP__