/// @file dpf/vec.hpp
/// @brief `dpf::vec<T, N>`, one output of `N` lanes added componentwise.
/// @details Each lane lives in the ring of `T` (`+` and `-` wrap the way `T`
///          wraps). A leaf adds every lane and does not carry from one lane
///          into the next. `T` is an ordinary output type: an integer,
///          `modint`, `fixedpoint`, `twobit`, `nyble`, or `xor_wrapper`.
/// @copyright Copyright (c) 2019-2026 Ryan Henry and [others](@ref authors)
/// @license Released under a GNU General Public v2.0 (GPLv2) license.

#ifndef LIBDPF_INCLUDE_DPF_VEC_HPP__
#define LIBDPF_INCLUDE_DPF_VEC_HPP__

#include <array>
#include <cstddef>
#include <cstring>
#include <type_traits>
#include <utility>

#include "hedley/hedley.h"

#include "dpf/leaf_arithmetic.hpp"
#include "dpf/utils.hpp"

namespace dpf
{

/// @brief `N` lanes of `T`, combined componentwise with no carry between lanes.
/// @tparam T lane type. Its `+`, `-`, and `*` are used per lane
/// @tparam N number of lanes
template <typename T, std::size_t N>
struct vec
{
    static_assert(N > 0, "dpf::vec needs at least one lane");
    static constexpr bool dpf_vec = true;
    /// @brief Number of lanes.
    static constexpr std::size_t lane_count = N;
    /// @brief Type of one lane.
    using lane_type = T;

    /// @brief Lane storage, index 0 in the least-significant lane.
    std::array<T, N> lanes{};

    /// @brief Value-initialize every lane.
    HEDLEY_NO_THROW
    constexpr vec() noexcept = default;

    /// @brief Mutable lane `i`.
    /// @param i the lane index
    /// @return a reference to that lane
    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    constexpr T & operator[](std::size_t i) noexcept { return lanes[i]; }

    /// @brief Lane `i`.
    /// @param i the lane index
    /// @return a reference to that lane
    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    constexpr const T & operator[](std::size_t i) const noexcept { return lanes[i]; }

    /// @brief Negate every lane.
    /// @return the negated vector
    HEDLEY_ALWAYS_INLINE
    constexpr vec operator-() const
        noexcept(noexcept(static_cast<T>(-std::declval<const T &>())))
    {
        vec out;
        for (std::size_t i = 0; i < N; ++i)
            out.lanes[i] = static_cast<T>(-lanes[i]);
        return out;
    }

    /// @brief Add each lane, with no carry into the next lane.
    /// @param rhs the right-hand vector
    /// @return the lane-wise sum
    HEDLEY_ALWAYS_INLINE
    constexpr vec operator+(const vec & rhs) const
        noexcept(noexcept(std::declval<const T &>() + std::declval<const T &>()))
    {
        vec out;
        for (std::size_t i = 0; i < N; ++i)
            out.lanes[i] = static_cast<T>(lanes[i] + rhs.lanes[i]);
        return out;
    }

    /// @brief Subtract each lane, with no borrow from the next lane.
    /// @param rhs the right-hand vector
    /// @return the lane-wise difference
    HEDLEY_ALWAYS_INLINE
    constexpr vec operator-(const vec & rhs) const
        noexcept(noexcept(std::declval<const T &>() - std::declval<const T &>()))
    {
        vec out;
        for (std::size_t i = 0; i < N; ++i)
            out.lanes[i] = static_cast<T>(lanes[i] - rhs.lanes[i]);
        return out;
    }

    /// @brief Multiply each lane.
    /// @param rhs the right-hand vector
    /// @return the lane-wise product
    HEDLEY_ALWAYS_INLINE
    constexpr vec operator*(const vec & rhs) const
        noexcept(noexcept(std::declval<const T &>() * std::declval<const T &>()))
    {
        vec out;
        for (std::size_t i = 0; i < N; ++i)
            out.lanes[i] = static_cast<T>(lanes[i] * rhs.lanes[i]);
        return out;
    }

    /// @brief Lane-wise equality.
    /// @param rhs the right-hand vector
    /// @return `true` when every lane matches
    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    constexpr bool operator==(const vec & rhs) const noexcept
    {
        for (std::size_t i = 0; i < N; ++i)
        {
            if (!(lanes[i] == rhs.lanes[i]))
                return false;
        }
        return true;
    }

    /// @brief Lane-wise inequality.
    /// @param rhs the right-hand vector
    /// @return `true` when some lane differs
    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    constexpr bool operator!=(const vec & rhs) const noexcept
    {
        return !(*this == rhs);
    }
};

namespace utils
{

template <typename T, std::size_t N>
struct bitlength_of<dpf::vec<T, N>>
  : std::integral_constant<std::size_t, bitlength_of_v<T> * N> {};

}  // namespace utils

namespace leaf_arithmetic
{
namespace vec_detail
{

template <typename U>
struct is_std_array : std::false_type {};
template <typename U, std::size_t M>
struct is_std_array<std::array<U, M>> : std::true_type {};

/// @brief Add or subtract every stored lane. A node wider than the vector is a
/// sequence of vectors, then a tail of leftover lanes. Bytes that do not
/// fill a lane are XOR-combined so a random pad still cancels.
/// @tparam T lane type
/// @tparam N number of lanes
/// @tparam Op lane operation
/// @param dst the destination
/// @param a the `a`
/// @param b the `b`
/// @param nbytes the number of bytes
/// @param op the `op`
template <typename T, std::size_t N, typename Op>
HEDLEY_ALWAYS_INLINE
void apply_bytes(unsigned char * dst, const unsigned char * a,
    const unsigned char * b, std::size_t nbytes, Op op)
{
    using V = dpf::vec<T, N>;
    std::size_t off = 0;
    const std::size_t nvec = nbytes / sizeof(V);
    for (std::size_t i = 0; i < nvec; ++i)
    {
        V va, vb;
        std::memcpy(&va, a + off, sizeof(V));
        std::memcpy(&vb, b + off, sizeof(V));
        V vc = op(va, vb);
        std::memcpy(dst + off, &vc, sizeof(V));
        off += sizeof(V);
    }
    while (off + sizeof(T) <= nbytes)
    {
        T va, vb;
        std::memcpy(&va, a + off, sizeof(T));
        std::memcpy(&vb, b + off, sizeof(T));
        T vc = op(va, vb);
        std::memcpy(dst + off, &vc, sizeof(T));
        off += sizeof(T);
    }
    for (; off < nbytes; ++off)
        dst[off] = static_cast<unsigned char>(a[off] ^ b[off]);
}

template <typename T, std::size_t N, typename NodeT, typename Op>
HEDLEY_ALWAYS_INLINE
NodeT apply_node(const NodeT & a, const NodeT & b, Op op)
{
    alignas(NodeT) unsigned char ca[sizeof(NodeT)];
    alignas(NodeT) unsigned char cb[sizeof(NodeT)];
    alignas(NodeT) unsigned char cc[sizeof(NodeT)];
    std::memcpy(ca, &a, sizeof(NodeT));
    std::memcpy(cb, &b, sizeof(NodeT));
    apply_bytes<T, N>(cc, ca, cb, sizeof(NodeT), op);
    NodeT out;
    std::memcpy(&out, cc, sizeof(NodeT));
    return out;
}

}  // namespace vec_detail

template <typename T, std::size_t N, typename NodeT>
struct add_t<dpf::vec<T, N>, NodeT,
    std::enable_if_t<!std::is_void_v<NodeT>
        && !vec_detail::is_std_array<NodeT>::value>>
{
    HEDLEY_ALWAYS_INLINE
    auto operator()(const NodeT & a, const NodeT & b) const
    {
        return vec_detail::apply_node<T, N>(a, b,
            [](const auto & x, const auto & y) { return x + y; });
    }
};

template <typename T, std::size_t N, typename Elem, std::size_t K>
struct add_t<dpf::vec<T, N>, std::array<Elem, K>>
{
    auto operator()(const std::array<Elem, K> & a, const std::array<Elem, K> & b) const
    {
        std::array<Elem, K> c{};
        auto * dst = reinterpret_cast<unsigned char *>(c.data());
        vec_detail::apply_bytes<T, N>(dst,
            reinterpret_cast<const unsigned char *>(a.data()),
            reinterpret_cast<const unsigned char *>(b.data()),
            sizeof(c),
            [](const auto & x, const auto & y) { return x + y; });
        return c;
    }
};

template <typename T, std::size_t N, typename NodeT>
struct subtract_t<dpf::vec<T, N>, NodeT,
    std::enable_if_t<!std::is_void_v<NodeT>
        && !vec_detail::is_std_array<NodeT>::value>>
{
    HEDLEY_ALWAYS_INLINE
    auto operator()(const NodeT & a, const NodeT & b) const
    {
        return vec_detail::apply_node<T, N>(a, b,
            [](const auto & x, const auto & y) { return x - y; });
    }
};

template <typename T, std::size_t N, typename Elem, std::size_t K>
struct subtract_t<dpf::vec<T, N>, std::array<Elem, K>>
{
    auto operator()(const std::array<Elem, K> & a, const std::array<Elem, K> & b) const
    {
        std::array<Elem, K> c{};
        auto * dst = reinterpret_cast<unsigned char *>(c.data());
        vec_detail::apply_bytes<T, N>(dst,
            reinterpret_cast<const unsigned char *>(a.data()),
            reinterpret_cast<const unsigned char *>(b.data()),
            sizeof(c),
            [](const auto & x, const auto & y) { return x - y; });
        return c;
    }
};

}  // namespace leaf_arithmetic

}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_VEC_HPP__