Tuple.h

// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2021 The Eigen Team
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
// SPDX-License-Identifier: MPL-2.0
 
#ifndef EIGEN_TUPLE_GPU
#define EIGEN_TUPLE_GPU
 
#include <type_traits>
#include <utility>
 
// This is a replacement of std::tuple that can be used in device code.
 
namespace Eigen {
namespace internal {
namespace tuple_impl {
 
// Internal tuple implementation.
template <size_t N, typename... Types>
class TupleImpl;
 
// Generic recursive tuple.
template <size_t N, typename T1, typename... Ts>
class TupleImpl<N, T1, Ts...> {
 public:
  // Tuple may contain Eigen types.
  EIGEN_MAKE_ALIGNED_OPERATOR_NEW
 
  // Default constructor, enable if all types are default-constructible.
  template <typename U1 = T1,
            typename EnableIf = std::enable_if_t<std::is_default_constructible<U1>::value &&
                                                 reduce_all<std::is_default_constructible<Ts>::value...>::value>>
  constexpr EIGEN_DEVICE_FUNC TupleImpl() : m_head{}, m_tail{} {}
 
  // Element constructor.
  template <typename U1, typename... Us,
            // Only enable if...
            typename EnableIf = std::enable_if_t<
                // the number of input arguments match, and ...
                sizeof...(Us) == sizeof...(Ts) && (
                                                      // this does not look like a copy/move constructor.
                                                      N > 1 || std::is_convertible<U1, T1>::value)>>
  constexpr EIGEN_DEVICE_FUNC TupleImpl(U1&& arg1, Us&&... args)
      : m_head(std::forward<U1>(arg1)), m_tail(std::forward<Us>(args)...) {}
 
  // The first stored value.
  EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE T1& head() { return m_head; }
 
  EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE const T1& head() const { return m_head; }
 
  // The tail values.
  EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE TupleImpl<N - 1, Ts...>& tail() { return m_tail; }
 
  EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE const TupleImpl<N - 1, Ts...>& tail() const { return m_tail; }
 
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void swap(TupleImpl& other) {
    using numext::swap;
    swap(m_head, other.m_head);
    swap(m_tail, other.m_tail);
  }
 
  template <typename... UTypes>
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TupleImpl& operator=(const TupleImpl<N, UTypes...>& other) {
    m_head = other.m_head;
    m_tail = other.m_tail;
    return *this;
  }
 
  template <typename... UTypes>
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TupleImpl& operator=(TupleImpl<N, UTypes...>&& other) {
    m_head = std::move(other.m_head);
    m_tail = std::move(other.m_tail);
    return *this;
  }
 
 private:
  // Allow related tuples to reference m_head/m_tail.
  template <size_t M, typename... UTypes>
  friend class TupleImpl;
 
  T1 m_head;
  TupleImpl<N - 1, Ts...> m_tail;
};
 
// Empty tuple specialization.
template <>
class TupleImpl<size_t(0)> {};
 
template <typename TupleType>
struct is_tuple : std::false_type {};
 
template <typename... Types>
struct is_tuple<TupleImpl<sizeof...(Types), Types...>> : std::true_type {};
 
// Gets an element from a tuple.
template <size_t Idx, typename T1, typename... Ts>
struct tuple_get_impl {
  using TupleType = TupleImpl<sizeof...(Ts) + 1, T1, Ts...>;
  using ReturnType = typename tuple_get_impl<Idx - 1, Ts...>::ReturnType;
 
  static constexpr EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE ReturnType& run(TupleType& tuple) {
    return tuple_get_impl<Idx - 1, Ts...>::run(tuple.tail());
  }
 
  static constexpr EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE const ReturnType& run(const TupleType& tuple) {
    return tuple_get_impl<Idx - 1, Ts...>::run(tuple.tail());
  }
};
 
// Base case, getting the head element.
template <typename T1, typename... Ts>
struct tuple_get_impl<0, T1, Ts...> {
  using TupleType = TupleImpl<sizeof...(Ts) + 1, T1, Ts...>;
  using ReturnType = T1;
 
  static constexpr EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE T1& run(TupleType& tuple) { return tuple.head(); }
 
  static constexpr EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE const T1& run(const TupleType& tuple) { return tuple.head(); }
};
 
// Concatenates N Tuples.
template <size_t NTuples, typename... Tuples>
struct tuple_cat_impl;
 
template <size_t NTuples, size_t N1, typename... Args1, size_t N2, typename... Args2, typename... Tuples>
struct tuple_cat_impl<NTuples, TupleImpl<N1, Args1...>, TupleImpl<N2, Args2...>, Tuples...> {
  using TupleType1 = TupleImpl<N1, Args1...>;
  using TupleType2 = TupleImpl<N2, Args2...>;
  using MergedTupleType = TupleImpl<N1 + N2, Args1..., Args2...>;
 
  using ReturnType = typename tuple_cat_impl<NTuples - 1, MergedTupleType, Tuples...>::ReturnType;
 
  // Uses the index sequences to extract and merge elements from tuple1 and tuple2,
  // then recursively calls again.
  template <typename Tuple1, size_t... I1s, typename Tuple2, size_t... I2s, typename... MoreTuples>
  static constexpr EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE ReturnType run(Tuple1&& tuple1, std::index_sequence<I1s...>,
                                                                        Tuple2&& tuple2, std::index_sequence<I2s...>,
                                                                        MoreTuples&&... tuples) {
    return tuple_cat_impl<NTuples - 1, MergedTupleType, Tuples...>::run(
        MergedTupleType(tuple_get_impl<I1s, Args1...>::run(std::forward<Tuple1>(tuple1))...,
                        tuple_get_impl<I2s, Args2...>::run(std::forward<Tuple2>(tuple2))...),
        std::forward<MoreTuples>(tuples)...);
  }
 
  // Concatenates the first two tuples.
  template <typename Tuple1, typename Tuple2, typename... MoreTuples>
  static constexpr EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE ReturnType run(Tuple1&& tuple1, Tuple2&& tuple2,
                                                                        MoreTuples&&... tuples) {
    return run(std::forward<Tuple1>(tuple1), std::make_index_sequence<N1>{}, std::forward<Tuple2>(tuple2),
               std::make_index_sequence<N2>{}, std::forward<MoreTuples>(tuples)...);
  }
};
 
// Base case with a single tuple.
template <size_t N, typename... Args>
struct tuple_cat_impl<1, TupleImpl<N, Args...>> {
  using ReturnType = TupleImpl<N, Args...>;
 
  template <typename Tuple1>
  static constexpr EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE ReturnType run(Tuple1&& tuple1) {
    return tuple1;
  }
};
 
// Special case of no tuples.
template <>
struct tuple_cat_impl<0> {
  using ReturnType = TupleImpl<0>;
  static constexpr EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE ReturnType run() { return ReturnType{}; }
};
 
// For use in make_tuple, unwraps a reference_wrapper.
template <typename T>
struct unwrap_reference_wrapper {
  using type = T;
};
 
template <typename T>
struct unwrap_reference_wrapper<std::reference_wrapper<T>> {
  using type = T&;
};
 
// For use in make_tuple, decays a type and unwraps a reference_wrapper.
template <typename T>
struct unwrap_decay {
  using type = typename unwrap_reference_wrapper<std::decay_t<T>>::type;
};

template <typename Tuple>
struct tuple_size;
 
template <typename... Types>
struct tuple_size<TupleImpl<sizeof...(Types), Types...>> : std::integral_constant<size_t, sizeof...(Types)> {};

template <size_t Idx, typename... Types>
constexpr EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const typename tuple_get_impl<Idx, Types...>::ReturnType& get(
    const TupleImpl<sizeof...(Types), Types...>& tuple) {
  return tuple_get_impl<Idx, Types...>::run(tuple);
}
 
template <size_t Idx, typename... Types>
constexpr EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE typename tuple_get_impl<Idx, Types...>::ReturnType& get(
    TupleImpl<sizeof...(Types), Types...>& tuple) {
  return tuple_get_impl<Idx, Types...>::run(tuple);
}

template <typename... Tuples,
          typename EnableIf = std::enable_if_t<internal::reduce_all<is_tuple<std::decay_t<Tuples>>::value...>::value>>
constexpr EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
    typename tuple_cat_impl<sizeof...(Tuples), std::decay_t<Tuples>...>::ReturnType
    tuple_cat(Tuples&&... tuples) {
  return tuple_cat_impl<sizeof...(Tuples), std::decay_t<Tuples>...>::run(std::forward<Tuples>(tuples)...);
}

template <typename... Args, typename ReturnType = TupleImpl<sizeof...(Args), Args&...>>
constexpr EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE ReturnType tie(Args&... args) noexcept {
  return ReturnType{args...};
}

template <typename... Args, typename ReturnType = TupleImpl<sizeof...(Args), typename unwrap_decay<Args>::type...>>
constexpr EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE ReturnType make_tuple(Args&&... args) {
  return ReturnType{std::forward<Args>(args)...};
}

template <typename... Args>
constexpr EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TupleImpl<sizeof...(Args), Args...> forward_as_tuple(Args&&... args) {
  return TupleImpl<sizeof...(Args), Args...>(std::forward<Args>(args)...);
}

template <typename... Types>
using tuple = TupleImpl<sizeof...(Types), Types...>;
 
}  // namespace tuple_impl
}  // namespace internal
}  // namespace Eigen
 
#endif  // EIGEN_TUPLE_GPU