tuple.hpp 2.51 KB
Newer Older
Chao Liu's avatar
Chao Liu committed
1
2
3
4
#ifndef CK_TUPLE_HPP
#define CK_TUPLE_HPP

#include "integral_constant.hpp"
Chao Liu's avatar
Chao Liu committed
5
#include "Sequence.hpp"
Chao Liu's avatar
Chao Liu committed
6
7
8

namespace ck {

Chao Liu's avatar
Chao Liu committed
9
namespace detail {
Chao Liu's avatar
Chao Liu committed
10

Chao Liu's avatar
Chao Liu committed
11
12
13
14
template <index_t>
struct TupleElementKey
{
};
Chao Liu's avatar
Chao Liu committed
15

Chao Liu's avatar
Chao Liu committed
16
17
18
19
20
template <typename Key, typename Data>
struct TupleElement
{
    template <typename T>
    __host__ __device__ explicit constexpr TupleElement(T&& v) : mData(static_cast<T&&>(v))
Chao Liu's avatar
Chao Liu committed
21
22
23
    {
    }

Chao Liu's avatar
Chao Liu committed
24
    Data mData;
Chao Liu's avatar
Chao Liu committed
25
26
};

Chao Liu's avatar
Chao Liu committed
27
28
template <typename Key, typename Data>
__host__ __device__ constexpr const Data& get_tuple_element(const TupleElement<Key, Data>& x)
Chao Liu's avatar
Chao Liu committed
29
{
Chao Liu's avatar
Chao Liu committed
30
31
    return x.mData;
}
Chao Liu's avatar
Chao Liu committed
32

Chao Liu's avatar
Chao Liu committed
33
34
template <typename Key, typename Data>
__host__ __device__ constexpr Data& get_tuple_element(TupleElement<Key, Data>& x)
Chao Liu's avatar
Chao Liu committed
35
{
Chao Liu's avatar
Chao Liu committed
36
37
    return x.mData;
}
Chao Liu's avatar
Chao Liu committed
38

Chao Liu's avatar
Chao Liu committed
39
40
template <typename Key, typename Data>
__host__ __device__ constexpr Data&& get_tuple_element(TupleElement<Key, Data>&& x)
Chao Liu's avatar
Chao Liu committed
41
{
Chao Liu's avatar
Chao Liu committed
42
43
    return static_cast<Data&&>(x.mData);
}
Chao Liu's avatar
Chao Liu committed
44

Chao Liu's avatar
Chao Liu committed
45
46
47
48
49
template <typename Indices, typename... Xs>
struct TupleImpl;

template <index_t... Is, typename... Xs>
struct TupleImpl<Sequence<Is...>, Xs...> : TupleElement<TupleElementKey<Is>, Xs>...
Chao Liu's avatar
Chao Liu committed
50
{
Chao Liu's avatar
Chao Liu committed
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
    template <typename... Ys>
    __host__ __device__ explicit constexpr TupleImpl(Ys&&... ys)
        : TupleElement<TupleElementKey<Is>, Xs>(static_cast<Ys&&>(ys))...
    {
    }

    __host__ __device__ static constexpr index_t Size() { return sizeof...(Xs); }

    template <index_t I>
    __host__ __device__ constexpr const auto& GetElementByKey(TupleElementKey<I>) const
    {
        return get_tuple_element<TupleElementKey<I>>(*this);
    }

    template <index_t I>
    __host__ __device__ constexpr auto& GetElementByKey(TupleElementKey<I>)
    {
        return get_tuple_element<TupleElementKey<I>>(*this);
    }
Chao Liu's avatar
Chao Liu committed
70
71
};

Chao Liu's avatar
Chao Liu committed
72
73
74
75
} // namespace detail

template <typename... Xs>
struct Tuple : detail::TupleImpl<typename arithmetic_sequence_gen<0, sizeof...(Xs), 1>::type, Xs...>
Chao Liu's avatar
Chao Liu committed
76
{
Chao Liu's avatar
Chao Liu committed
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
    using base =
        detail::TupleImpl<typename arithmetic_sequence_gen<0, sizeof...(Xs), 1>::type, Xs...>;

    template <typename... Ys>
    __host__ __device__ explicit constexpr Tuple(Ys&&... ys) : base(static_cast<Ys&&>(ys)...)
    {
    }

    template <index_t I>
    __host__ __device__ constexpr const auto& At(Number<I>) const
    {
        static_assert(I < base::Size(), "wrong! out of range");
        return GetElementByKey(detail::TupleElementKey<I>{});
    }

    template <index_t I>
    __host__ __device__ constexpr auto& At(Number<I>)
    {
        static_assert(I < base::Size(), "wrong! out of range");
        return GetElementByKey(detail::TupleElementKey<I>{});
    }
Chao Liu's avatar
Chao Liu committed
98
99
100
101
};

} // namespace ck
#endif