/// @file dpf/eval_common.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_EVAL_COMMON_HPP__
#define LIBDPF_INCLUDE_DPF_EVAL_COMMON_HPP__

#include "hedley/hedley.h"

#include <cstddef>
#include <cstring>
#include <limits>
#include <stdexcept>
#include <type_traits>

#include "dpf/leaf_node.hpp"
#include "dpf/secret_share.hpp"

namespace dpf
{

/// Sentinel: `dpf_output` converts to a bare `OutputT` (no party tag).
inline constexpr std::size_t no_party = std::numeric_limits<std::size_t>::max();

/// Eval result type for a leaf output of `KeyT`: subtractive share when the
/// key is party-tagged, otherwise the concrete output type.
template <typename KeyT, typename OutputT, bool = is_party_key_v<KeyT>>
struct eval_leaf_result
{
    using type = OutputT;
};
template <typename KeyT, typename OutputT>
struct eval_leaf_result<KeyT, OutputT, true>
{
    using type = subtractive_share<OutputT, party_of_v<KeyT>>;
};
template <typename KeyT, typename OutputT>
using eval_leaf_result_t = typename eval_leaf_result<KeyT, OutputT>::type;

/// Eval result type for a comparison output of `KeyT`: additive share when
/// the key is party-tagged, otherwise `Beta`.
template <typename KeyT, typename Beta, bool = is_party_key_v<KeyT>>
struct eval_cmp_result
{
    using type = Beta;
};
template <typename KeyT, typename Beta>
struct eval_cmp_result<KeyT, Beta, true>
{
    using type = additive_share<Beta, party_of_v<KeyT>>;
};
template <typename KeyT, typename Beta>
using eval_cmp_result_t = typename eval_cmp_result<KeyT, Beta>::type;

template <typename OutputT,
          typename NodeT,
          std::size_t Party = no_party>
struct alignas(utils::max_align_v) dpf_output
{
    using result_type = std::conditional_t<
        Party == no_party,
        OutputT,
        subtractive_share<OutputT, Party>>;

    dpf_output(const dpf_output &) = default;
    HEDLEY_NO_THROW
    dpf_output(dpf_output &&) noexcept = default;
    dpf_output & operator=(const dpf_output &) = default;
    HEDLEY_NO_THROW
    dpf_output & operator=(dpf_output &&) noexcept = default;
    ~dpf_output() = default;

    HEDLEY_ALWAYS_INLINE
    HEDLEY_PURE
    constexpr operator result_type() const
    {
        OutputT v = extract_leaf<NodeT, OutputT>(node, offset);
        if constexpr (Party == no_party)
            return v;
        else
            return subtractive_share<OutputT, Party>::from_raw(v);
    }

    HEDLEY_ALWAYS_INLINE
    HEDLEY_PURE
    auto operator*() const
    {
        return static_cast<result_type>(*this);
    }

    NodeT node;
    std::size_t offset;

  private:
    dpf_output(NodeT leaf_node, std::size_t off)
      : node{leaf_node}, offset{off} { }

  public:
    template <typename Output,
              typename Input,
              typename Node>
    friend auto make_dpf_output(const Node & node, Input x);

    template <std::size_t P,
              typename Output,
              typename Input,
              typename Node>
    friend auto make_dpf_output(const Node & node, Input x);
};

/// Copy one packed leaf into a buffer whose element type may differ
/// from `LeafT` (bit arrays store `word_type`, not the exterior node).
/// Byte destination keeps the store free of strict-aliasing UB.
template <typename LeafT, typename Buffer>
HEDLEY_ALWAYS_INLINE
HEDLEY_NO_THROW
void store_leaf_bytes(Buffer && buf, std::size_t index, const LeafT & leaf) noexcept
{
    auto * dst = reinterpret_cast<unsigned char *>(utils::data(buf))
        + index * sizeof(LeafT);
    std::memcpy(dst, std::addressof(leaf), sizeof(LeafT));
}

template <typename Output,
          typename Input,
          typename Node>
auto make_dpf_output(const Node & node, Input x)
{
    return dpf_output<concrete_type_t<Output>, Node>{node,
        offset_within_block<concrete_type_t<Output>, Node>(x)};
}

template <std::size_t Party,
          typename Output,
          typename Input,
          typename Node>
auto make_dpf_output(const Node & node, Input x)
{
    return dpf_output<concrete_type_t<Output>, Node, Party>{node,
        offset_within_block<concrete_type_t<Output>, Node>(x)};
}

/// Wrap a raw leaf node into a party-tagged `dpf_output` when `KeyT` is a
/// `party_key`, otherwise a bare `dpf_output`.
template <typename KeyT, typename Output, typename Input, typename Node>
auto make_eval_dpf_output(const Node & node, Input x)
{
    if constexpr (is_party_key_v<KeyT>)
        return make_dpf_output<party_of_v<KeyT>, Output>(node, x);
    else
        return make_dpf_output<Output>(node, x);
}

/// Wrap a raw comparison `Beta` value as an additive share when `KeyT` is a
/// `party_key`.
template <typename KeyT, typename Beta>
HEDLEY_NO_THROW
auto make_eval_cmp_result(Beta raw) noexcept
{
    if constexpr (is_party_key_v<KeyT>)
        return additive_share<Beta, party_of_v<KeyT>>::from_raw(raw);
    else
        return raw;
}

}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_EVAL_COMMON_HPP__