/// @file dpf/output_buffer.hpp
/// @brief Move-only storage for shares written by multi-point evaluation.
/// @details Slot type follows the key. A `party_key` leaf buffer holds
///          `subtractive_share`s; a comparison buffer holds
///          `additive_share`s. `bit`, `twobit`, and `nyble` slots are packed.
///          Trivially default-constructible slots are left uninitialized
///          because evaluation overwrites every slot it is responsible for.
///
///          `bit` / `gf2`, `twobit` / `gf22`, and `nyble` / `gf24` slots are packed.
///          `eval_interval` and recipe `eval_sequence` take the buffer by
///          non-const reference. The returned iterable refers into it.
///          @snippet evaluation/output_buffers.cpp output-buffer
/// @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_OUTPUT_BUFFER_HPP__
#define LIBDPF_INCLUDE_DPF_OUTPUT_BUFFER_HPP__

#include "hedley/hedley.h"

#include <cstddef>
#include <algorithm>
#include <tuple>
#include <limits>
#include <iterator>
#include <new>
#include <type_traits>
#include <vector>

#include "dpf/aligned_allocator.hpp"
#include "dpf/leaf_node.hpp"
#include "dpf/utils.hpp"
#include "dpf/bit.hpp"
#include "dpf/bit_array.hpp"
#include "dpf/packed_array.hpp"
#include "dpf/secret_share.hpp"
#include "dpf/sequence_recipe.hpp"
#include "dpf/sequence_utils.hpp"

namespace dpf
{

/// @brief Buffer element type for leaf eval of `KeyT`: party-tagged subtractive
/// share when `KeyT` is a `party_key`, otherwise the concrete output.
/// @tparam KeyT key type
/// @tparam OutputT output type
template <typename KeyT, typename OutputT, bool = is_party_key_v<KeyT>>
struct leaf_buffer_elem
{
    using type = OutputT;
};
template <typename KeyT, typename OutputT>
struct leaf_buffer_elem<KeyT, OutputT, true>
{
    using type = subtractive_share<OutputT, party_of_v<KeyT>>;
};
template <typename KeyT, typename OutputT>
using leaf_buffer_elem_t = typename leaf_buffer_elem<KeyT, OutputT>::type;

/// @brief Buffer element type for comparison eval of `KeyT`.
/// @tparam KeyT key type
/// @tparam Beta payload type
template <typename KeyT, typename Beta, bool = is_party_key_v<KeyT>>
struct cmp_buffer_elem
{
    using type = Beta;
};
template <typename KeyT, typename Beta>
struct cmp_buffer_elem<KeyT, Beta, true>
{
    using type = additive_share<Beta, party_of_v<KeyT>>;
};
template <typename KeyT, typename Beta>
using cmp_buffer_elem_t = typename cmp_buffer_elem<KeyT, Beta>::type;

/// @brief `std::vector(n)` value-initializes every slot. Interval / full eval
/// overwrites the whole buffer, so skip default-construction for trivial
/// `T`. Non-trivial outputs still run their default constructor.
/// @tparam T value type
/// @tparam Alignment allocation alignment
template <typename T,
          std::size_t Alignment>
class output_buffer_allocator : public aligned_allocator<T, Alignment>
{
  public:
    using is_always_equal = std::true_type;
    using propagate_on_container_move_assignment = std::true_type;

    template <typename U>
    struct rebind
    {
        using other = output_buffer_allocator<U, Alignment>;
    };

    HEDLEY_NO_THROW
    output_buffer_allocator() noexcept = default;
    HEDLEY_NO_THROW
    output_buffer_allocator(const output_buffer_allocator &) noexcept = default;
    template <typename U>
    HEDLEY_NO_THROW
    output_buffer_allocator(const output_buffer_allocator<U, Alignment> &) noexcept {}

    template <typename U>
    void construct(U * p)
        noexcept(std::is_nothrow_default_constructible_v<U>)
    {
        if constexpr (!std::is_trivially_default_constructible_v<U>)
        {
            ::new (static_cast<void *>(p)) U();
        }
    }

    template <typename U, typename A0, typename ...Args>
    void construct(U * p, A0 && a0, Args && ...args)
    {
        ::new (static_cast<void *>(p)) U(std::forward<A0>(a0),
            std::forward<Args>(args)...);
    }

    template <typename U>
    HEDLEY_NO_THROW
    void destroy(U * p) noexcept
    {
        if constexpr (!std::is_trivially_destructible_v<U>)
        {
            p->~U();
        }
    }
};

template <typename T, std::size_t A, typename U, std::size_t B>
HEDLEY_NO_THROW
constexpr bool operator==(const output_buffer_allocator<T, A> &,
    const output_buffer_allocator<U, B> &) noexcept
{
    return A == B;
}

template <typename T, std::size_t A, typename U, std::size_t B>
HEDLEY_NO_THROW
constexpr bool operator!=(const output_buffer_allocator<T, A> & lhs,
    const output_buffer_allocator<U, B> & rhs) noexcept
{
    return !(lhs == rhs);
}

/// @brief Move-only vector of `T`. Copy construction and copy assignment are
/// deleted. `at`, `operator[]`, `data`, iterators, and `size` are public.
/// @tparam T value type
/// @tparam Alignment allocation alignment
template <typename T,
          std::size_t Alignment = utils::max_align_v>
class output_buffer final
  : private std::vector<T, dpf::output_buffer_allocator<T, Alignment>>
{
  private:
    using vector = std::vector<T, dpf::output_buffer_allocator<T, Alignment>>;
  public:
    using value_type = typename vector::value_type;
    using iterator = typename vector::iterator;
    using const_iterator = typename vector::const_iterator;
    using size_type = typename vector::size_type;
    HEDLEY_NO_THROW
    output_buffer() noexcept = default;
    explicit output_buffer(size_type size) : vector(size) { }
    HEDLEY_NO_THROW
    output_buffer(output_buffer &&) noexcept = default;
    output_buffer(const output_buffer &) = delete;
    HEDLEY_NO_THROW
    output_buffer & operator=(output_buffer &&) noexcept = default;
    output_buffer & operator=(const output_buffer &) = delete;
    HEDLEY_NO_THROW
    ~output_buffer() noexcept = default;

    // "selectively public" inheritance
    using vector::at;
    using vector::operator[];
    using vector::data;
    using vector::begin;
    using vector::cbegin;
    using vector::end;
    using vector::cend;
    using vector::size;
};

template <>
class output_buffer<dpf::bit> : public dpf::dynamic_bit_array<>
{
  private:
    using size_type = typename dpf::dynamic_bit_array<>::size_type;
  public:
    explicit output_buffer(size_type size) : dynamic_bit_array(size) { }
    HEDLEY_NO_THROW
    output_buffer(output_buffer &&) noexcept = default;
    output_buffer(const output_buffer &) = delete;
    HEDLEY_NO_THROW
    output_buffer & operator=(output_buffer &&) noexcept = default;
    output_buffer & operator=(const output_buffer &) = delete;
    HEDLEY_NO_THROW
    ~output_buffer() noexcept = default;
};

template <>
class output_buffer<dpf::twobit> : public dpf::dynamic_packed_array<dpf::twobit>
{
    using base = dpf::dynamic_packed_array<dpf::twobit>;
  public:
    using size_type = typename base::size_type;
    explicit output_buffer(size_type size) : base(size) { }
    HEDLEY_NO_THROW
    output_buffer(output_buffer &&) noexcept = default;
    output_buffer(const output_buffer &) = delete;
    HEDLEY_NO_THROW
    output_buffer & operator=(output_buffer &&) noexcept = default;
    output_buffer & operator=(const output_buffer &) = delete;
    HEDLEY_NO_THROW
    ~output_buffer() noexcept = default;
};

template <>
class output_buffer<dpf::nyble> : public dpf::dynamic_packed_array<dpf::nyble>
{
    using base = dpf::dynamic_packed_array<dpf::nyble>;
  public:
    using size_type = typename base::size_type;
    explicit output_buffer(size_type size) : base(size) { }
    HEDLEY_NO_THROW
    output_buffer(output_buffer &&) noexcept = default;
    output_buffer(const output_buffer &) = delete;
    HEDLEY_NO_THROW
    output_buffer & operator=(output_buffer &&) noexcept = default;
    output_buffer & operator=(const output_buffer &) = delete;
    HEDLEY_NO_THROW
    ~output_buffer() noexcept = default;
};

#define LIBDPF_PACKED_SHARE_BUFFER(LANE, PARTY)                                   \
template <>                                                                       \
class output_buffer<subtractive_share<LANE, PARTY>>                              \
  : public packed_share_output<LANE, PARTY>                                      \
{                                                                                \
    using base = packed_share_output<LANE, PARTY>;                               \
  public:                                                                        \
    using size_type = typename base::size_type;                                  \
    explicit output_buffer(size_type size) : base(size) {}                       \
    HEDLEY_NO_THROW                                                              \
    output_buffer(output_buffer &&) noexcept = default;                          \
    output_buffer(const output_buffer &) = delete;                               \
    HEDLEY_NO_THROW                                                              \
    output_buffer & operator=(output_buffer &&) noexcept = default;              \
    output_buffer & operator=(const output_buffer &) = delete;                   \
    HEDLEY_NO_THROW                                                              \
    ~output_buffer() noexcept = default;                                         \
};

LIBDPF_PACKED_SHARE_BUFFER(dpf::twobit, 0);
LIBDPF_PACKED_SHARE_BUFFER(dpf::twobit, 1);
LIBDPF_PACKED_SHARE_BUFFER(dpf::nyble, 0);
LIBDPF_PACKED_SHARE_BUFFER(dpf::nyble, 1);
#undef LIBDPF_PACKED_SHARE_BUFFER

/// @brief Packed bit share buffers reuse the bit-array image; iterators yield shares.
#define LIBDPF_BIT_SHARE_BUFFER(PARTY)                                            \
template <>                                                                       \
class output_buffer<subtractive_share<dpf::bit, PARTY>>                          \
  : public dpf::dynamic_bit_array<>                                              \
{                                                                                \
  private:                                                                        \
    using size_type = typename dpf::dynamic_bit_array<>::size_type;              \
  public:                                                                         \
    explicit output_buffer(size_type size) : dynamic_bit_array(size) {}          \
    output_buffer(output_buffer &&) noexcept = default;                          \
    output_buffer(const output_buffer &) = delete;                               \
    output_buffer & operator=(output_buffer &&) noexcept = default;              \
    output_buffer & operator=(const output_buffer &) = delete;                   \
    ~output_buffer() noexcept = default;                                                  \
};
LIBDPF_BIT_SHARE_BUFFER(0);
LIBDPF_BIT_SHARE_BUFFER(1);
#undef LIBDPF_BIT_SHARE_BUFFER

/// @brief Buffer sized for the closed interval `[from, to]` of output `I`.
/// @details On a `party_key`, elements are subtractive shares of that output.
/// @tparam DpfKey DPF key type
/// @tparam I output index
/// @tparam InputT input domain type
/// @param from the inclusive start of the range
/// @param to the `to`
/// @return Buffer sized for the closed interval `[from, to]` of output `I`
template <typename DpfKey,
          std::size_t I = 0,
          typename InputT>
auto make_output_buffer_for_interval(InputT from, InputT to)
{
    using dpf_type = DpfKey;
    using output_type = typename DpfKey::concrete_output_type<I>;
    using buffer_elem = leaf_buffer_elem_t<DpfKey, output_type>;

    std::size_t nodes_in_interval = utils::get_leafnodes_in_output_interval<dpf_type>(from, to);
    return dpf::output_buffer<buffer_elem>(nodes_in_interval*dpf_type::outputs_per_leaf);
}

template <typename DpfKey,
          std::size_t I0,
          std::size_t I1,
          std::size_t ...Is,
          typename InputT>
auto make_output_buffer_for_interval(InputT from, InputT to)
{
    return std::make_tuple(
        make_output_buffer_for_interval<DpfKey, I0>(from, to),
        make_output_buffer_for_interval<DpfKey, I1>(from, to),
        make_output_buffer_for_interval<DpfKey, Is>(from, to)...);
}

/// @brief Buffer sized for the tree walk of `[from, to]` after `offset_x`.
/// @details Matches `make_basic_interval_memoizer(dpf, from, to)`: leaf packing
///          depends on alignment once a wildcard offset is ready.
template <std::size_t I = 0,
          typename DpfKey,
          typename InputT>
inline auto make_output_buffer_for_interval(const DpfKey & dpf, InputT from, InputT to)
{
    using input_type = typename DpfKey::input_type;
    if (dpf.offset_x.is_ready())
    {
        return make_output_buffer_for_interval<DpfKey, I>(
            dpf.offset_x(static_cast<input_type>(from)),
            dpf.offset_x(static_cast<input_type>(to)));
    }
    return make_output_buffer_for_interval<DpfKey, I>(from, to);
}

template <std::size_t I0,
          std::size_t I1,
          std::size_t ...Is,
          typename DpfKey,
          typename InputT>
inline auto make_output_buffer_for_interval(const DpfKey & dpf, InputT from, InputT to)
{
    using input_type = typename DpfKey::input_type;
    if (dpf.offset_x.is_ready())
    {
        const auto tfrom = dpf.offset_x(static_cast<input_type>(from));
        const auto tto = dpf.offset_x(static_cast<input_type>(to));
        return make_output_buffer_for_interval<DpfKey, I0, I1, Is...>(tfrom, tto);
    }
    return make_output_buffer_for_interval<DpfKey, I0, I1, Is...>(from, to);
}

/// @brief Buffer sized for every input of output `I`.
/// @tparam DpfKey DPF key type
/// @tparam I output index
/// @return Buffer sized for every input of output `I`
template <typename DpfKey,
          std::size_t I = 0>
auto make_output_buffer_for_full()
{
    using dpf_type = DpfKey;
    using input_type = typename dpf_type::input_type;

    return make_output_buffer_for_interval<dpf_type, I>(
        std::numeric_limits<input_type>::min(),
        std::numeric_limits<input_type>::max());
}

template <typename DpfKey,
          std::size_t I0,
          std::size_t I1,
          std::size_t ...Is>
auto make_output_buffer_for_full()
{
    return std::make_tuple(
        make_output_buffer_for_full<DpfKey, I0>(),
        make_output_buffer_for_full<DpfKey, I1>(),
        make_output_buffer_for_full<DpfKey, Is>()...);
}

template <std::size_t I = 0,
          typename DpfKey>
inline auto make_output_buffer_for_full(const DpfKey &)
{
    return make_output_buffer_for_full<DpfKey, I>();
}

template <std::size_t I0,
          std::size_t I1,
          std::size_t ...Is,
          typename DpfKey>
inline auto make_output_buffer_for_full(const DpfKey &)
{
    return make_output_buffer_for_full<DpfKey, I0, I1, Is...>();
}

template <typename DpfKey,
          std::size_t I = 0,
          typename ForwardIterator,
          typename ReturnType = return_entire_node_tag_>
auto make_output_buffer_for_subsequence(ForwardIterator begin, ForwardIterator end, ReturnType return_type = ReturnType{})
{
    (void)return_type;
    using dpf_type = DpfKey;
    using output_type = typename DpfKey::concrete_output_type<I>;
    using buffer_elem = leaf_buffer_elem_t<DpfKey, output_type>;
    std::size_t points_in_sequence = std::distance(begin, end);

    static_assert(std::is_same_v<ReturnType, return_entire_node_tag_> ||
                    std::is_same_v<ReturnType, return_output_only_tag_>);
    if constexpr(std::is_same_v<ReturnType, return_entire_node_tag_>)
    {
        return dpf::output_buffer<buffer_elem>(points_in_sequence*dpf_type::outputs_per_leaf);
    }
    else
    {
        if constexpr(std::is_same_v<typename DpfKey::concrete_output_type<0>, dpf::bit>)
        {
            auto tmp = dpf::output_buffer<buffer_elem>(points_in_sequence);
            tmp.unset();
            return std::move(tmp);
        }
        else
        {
            return dpf::output_buffer<buffer_elem>(points_in_sequence);
        }
    }
}

template <typename DpfKey,
          std::size_t I0,
          std::size_t I1,
          std::size_t ...Is,
          typename ForwardIterator,
          typename ReturnType = return_entire_node_tag_>
auto make_output_buffer_for_subsequence(ForwardIterator begin, ForwardIterator end, ReturnType return_type = ReturnType{})
{
    return std::make_tuple(
        make_output_buffer_for_subsequence<DpfKey, I0>(begin, end, return_type),
        make_output_buffer_for_subsequence<DpfKey, I1>(begin, end, return_type),
        make_output_buffer_for_subsequence<DpfKey, Is>(begin, end, return_type)...);
}

template <std::size_t I = 0,
          typename DpfKey,
          typename ForwardIterator,
          typename ReturnType = return_entire_node_tag_>
inline auto make_output_buffer_for_subsequence(const DpfKey &, ForwardIterator begin, ForwardIterator end, ReturnType return_type = ReturnType{})
{
    return make_output_buffer_for_subsequence<DpfKey, I>(begin, end, return_type);
}

template <std::size_t I0,
          std::size_t I1,
          std::size_t ...Is,
          typename DpfKey,
          typename ForwardIterator,
          typename ReturnType = return_entire_node_tag_>
inline auto make_output_buffer_for_subsequence(const DpfKey &, ForwardIterator begin, ForwardIterator end, ReturnType return_type = ReturnType{})
{
    return make_output_buffer_for_subsequence<DpfKey, I0, I1, Is...>(begin, end, return_type);
}

template <typename DpfKey,
          std::size_t I = 0,
          typename ReturnType = return_entire_node_tag_>
auto make_output_buffer_for_recipe_subsequence(const sequence_recipe & recipe, ReturnType return_type = ReturnType{})
{
    (void)return_type;
    using dpf_type = DpfKey;
    using output_type = typename DpfKey::concrete_output_type<I>;
    using buffer_elem = leaf_buffer_elem_t<DpfKey, output_type>;

    static_assert(std::is_same_v<ReturnType, return_entire_node_tag_> ||
                    std::is_same_v<ReturnType, return_output_only_tag_>);
    if constexpr(std::is_same_v<ReturnType, return_entire_node_tag_>)
    {
        return dpf::output_buffer<buffer_elem>(recipe.num_leaf_nodes()*dpf_type::outputs_per_leaf);
    }
    else
    {
        if constexpr(std::is_same_v<typename DpfKey::concrete_output_type<0>, dpf::bit>)
        {
            auto tmp = dpf::output_buffer<buffer_elem>(recipe.output_indices().size());
            tmp.unset();
            return std::move(tmp);
        }
        else
        {
            return dpf::output_buffer<buffer_elem>(recipe.output_indices().size());
        }
    }
}

template <typename DpfKey,
          std::size_t I0,
          std::size_t I1,
          std::size_t ...Is,
          typename ReturnType = return_entire_node_tag_>
inline auto make_output_buffer_for_recipe_subsequence(const sequence_recipe & recipe, ReturnType return_type = ReturnType{})
{
    return std::make_tuple(
        make_output_buffer_for_recipe_subsequence<DpfKey, I0>(recipe, return_type),
        make_output_buffer_for_recipe_subsequence<DpfKey, I1>(recipe, return_type),
        make_output_buffer_for_recipe_subsequence<DpfKey, Is>(recipe, return_type)...);
}

template <std::size_t I = 0,
          typename DpfKey,
          typename ReturnType = return_entire_node_tag_>
inline auto make_output_buffer_for_recipe_subsequence(const DpfKey &, const sequence_recipe & recipe, ReturnType return_type = ReturnType{})
{
    return make_output_buffer_for_recipe_subsequence<DpfKey, I>(recipe, return_type);
}

template <std::size_t I0,
          std::size_t I1,
          std::size_t ...Is,
          typename DpfKey,
          typename ReturnType = return_entire_node_tag_>
inline auto make_output_buffer_for_recipe_subsequence(const DpfKey &, const sequence_recipe & recipe, ReturnType return_type = ReturnType{})
{
    return make_output_buffer_for_recipe_subsequence<DpfKey, I0, I1, Is...>(recipe, return_type);
}

namespace utils
{

template <>
struct is_bit_array<output_buffer<bit>> : std::true_type {};

}  // namespace utils

}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_OUTPUT_BUFFER_HPP__