Use proper index_pair_like constraint for submdspan

kokkos · Sep 3, 2024 · 7da7385 · 7da7385
1 parent 17ba6af
commit 7da7385
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 4 deletions.
diff --git a/include/experimental/__p2630_bits/submdspan_extents.hpp b/include/experimental/__p2630_bits/submdspan_extents.hpp
@@ -17,6 +17,7 @@
 #pragma once
 
 #include <tuple>
+#include <complex>
 
 #include "strided_slice.hpp"
 namespace MDSPAN_IMPL_STANDARD_NAMESPACE {
@@ -52,6 +53,33 @@ template <class OffsetType, class ExtentType, class StrideType>
 struct is_strided_slice<
     strided_slice<OffsetType, ExtentType, StrideType>> : std::true_type {};
 
+// Helper for identifying valid pair like things
+template <class T, class IndexType> struct index_pair_like : std::false_type {};
+
+template <class IdxT1, class IdxT2, class IndexType>
+struct index_pair_like<std::pair<IdxT1, IdxT2>, IndexType> {
+  static constexpr bool value = std::is_convertible_v<IdxT1, IndexType> &&
+                                std::is_convertible_v<IdxT2, IndexType>;
+};
+
+template <class IdxT1, class IdxT2, class IndexType>
+struct index_pair_like<std::tuple<IdxT1, IdxT2>, IndexType> {
+  static constexpr bool value = std::is_convertible_v<IdxT1, IndexType> &&
+                                std::is_convertible_v<IdxT2, IndexType>;
+};
+
+template <class IdxT, class IndexType>
+struct index_pair_like<std::complex<IdxT>, IndexType> {
+  static constexpr bool value = std::is_convertible_v<IdxT, IndexType>;
+};
+
+template <class IdxT, class IndexType>
+struct index_pair_like<std::array<IdxT, 2>, IndexType> {
+  static constexpr bool value = std::is_convertible_v<IdxT, IndexType>;
+};
+
+// FIXME: we actually need to pass IndexType into all of these
+
 // first_of(slice): getting begin of slice specifier range
 MDSPAN_TEMPLATE_REQUIRES(
   class Integral,
@@ -70,7 +98,7 @@ first_of(const ::MDSPAN_IMPL_STANDARD_NAMESPACE::full_extent_t &) {
 
 MDSPAN_TEMPLATE_REQUIRES(
   class Slice,
-  /* requires */(std::is_convertible_v<Slice, std::tuple<size_t, size_t>>)
+  /* requires */(index_pair_like<Slice, size_t>::value)
 )
 MDSPAN_INLINE_FUNCTION
 constexpr auto first_of(const Slice &i) {
@@ -100,7 +128,7 @@ constexpr Integral
 
 MDSPAN_TEMPLATE_REQUIRES(
   size_t k, class Extents, class Slice,
-  /* requires */(std::is_convertible_v<Slice, std::tuple<size_t, size_t>>)
+  /* requires */(index_pair_like<Slice, size_t>::value)
 )
 MDSPAN_INLINE_FUNCTION
 constexpr auto last_of(std::integral_constant<size_t, k>, const Extents &,

diff --git a/include/experimental/__p2630_bits/submdspan_mapping.hpp b/include/experimental/__p2630_bits/submdspan_mapping.hpp
@@ -98,8 +98,7 @@ template<class SliceSpecifier, class IndexType>
 struct is_range_slice {
   constexpr static bool value =
     std::is_same_v<SliceSpecifier, full_extent_t> ||
-    std::is_convertible_v<SliceSpecifier,
-                          std::tuple<IndexType, IndexType>>;
+    index_pair_like<SliceSpecifier, IndexType>::value;
 };
 
 template<class SliceSpecifier, class IndexType>