dim.hpp

// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
//
// SPDX-License-Identifier: BSD-3-Clause
 
#ifndef GKO_PUBLIC_CORE_BASE_DIM_HPP_
#define GKO_PUBLIC_CORE_BASE_DIM_HPP_
 
 
#include <iostream>
 
#include <ginkgo/core/base/types.hpp>
 
 
namespace gko {
 
 
template <size_type Dimensionality, typename DimensionType = size_type>
struct dim {
    static constexpr size_type dimensionality = Dimensionality;
    friend struct dim<dimensionality + 1>;
 
    using dimension_type = DimensionType;
 
    constexpr GKO_ATTRIBUTES dim() : dim{dimension_type{}} {}
 
    explicit constexpr GKO_ATTRIBUTES dim(const dimension_type& size)
        : first_{size}, rest_{size}
    {}
 
    template <typename... Rest, std::enable_if_t<sizeof...(Rest) ==
                                                 Dimensionality - 1>* = nullptr>
    constexpr GKO_ATTRIBUTES dim(const dimension_type& first,
                                 const Rest&... rest)
        : first_{first}, rest_{static_cast<dimension_type>(rest)...}
    {}
 
    constexpr GKO_ATTRIBUTES const dimension_type& operator[](
        const size_type& dimension) const noexcept
    {
        return GKO_ASSERT(dimension < dimensionality),
               dimension == 0 ? first_ : rest_[dimension - 1];
    }
 
    GKO_ATTRIBUTES dimension_type& operator[](
        const size_type& dimension) noexcept
    {
        return GKO_ASSERT(dimension < dimensionality),
               dimension == 0 ? first_ : rest_[dimension - 1];
    }
 
    explicit constexpr GKO_ATTRIBUTES operator bool() const
    {
        return static_cast<bool>(first_) && static_cast<bool>(rest_);
    }
 
    friend constexpr GKO_ATTRIBUTES bool operator==(const dim& x, const dim& y)
    {
        return x.first_ == y.first_ && x.rest_ == y.rest_;
    }
 
    friend constexpr GKO_ATTRIBUTES bool operator!=(const dim& x, const dim& y)
    {
        return !(x == y);
    }
 
    friend constexpr GKO_ATTRIBUTES dim operator*(const dim& x, const dim& y)
    {
        return dim(x.first_ * y.first_, x.rest_ * y.rest_);
    }
 
    friend std::ostream& operator<<(std::ostream& os, const dim& x)
    {
        os << "(";
        x.print_to(os);
        os << ")";
        return os;
    }
 
private:
    void inline print_to(std::ostream& os) const
    {
        os << first_ << ", ";
        rest_.print_to(os);
    }
 
 
    constexpr GKO_ATTRIBUTES dim(const dimension_type first,
                                 dim<dimensionality - 1> rest)
        : first_{first}, rest_{rest}
    {}
 
    dimension_type first_;
    dim<dimensionality - 1, dimension_type> rest_;
};
 
 
// base case for dim recursive template
template <typename DimensionType>
struct dim<1u, DimensionType> {
    static constexpr size_type dimensionality = 1u;
    friend struct dim<2>;
 
    using dimension_type = DimensionType;
 
    constexpr GKO_ATTRIBUTES dim(const dimension_type& size = dimension_type{})
        : first_{size}
    {}
 
    constexpr GKO_ATTRIBUTES const dimension_type& operator[](
        const size_type& dimension) const noexcept
    {
        return GKO_ASSERT(dimension == 0), first_;
    }
 
    GKO_ATTRIBUTES dimension_type& operator[](const size_type& dimension)
    {
        return GKO_ASSERT(dimension == 0), first_;
    }
 
    explicit constexpr GKO_ATTRIBUTES operator bool() const
    {
        return static_cast<bool>(first_);
    }
 
    friend constexpr GKO_ATTRIBUTES bool operator==(const dim& x, const dim& y)
    {
        return x.first_ == y.first_;
    }
 
    friend constexpr GKO_ATTRIBUTES bool operator!=(const dim& x, const dim& y)
    {
        return !(x == y);
    }
 
    friend constexpr GKO_ATTRIBUTES dim operator*(const dim& x, const dim& y)
    {
        return dim(x.first_ * y.first_);
    }
 
    friend std::ostream& operator<<(std::ostream& os, const dim& x)
    {
        os << "(";
        x.print_to(os);
        os << ")";
        return os;
    }
 
private:
    void inline print_to(std::ostream& os) const { os << first_; }
 
    dimension_type first_;
};
 
 
template <typename DimensionType>
constexpr GKO_ATTRIBUTES GKO_INLINE dim<2, DimensionType> transpose(
    const dim<2, DimensionType>& dimensions) noexcept
{
    return {dimensions[1], dimensions[0]};
}
 
 
}  // namespace gko
 
 
#endif  // GKO_PUBLIC_CORE_BASE_DIM_HPP_