discrete_vector.hpp

// Copyright (C) The DDC development team, see COPYRIGHT.md file
//
// SPDX-License-Identifier: MIT
 
#pragma once
 
#include <array>
#include <concepts>
#include <cstddef>
#include <iosfwd>
#include <span>
#include <type_traits>
#include <utility>
 
#include <Kokkos_Macros.hpp>
 
#include "detail/macros.hpp"
#include "detail/utils.hpp"
 
#include "discrete_dimension.hpp"
#include "type_seq.hpp"
 
namespace ddc {
 
template <concepts::discrete_dimension...>
class DiscreteVector;
 
template <class T>
struct is_discrete_vector : std::false_type
{
};
 
template <concepts::discrete_dimension... Tags>
struct is_discrete_vector<DiscreteVector<Tags...>> : std::true_type
{
};
 
template <class T>
inline constexpr bool is_discrete_vector_v = is_discrete_vector<T>::value;
 
namespace concepts {
 
template <class T>
concept discrete_vector = is_discrete_vector_v<T>;
 
}
 
namespace detail {
 
template <ddc::concepts::discrete_dimension... Tags>
struct ToTypeSeq<DiscreteVector<Tags...>>
{
    using type = ddc::TypeSeq<Tags...>;
};
 
} // namespace detail
 
/** A DiscreteVectorElement is a scalar that represents the difference between two coordinates.
 */
using DiscreteVectorElement = std::ptrdiff_t;
 
template <concepts::discrete_dimension QueryTag, concepts::discrete_dimension... Tags>
KOKKOS_FUNCTION constexpr DiscreteVectorElement const& get(
        DiscreteVector<Tags...> const& tuple) noexcept
{
    return tuple.template get<QueryTag>();
}
 
template <concepts::discrete_dimension QueryTag, concepts::discrete_dimension... Tags>
KOKKOS_FUNCTION constexpr DiscreteVectorElement& get(DiscreteVector<Tags...>& tuple) noexcept
{
    return tuple.template get<QueryTag>();
}
 
template <concepts::discrete_dimension QueryTag, concepts::discrete_dimension... Tags>
KOKKOS_FUNCTION constexpr DiscreteVectorElement const& get_or(
        DiscreteVector<Tags...> const& tuple,
        DiscreteVectorElement const& default_value) noexcept
{
    return tuple.template get_or<QueryTag>(default_value);
}
 
template <concepts::discrete_dimension QueryTag, concepts::discrete_dimension... Tags>
KOKKOS_FUNCTION constexpr DiscreteVector<QueryTag> select_or(
        DiscreteVector<Tags...> const& arr,
        DiscreteVector<QueryTag> const& default_value) noexcept
{
    if constexpr (in_tags_v<QueryTag, TypeSeq<Tags...>>) {
        return DiscreteVector<QueryTag>(arr);
    } else {
        return default_value;
    }
}
 
/// Unary operators: +, -
 
template <concepts::discrete_dimension... Tags>
KOKKOS_FUNCTION constexpr DiscreteVector<Tags...> operator+(DiscreteVector<Tags...> const& x)
{
    return x;
}
 
template <concepts::discrete_dimension... Tags>
KOKKOS_FUNCTION constexpr DiscreteVector<Tags...> operator-(DiscreteVector<Tags...> const& x)
{
    return DiscreteVector<Tags...>((-get<Tags>(x))...);
}
 
/// Internal binary operators: +, -
 
template <concepts::discrete_dimension... Tags, concepts::discrete_dimension... OTags>
KOKKOS_FUNCTION constexpr auto operator+(
        DiscreteVector<Tags...> const& lhs,
        DiscreteVector<OTags...> const& rhs)
{
    if constexpr (sizeof...(Tags) >= sizeof...(OTags)) {
        static_assert((type_seq_contains_v<TypeSeq<OTags>, TypeSeq<Tags...>> && ...));
        DiscreteVector<Tags...> result(lhs);
        result += rhs;
        return result;
    } else {
        static_assert((type_seq_contains_v<TypeSeq<Tags>, TypeSeq<OTags...>> && ...));
        DiscreteVector<OTags...> result(rhs);
        result += lhs;
        return result;
    }
}
 
template <concepts::discrete_dimension... Tags, concepts::discrete_dimension... OTags>
KOKKOS_FUNCTION constexpr auto operator-(
        DiscreteVector<Tags...> const& lhs,
        DiscreteVector<OTags...> const& rhs)
{
    if constexpr (sizeof...(Tags) >= sizeof...(OTags)) {
        static_assert((type_seq_contains_v<TypeSeq<OTags>, TypeSeq<Tags...>> && ...));
        DiscreteVector<Tags...> result(lhs);
        result -= rhs;
        return result;
    } else {
        static_assert((type_seq_contains_v<TypeSeq<Tags>, TypeSeq<OTags...>> && ...));
        DiscreteVector<OTags...> result(-rhs);
        result += lhs;
        return result;
    }
}
 
template <concepts::discrete_dimension Tag, std::integral IntegralType>
KOKKOS_FUNCTION constexpr DiscreteVector<Tag> operator+(
        DiscreteVector<Tag> const& lhs,
        IntegralType const& rhs)
{
    return DiscreteVector<Tag>(get<Tag>(lhs) + rhs);
}
 
template <std::integral IntegralType, concepts::discrete_dimension Tag>
KOKKOS_FUNCTION constexpr DiscreteVector<Tag> operator+(
        IntegralType const& lhs,
        DiscreteVector<Tag> const& rhs)
{
    return DiscreteVector<Tag>(lhs + get<Tag>(rhs));
}
 
template <concepts::discrete_dimension Tag, std::integral IntegralType>
KOKKOS_FUNCTION constexpr DiscreteVector<Tag> operator-(
        DiscreteVector<Tag> const& lhs,
        IntegralType const& rhs)
{
    return DiscreteVector<Tag>(get<Tag>(lhs) - rhs);
}
 
template <std::integral IntegralType, concepts::discrete_dimension Tag>
KOKKOS_FUNCTION constexpr DiscreteVector<Tag> operator-(
        IntegralType const& lhs,
        DiscreteVector<Tag> const& rhs)
{
    return DiscreteVector<Tag>(lhs - get<Tag>(rhs));
}
 
/// external left binary operator: *
 
template <std::integral IntegralType, concepts::discrete_dimension... Tags>
KOKKOS_FUNCTION constexpr auto operator*(
        IntegralType const& lhs,
        DiscreteVector<Tags...> const& rhs)
{
    return DiscreteVector<Tags...>((lhs * get<Tags>(rhs))...);
}
 
template <concepts::discrete_dimension... QueryTags, concepts::discrete_dimension... Tags>
KOKKOS_FUNCTION constexpr DiscreteVector<QueryTags...> select(
        DiscreteVector<Tags...> const& arr) noexcept
{
    return DiscreteVector<QueryTags...>(arr);
}
 
template <concepts::discrete_dimension... QueryTags, concepts::discrete_dimension... Tags>
KOKKOS_FUNCTION constexpr DiscreteVector<QueryTags...> select(
        DiscreteVector<Tags...>&& arr) noexcept
{
    return DiscreteVector<QueryTags...>(std::move(arr));
}
 
/// Returns a reference towards the DiscreteVector that contains the QueryTag
template <
        concepts::discrete_dimension QueryTag,
        concepts::discrete_vector HeadDVect,
        concepts::discrete_vector... TailDVects>
KOKKOS_FUNCTION constexpr auto const& take(HeadDVect const& head, TailDVects const&... tail)
{
    DDC_IF_NVCC_THEN_PUSH_AND_SUPPRESS(implicit_return_from_non_void_function)
    if constexpr (type_seq_contains_v<TypeSeq<QueryTag>, to_type_seq_t<HeadDVect>>) {
        static_assert(
                (!type_seq_contains_v<TypeSeq<QueryTag>, to_type_seq_t<TailDVects>> && ...),
                "ERROR: tag redundant");
        return head;
    } else {
        static_assert(sizeof...(TailDVects) > 0, "ERROR: tag not found");
        return take<QueryTag>(tail...);
    }
    DDC_IF_NVCC_THEN_POP
}
 
namespace detail {
 
template <class T>
class DiscreteVectorConversionOperators
{
};
 
template <ddc::concepts::discrete_dimension Tag>
class DiscreteVectorConversionOperators<DiscreteVector<Tag>>
{
public:
    // NOLINTBEGIN(google-explicit-constructor)
    KOKKOS_FUNCTION constexpr operator DiscreteVectorElement const&() const noexcept
    {
        return static_cast<DiscreteVector<Tag> const*>(this)->m_values[0];
    }
 
    KOKKOS_FUNCTION constexpr operator DiscreteVectorElement&() noexcept
    {
        return static_cast<DiscreteVector<Tag>*>(this)->m_values[0];
    }
    // NOLINTEND(google-explicit-constructor)
};
 
/// Returns a reference to the underlying `std::array`
template <ddc::concepts::discrete_dimension... Tags>
KOKKOS_FUNCTION constexpr std::array<DiscreteVectorElement, sizeof...(Tags)>& array(
        DiscreteVector<Tags...>& v) noexcept
{
    return v.m_values;
}
 
/// Returns a reference to the underlying `std::array`
template <ddc::concepts::discrete_dimension... Tags>
KOKKOS_FUNCTION constexpr std::array<DiscreteVectorElement, sizeof...(Tags)> const& array(
        DiscreteVector<Tags...> const& v) noexcept
{
    return v.m_values;
}
 
} // namespace detail
 
/** A DiscreteVector is a vector in the discrete dimension
 *
 * Each is tagged by its associated dimensions.
 */
template <concepts::discrete_dimension... Tags>
class DiscreteVector : public detail::DiscreteVectorConversionOperators<DiscreteVector<Tags...>>
{
    friend class detail::DiscreteVectorConversionOperators<DiscreteVector<Tags...>>;
 
    friend KOKKOS_FUNCTION constexpr std::array<DiscreteVectorElement, sizeof...(Tags)>& detail::
            array<Tags...>(DiscreteVector<Tags...>& v) noexcept;
    friend KOKKOS_FUNCTION constexpr std::array<DiscreteVectorElement, sizeof...(Tags)> const&
    detail::array<Tags...>(DiscreteVector<Tags...> const& v) noexcept;
 
    using tags_seq = TypeSeq<Tags...>;
 
    static_assert(
            type_seq_is_unique_v<tags_seq>,
            "The dimensions of a DiscreteVector must be unique");
 
private:
    std::array<DiscreteVectorElement, sizeof...(Tags)> m_values;
 
public:
    using value_type = DiscreteVectorElement;
 
    static KOKKOS_FUNCTION constexpr std::size_t size() noexcept
    {
        return sizeof...(Tags);
    }
 
public:
    KOKKOS_DEFAULTED_FUNCTION constexpr DiscreteVector() = default;
 
    KOKKOS_DEFAULTED_FUNCTION constexpr DiscreteVector(DiscreteVector const&) = default;
 
    KOKKOS_DEFAULTED_FUNCTION constexpr DiscreteVector(DiscreteVector&&) = default;
 
    template <concepts::discrete_vector... DVects>
    KOKKOS_FUNCTION constexpr explicit DiscreteVector(DVects const&... delems) noexcept
        : m_values {take<Tags>(delems...).template get<Tags>()...}
    {
    }
 
    template <std::convertible_to<DiscreteVectorElement> IntegerType>
    KOKKOS_FUNCTION constexpr explicit DiscreteVector(
            std::array<IntegerType, sizeof...(Tags)> const& values) noexcept
        : m_values(
                  detail::convert_array_to<DiscreteVectorElement>(
                          values,
                          std::make_index_sequence<sizeof...(Tags)>()))
    {
    }
 
    template <std::convertible_to<DiscreteVectorElement>... Params>
    KOKKOS_FUNCTION constexpr explicit DiscreteVector(Params const&... params) noexcept
        requires((!is_discrete_vector_v<Params> && ...) && sizeof...(Params) == sizeof...(Tags))
        : m_values {static_cast<DiscreteVectorElement>(params)...}
    {
    }
 
    KOKKOS_DEFAULTED_FUNCTION ~DiscreteVector() = default;
 
    KOKKOS_DEFAULTED_FUNCTION constexpr DiscreteVector& operator=(DiscreteVector const& other)
            = default;
 
    KOKKOS_DEFAULTED_FUNCTION constexpr DiscreteVector& operator=(DiscreteVector&& other) = default;
 
    template <concepts::discrete_dimension... OTags>
    KOKKOS_FUNCTION constexpr bool operator==(DiscreteVector<OTags...> const& rhs) const noexcept
    {
        return ((m_values[type_seq_rank_v<Tags, tags_seq>] == rhs.template get<Tags>()) && ...);
    }
 
#if !defined(__cpp_impl_three_way_comparison) || __cpp_impl_three_way_comparison < 201902L
    // In C++20, `a!=b` shall be automatically translated by the compiler to `!(a==b)`
    template <concepts::discrete_dimension... OTags>
    KOKKOS_FUNCTION constexpr bool operator!=(DiscreteVector<OTags...> const& rhs) const noexcept
    {
        return !(*this == rhs);
    }
#endif
 
    template <concepts::discrete_dimension QueryTag>
    KOKKOS_FUNCTION constexpr DiscreteVectorElement& get() noexcept
    {
        static_assert(in_tags_v<QueryTag, tags_seq>, "requested Tag absent from DiscreteVector");
        return m_values[type_seq_rank_v<QueryTag, tags_seq>];
    }
 
    template <concepts::discrete_dimension QueryTag>
    KOKKOS_FUNCTION constexpr DiscreteVectorElement const& get() const noexcept
    {
        static_assert(in_tags_v<QueryTag, tags_seq>, "requested Tag absent from DiscreteVector");
        return m_values[type_seq_rank_v<QueryTag, tags_seq>];
    }
 
    template <concepts::discrete_dimension QueryTag>
    KOKKOS_FUNCTION constexpr DiscreteVectorElement const& get_or(
            DiscreteVectorElement const& default_value) const&
    {
        DDC_IF_NVCC_THEN_PUSH_AND_SUPPRESS(implicit_return_from_non_void_function)
        if constexpr (in_tags_v<QueryTag, tags_seq>) {
            return m_values[type_seq_rank_v<QueryTag, tags_seq>];
        } else {
            return default_value;
        }
        DDC_IF_NVCC_THEN_POP
    }
 
    KOKKOS_FUNCTION constexpr DiscreteVectorElement const& value() const noexcept
        requires(sizeof...(Tags) == 1)
    {
        return m_values[0];
    }
 
    KOKKOS_FUNCTION constexpr DiscreteVector& operator++()
        requires(sizeof...(Tags) == 1)
    {
        ++m_values[0];
        return *this;
    }
 
    KOKKOS_FUNCTION constexpr DiscreteVector operator++(int)
        requires(sizeof...(Tags) == 1)
    {
        DiscreteVector const tmp = *this;
        ++m_values[0];
        return tmp;
    }
 
    KOKKOS_FUNCTION constexpr DiscreteVector& operator--()
        requires(sizeof...(Tags) == 1)
    {
        --m_values[0];
        return *this;
    }
 
    KOKKOS_FUNCTION constexpr DiscreteVector operator--(int)
        requires(sizeof...(Tags) == 1)
    {
        DiscreteVector const tmp = *this;
        --m_values[0];
        return tmp;
    }
 
    template <concepts::discrete_dimension... OTags>
    KOKKOS_FUNCTION constexpr DiscreteVector& operator+=(DiscreteVector<OTags...> const& rhs)
    {
        static_assert((type_seq_contains_v<TypeSeq<OTags>, tags_seq> && ...));
        ((m_values[type_seq_rank_v<OTags, tags_seq>] += rhs.template get<OTags>()), ...);
        return *this;
    }
 
    template <std::integral IntegralType>
    KOKKOS_FUNCTION constexpr DiscreteVector& operator+=(IntegralType const& rhs)
        requires(sizeof...(Tags) == 1)
    {
        m_values[0] += rhs;
        return *this;
    }
 
    template <concepts::discrete_dimension... OTags>
    KOKKOS_FUNCTION constexpr DiscreteVector& operator-=(DiscreteVector<OTags...> const& rhs)
    {
        static_assert((type_seq_contains_v<TypeSeq<OTags>, tags_seq> && ...));
        ((m_values[type_seq_rank_v<OTags, tags_seq>] -= rhs.template get<OTags>()), ...);
        return *this;
    }
 
    template <std::integral IntegralType>
    KOKKOS_FUNCTION constexpr DiscreteVector& operator-=(IntegralType const& rhs)
        requires(sizeof...(Tags) == 1)
    {
        m_values[0] -= rhs;
        return *this;
    }
 
    template <concepts::discrete_dimension... OTags>
    KOKKOS_FUNCTION constexpr DiscreteVector& operator*=(DiscreteVector<OTags...> const& rhs)
    {
        static_assert((type_seq_contains_v<TypeSeq<OTags>, tags_seq> && ...));
        ((m_values[type_seq_rank_v<OTags, tags_seq>] *= rhs.template get<OTags>()), ...);
        return *this;
    }
};
 
template <concepts::discrete_dimension Tag>
KOKKOS_FUNCTION constexpr bool operator<(
        DiscreteVector<Tag> const& lhs,
        DiscreteVector<Tag> const& rhs)
{
    return lhs.value() < rhs.value();
}
 
template <concepts::discrete_dimension Tag, class IntegralType>
KOKKOS_FUNCTION constexpr bool operator<(DiscreteVector<Tag> const& lhs, IntegralType const& rhs)
{
    return lhs.value() < rhs;
}
 
namespace detail {
 
void print_discrete_vector(std::ostream& os, std::span<DiscreteVectorElement const> view);
 
} // namespace detail
 
template <concepts::discrete_dimension... Tags>
std::ostream& operator<<(std::ostream& os, DiscreteVector<Tags...> const& arr)
{
    detail::print_discrete_vector(os, detail::array(arr));
    return os;
}
 
} // namespace ddc