discrete_element.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 "discrete_vector.hpp"
#include "type_seq.hpp"
 
namespace ddc {
 
template <concepts::discrete_dimension...>
class DiscreteElement;
 
template <class T>
struct is_discrete_element : std::false_type
{
};
 
template <concepts::discrete_dimension... Tags>
struct is_discrete_element<DiscreteElement<Tags...>> : std::true_type
{
};
 
template <class T>
inline constexpr bool is_discrete_element_v = is_discrete_element<T>::value;
 
namespace concepts {
 
template <class T>
concept discrete_element = is_discrete_element_v<T>;
 
}
 
namespace detail {
 
template <ddc::concepts::discrete_dimension... Tags>
struct ToTypeSeq<DiscreteElement<Tags...>>
{
    using type = ddc::TypeSeq<Tags...>;
};
 
} // namespace detail
 
/** A DiscreteCoordElement is a scalar that identifies an element of the discrete dimension
 */
using DiscreteElementType = std::size_t;
 
template <concepts::discrete_dimension Tag>
KOKKOS_FUNCTION constexpr DiscreteElementType const& uid(DiscreteElement<Tag> const& tuple) noexcept
{
    return tuple.uid();
}
 
template <concepts::discrete_dimension Tag>
KOKKOS_FUNCTION constexpr DiscreteElementType& uid(DiscreteElement<Tag>& tuple) noexcept
{
    return tuple.uid();
}
 
template <concepts::discrete_dimension QueryTag, concepts::discrete_dimension... Tags>
KOKKOS_FUNCTION constexpr DiscreteElementType const& uid(
        DiscreteElement<Tags...> const& tuple) noexcept
{
    return tuple.template uid<QueryTag>();
}
 
template <concepts::discrete_dimension QueryTag, concepts::discrete_dimension... Tags>
KOKKOS_FUNCTION constexpr DiscreteElementType& uid(DiscreteElement<Tags...>& tuple) noexcept
{
    return tuple.template uid<QueryTag>();
}
 
template <concepts::discrete_dimension QueryTag, concepts::discrete_dimension... Tags>
KOKKOS_FUNCTION constexpr DiscreteElement<QueryTag> select_or(
        DiscreteElement<Tags...> const& arr,
        DiscreteElement<QueryTag> const& default_value) noexcept
{
    if constexpr (in_tags_v<QueryTag, TypeSeq<Tags...>>) {
        return DiscreteElement<QueryTag>(arr);
    } else {
        return default_value;
    }
}
 
template <concepts::discrete_dimension... QueryTags, concepts::discrete_dimension... Tags>
KOKKOS_FUNCTION constexpr DiscreteElement<QueryTags...> select(
        DiscreteElement<Tags...> const& arr) noexcept
{
    return DiscreteElement<QueryTags...>(arr);
}
 
template <concepts::discrete_dimension... QueryTags, concepts::discrete_dimension... Tags>
KOKKOS_FUNCTION constexpr DiscreteElement<QueryTags...> select(
        DiscreteElement<Tags...>&& arr) noexcept
{
    return DiscreteElement<QueryTags...>(std::move(arr));
}
 
/// Returns a reference towards the DiscreteElement that contains the QueryTag
template <
        concepts::discrete_dimension QueryTag,
        concepts::discrete_element HeadDElem,
        concepts::discrete_element... TailDElems>
KOKKOS_FUNCTION constexpr auto const& take(HeadDElem const& head, TailDElems 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<HeadDElem>>) {
        static_assert(
                (!type_seq_contains_v<TypeSeq<QueryTag>, to_type_seq_t<TailDElems>> && ...),
                "ERROR: tag redundant");
        return head;
    } else {
        static_assert(sizeof...(TailDElems) > 0, "ERROR: tag not found");
        return take<QueryTag>(tail...);
    }
    DDC_IF_NVCC_THEN_POP
}
 
namespace detail {
 
/// Returns a reference to the underlying `std::array`
template <ddc::concepts::discrete_dimension... Tags>
KOKKOS_FUNCTION constexpr std::array<DiscreteElementType, sizeof...(Tags)>& array(
        DiscreteElement<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<DiscreteElementType, sizeof...(Tags)> const& array(
        DiscreteElement<Tags...> const& v) noexcept
{
    return v.m_values;
}
 
} // namespace detail
 
template <concepts::discrete_dimension DDim>
constexpr DiscreteElement<DDim> create_reference_discrete_element() noexcept
{
    return DiscreteElement<DDim>(0);
}
 
/** A DiscreteElement identifies an element of the discrete dimension
 *
 * Each one is tagged by its associated dimensions.
 */
template <concepts::discrete_dimension... Tags>
class DiscreteElement
{
    using tags_seq = TypeSeq<Tags...>;
 
    static_assert(
            type_seq_is_unique_v<tags_seq>,
            "The dimensions of a DiscreteElement must be unique");
 
    friend KOKKOS_FUNCTION constexpr std::array<DiscreteElementType, sizeof...(Tags)>& detail::
            array<Tags...>(DiscreteElement<Tags...>& v) noexcept;
 
    friend KOKKOS_FUNCTION constexpr std::array<DiscreteElementType, sizeof...(Tags)> const&
    detail::array<Tags...>(DiscreteElement<Tags...> const& v) noexcept;
 
private:
    std::array<DiscreteElementType, sizeof...(Tags)> m_values;
 
public:
    using value_type = DiscreteElementType;
 
    static KOKKOS_FUNCTION constexpr std::size_t size() noexcept
    {
        return sizeof...(Tags);
    }
 
public:
    KOKKOS_DEFAULTED_FUNCTION constexpr DiscreteElement() = default;
 
    KOKKOS_DEFAULTED_FUNCTION constexpr DiscreteElement(DiscreteElement const&) = default;
 
    KOKKOS_DEFAULTED_FUNCTION constexpr DiscreteElement(DiscreteElement&&) = default;
 
    template <concepts::discrete_element... DElems>
    KOKKOS_FUNCTION constexpr explicit DiscreteElement(DElems const&... delems) noexcept
        : m_values {take<Tags>(delems...).template uid<Tags>()...}
    {
    }
 
    template <std::convertible_to<DiscreteElementType> IntegerType>
    KOKKOS_FUNCTION constexpr explicit DiscreteElement(
            std::array<IntegerType, sizeof...(Tags)> const& values) noexcept
        : m_values(
                  detail::convert_array_to<
                          DiscreteElementType>(values, std::make_index_sequence<sizeof...(Tags)>()))
    {
    }
 
    template <std::convertible_to<DiscreteElementType>... Params>
    KOKKOS_FUNCTION constexpr explicit DiscreteElement(Params const&... params) noexcept
        requires((!is_discrete_element_v<Params> && ...) && sizeof...(Params) == sizeof...(Tags))
        : m_values {static_cast<DiscreteElementType>(params)...}
    {
    }
 
    KOKKOS_DEFAULTED_FUNCTION ~DiscreteElement() = default;
 
    KOKKOS_DEFAULTED_FUNCTION DiscreteElement& operator=(DiscreteElement const& other) = default;
 
    KOKKOS_DEFAULTED_FUNCTION DiscreteElement& operator=(DiscreteElement&& other) = default;
 
    template <concepts::discrete_dimension QueryTag>
    KOKKOS_FUNCTION constexpr value_type& uid() noexcept
    {
        static_assert(in_tags_v<QueryTag, tags_seq>, "requested Tag absent from DiscreteElement");
        return m_values[type_seq_rank_v<QueryTag, tags_seq>];
    }
 
    template <concepts::discrete_dimension QueryTag>
    KOKKOS_FUNCTION constexpr value_type const& uid() const noexcept
    {
        static_assert(in_tags_v<QueryTag, tags_seq>, "requested Tag absent from DiscreteElement");
        return m_values[type_seq_rank_v<QueryTag, tags_seq>];
    }
 
    KOKKOS_FUNCTION constexpr value_type& uid() noexcept
        requires(sizeof...(Tags) == 1)
    {
        return m_values[0];
    }
 
    KOKKOS_FUNCTION constexpr value_type const& uid() const noexcept
        requires(sizeof...(Tags) == 1)
    {
        return m_values[0];
    }
 
    KOKKOS_FUNCTION constexpr DiscreteElement& operator++()
        requires(sizeof...(Tags) == 1)
    {
        ++m_values[0];
        return *this;
    }
 
    KOKKOS_FUNCTION constexpr DiscreteElement operator++(int)
        requires(sizeof...(Tags) == 1)
    {
        DiscreteElement const tmp = *this;
        ++m_values[0];
        return tmp;
    }
 
    KOKKOS_FUNCTION constexpr DiscreteElement& operator--()
        requires(sizeof...(Tags) == 1)
    {
        --m_values[0];
        return *this;
    }
 
    KOKKOS_FUNCTION constexpr DiscreteElement operator--(int)
        requires(sizeof...(Tags) == 1)
    {
        DiscreteElement const tmp = *this;
        --m_values[0];
        return tmp;
    }
 
    template <concepts::discrete_dimension... OTags>
    KOKKOS_FUNCTION constexpr DiscreteElement& 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 DiscreteElement& operator+=(IntegralType const& rhs)
        requires(sizeof...(Tags) == 1)
    {
        m_values[0] += rhs;
        return *this;
    }
 
    template <concepts::discrete_dimension... OTags>
    KOKKOS_FUNCTION constexpr DiscreteElement& 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 DiscreteElement& operator-=(IntegralType const& rhs)
        requires(sizeof...(Tags) == 1)
    {
        m_values[0] -= rhs;
        return *this;
    }
};
 
namespace detail {
 
void print_discrete_element(std::ostream& os, std::span<DiscreteElementType const> view);
 
} // namespace detail
 
template <concepts::discrete_dimension... Tags>
std::ostream& operator<<(std::ostream& os, DiscreteElement<Tags...> const& arr)
{
    detail::print_discrete_element(os, detail::array(arr));
    return os;
}
 
 
template <concepts::discrete_dimension... Tags, concepts::discrete_dimension... OTags>
KOKKOS_FUNCTION constexpr bool operator==(
        DiscreteElement<Tags...> const& lhs,
        DiscreteElement<OTags...> const& rhs) noexcept
{
    return ((lhs.template uid<Tags>() == rhs.template uid<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... Tags, concepts::discrete_dimension... OTags>
KOKKOS_FUNCTION constexpr bool operator!=(
        DiscreteElement<Tags...> const& lhs,
        DiscreteElement<OTags...> const& rhs) noexcept
{
    return !(lhs == rhs);
}
#endif
 
template <concepts::discrete_dimension Tag>
KOKKOS_FUNCTION constexpr bool operator<(
        DiscreteElement<Tag> const& lhs,
        DiscreteElement<Tag> const& rhs)
{
    return lhs.uid() < rhs.uid();
}
 
template <concepts::discrete_dimension Tag>
KOKKOS_FUNCTION constexpr bool operator<=(
        DiscreteElement<Tag> const& lhs,
        DiscreteElement<Tag> const& rhs)
{
    return lhs.uid() <= rhs.uid();
}
 
template <concepts::discrete_dimension Tag>
KOKKOS_FUNCTION constexpr bool operator>(
        DiscreteElement<Tag> const& lhs,
        DiscreteElement<Tag> const& rhs)
{
    return lhs.uid() > rhs.uid();
}
 
template <concepts::discrete_dimension Tag>
KOKKOS_FUNCTION constexpr bool operator>=(
        DiscreteElement<Tag> const& lhs,
        DiscreteElement<Tag> const& rhs)
{
    return lhs.uid() >= rhs.uid();
}
 
/// right external binary operators: +, -
 
template <concepts::discrete_dimension... Tags, concepts::discrete_dimension... OTags>
KOKKOS_FUNCTION constexpr DiscreteElement<Tags...> operator+(
        DiscreteElement<Tags...> const& lhs,
        DiscreteVector<OTags...> const& rhs)
{
    static_assert((type_seq_contains_v<TypeSeq<OTags>, TypeSeq<Tags...>> && ...));
    DiscreteElement<Tags...> result(lhs);
    result += rhs;
    return result;
}
 
template <concepts::discrete_dimension Tag, std::integral IntegralType>
KOKKOS_FUNCTION constexpr DiscreteElement<Tag> operator+(
        DiscreteElement<Tag> const& lhs,
        IntegralType const& rhs)
    requires(!is_discrete_vector_v<IntegralType>)
{
    DiscreteElement<Tag> result(lhs);
    result += rhs;
    return result;
}
 
template <concepts::discrete_dimension... Tags, concepts::discrete_dimension... OTags>
KOKKOS_FUNCTION constexpr DiscreteElement<Tags...> operator-(
        DiscreteElement<Tags...> const& lhs,
        DiscreteVector<OTags...> const& rhs)
{
    static_assert((type_seq_contains_v<TypeSeq<OTags>, TypeSeq<Tags...>> && ...));
    DiscreteElement<Tags...> result(lhs);
    result -= rhs;
    return result;
}
 
template <concepts::discrete_dimension Tag, std::integral IntegralType>
KOKKOS_FUNCTION constexpr DiscreteElement<Tag> operator-(
        DiscreteElement<Tag> const& lhs,
        IntegralType const& rhs)
    requires(!is_discrete_vector_v<IntegralType>)
{
    DiscreteElement<Tag> result(lhs);
    result -= rhs;
    return result;
}
 
/// binary operator: -
 
template <concepts::discrete_dimension... Tags, concepts::discrete_dimension... OTags>
KOKKOS_FUNCTION constexpr DiscreteVector<Tags...> operator-(
        DiscreteElement<Tags...> const& lhs,
        DiscreteElement<OTags...> const& rhs)
{
    static_assert(type_seq_same_v<TypeSeq<Tags...>, TypeSeq<OTags...>>);
    return DiscreteVector<Tags...>((uid<Tags>(lhs) - uid<Tags>(rhs))...);
}
 
} // namespace ddc