shape.hh 1.21 KB
Newer Older
Philip Trettner's avatar
Philip Trettner committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
#pragma once

#include <array>
#include <cassert>

namespace tg
{
template <int Order>
struct shape
{
    static inline constexpr int order = Order;

    std::array<int, order> dims;

public:
    shape() = default;
    constexpr shape(std::array<int, order> const& dims) : dims(dims) {}

public:
    constexpr int operator[](int o) const
    {
        assert(0 <= o && o < order);
        return dims[o];
    }

    template <int Order2>
    constexpr bool operator==(shape<Order2> const& rhs) const
    {
        if constexpr (Order == Order2)
        {
            for (auto i = 0; i < Order; ++i)
                if (dims[i] != rhs[i])
                    return false;
            return true;
        }
        else
            return false;
    }
    template <int Order2>
    constexpr bool operator!=(shape<Order2> const& rhs) const
    {
        return !operator==(rhs);
    }
};

inline constexpr shape<1> make_shape(int w) { return {{w}}; }
inline constexpr shape<2> make_shape(int w, int h) { return {{w, h}}; }
inline constexpr shape<3> make_shape(int w, int h, int d) { return {{w, h, d}}; }
template <int Order>
constexpr shape<Order> make_shape(std::array<int, Order> const& dims)
{
    return {dims};
}

} // namespace tg