From e1250ce16f5097793837230c9b63e925a7c43762 Mon Sep 17 00:00:00 2001
From: Andrey Prokopenko <prokopenkoav@ornl.gov>
Date: Mon, 11 Nov 2024 10:35:05 -0500
Subject: [PATCH] Simplify Iota struct usage

---
 .../bvh_driver/benchmark_registration.hpp     | 21 ++++++-------------
 examples/brute_force/example_brute_force.cpp  | 18 ++++------------
 test/tstQueryTreeIntersectsKDOP.cpp           | 20 +++++-------------
 3 files changed, 15 insertions(+), 44 deletions(-)

diff --git a/benchmarks/bvh_driver/benchmark_registration.hpp b/benchmarks/bvh_driver/benchmark_registration.hpp
index 179f831fd..f22991d83 100644
--- a/benchmarks/bvh_driver/benchmark_registration.hpp
+++ b/benchmarks/bvh_driver/benchmark_registration.hpp
@@ -32,19 +32,12 @@ struct is_boost_rtree<BoostExt::RTree<Geometry>> : std::true_type
 template <typename Geometry>
 inline constexpr bool is_boost_rtree_v = is_boost_rtree<Geometry>::value;
 
-template <typename MemorySpace, typename Index = int>
+template <typename MemorySpace>
 struct Iota
 {
+  static_assert(Kokkos::is_memory_space_v<MemorySpace>);
   using memory_space = MemorySpace;
-  using index_type = Index;
-
-  size_t _n;
-
-  template <typename T,
-            typename Enable = std::enable_if_t<std::is_integral_v<T>>>
-  Iota(T n)
-      : _n(n)
-  {}
+  int _n;
 };
 
 template <typename MemorySpace>
@@ -54,10 +47,7 @@ struct ArborX::AccessTraits<Iota<MemorySpace>, ArborX::PrimitivesTag>
 
   using memory_space = typename Self::memory_space;
   static KOKKOS_FUNCTION size_t size(Self const &self) { return self._n; }
-  static KOKKOS_FUNCTION auto get(Self const &, size_t i)
-  {
-    return (typename Self::index_type)i;
-  }
+  static KOKKOS_FUNCTION auto get(Self const &, int i) { return i; }
 };
 
 struct Spec
@@ -164,7 +154,8 @@ auto makeTree(ExecutionSpace const &space, Primitives const &primitives)
     return TreeType(space, primitives);
   else
     return TreeType(space,
-                    Iota<typename TreeType::memory_space>{primitives.size()},
+                    Iota<typename TreeType::memory_space>{
+                        static_cast<int>(primitives.size())},
                     primitives);
 }
 
diff --git a/examples/brute_force/example_brute_force.cpp b/examples/brute_force/example_brute_force.cpp
index 46cfb65a3..22445f79a 100644
--- a/examples/brute_force/example_brute_force.cpp
+++ b/examples/brute_force/example_brute_force.cpp
@@ -23,19 +23,12 @@ struct Dummy
 using ExecutionSpace = Kokkos::DefaultExecutionSpace;
 using MemorySpace = ExecutionSpace::memory_space;
 
-template <typename MemorySpace, typename Index = int>
+template <typename MemorySpace>
 struct Iota
 {
+  static_assert(Kokkos::is_memory_space_v<MemorySpace>);
   using memory_space = MemorySpace;
-  using index_type = Index;
-
-  size_t _n;
-
-  template <typename T,
-            typename Enable = std::enable_if_t<std::is_integral_v<T>>>
-  Iota(T n)
-      : _n(n)
-  {}
+  int _n;
 };
 
 template <typename MemorySpace>
@@ -45,10 +38,7 @@ struct ArborX::AccessTraits<Iota<MemorySpace>, ArborX::PrimitivesTag>
 
   using memory_space = typename Self::memory_space;
   static KOKKOS_FUNCTION size_t size(Self const &self) { return self._n; }
-  static KOKKOS_FUNCTION auto get(Self const &, size_t i)
-  {
-    return (typename Self::index_type)i;
-  }
+  static KOKKOS_FUNCTION auto get(Self const &, int i) { return i; }
 };
 
 struct DummyIndexableGetter
diff --git a/test/tstQueryTreeIntersectsKDOP.cpp b/test/tstQueryTreeIntersectsKDOP.cpp
index a62c885e9..21dac3737 100644
--- a/test/tstQueryTreeIntersectsKDOP.cpp
+++ b/test/tstQueryTreeIntersectsKDOP.cpp
@@ -18,19 +18,12 @@
 
 #include <boost/test/unit_test.hpp>
 
-template <typename MemorySpace, typename Index = int>
+template <typename MemorySpace>
 struct Iota
 {
+  static_assert(Kokkos::is_memory_space_v<MemorySpace>);
   using memory_space = MemorySpace;
-  using index_type = Index;
-
-  size_t _n;
-
-  template <typename T,
-            typename Enable = std::enable_if_t<std::is_integral_v<T>>>
-  Iota(T n)
-      : _n(n)
-  {}
+  int _n;
 };
 
 template <typename MemorySpace>
@@ -40,10 +33,7 @@ struct ArborX::AccessTraits<Iota<MemorySpace>, ArborX::PrimitivesTag>
 
   using memory_space = typename Self::memory_space;
   static KOKKOS_FUNCTION size_t size(Self const &self) { return self._n; }
-  static KOKKOS_FUNCTION auto get(Self const &, size_t i)
-  {
-    return (typename Self::index_type)i;
-  }
+  static KOKKOS_FUNCTION auto get(Self const &, int i) { return i; }
 };
 
 #include <vector>
@@ -73,7 +63,7 @@ BOOST_AUTO_TEST_CASE_TEMPLATE(intersects_kdop, DeviceType, ARBORX_DEVICE_TYPES)
       {{0, 0, 3}}, // 12
   };
   ArborX::BoundingVolumeHierarchy const tree(
-      ExecutionSpace{}, Iota<MemorySpace>{primitives.size()},
+      ExecutionSpace{}, Iota<MemorySpace>{static_cast<int>(primitives.size())},
       Kokkos::create_mirror_view_and_copy(
           MemorySpace{},
           Kokkos::View<Point *, Kokkos::HostSpace, Kokkos::MemoryUnmanaged>(