Hide ArborX::PairIndexRank

benchmarks/distributed_tree_driver/distributed_tree_driver.cpp

@@ -11,7 +11,6 @@
 
 #include <ArborXBenchmark_TimeMonitor.hpp>
 #include <ArborX_DistributedTree.hpp>
-#include <ArborX_PairIndexRank.hpp>
 #include <ArborX_Point.hpp>
 #include <ArborX_Version.hpp>
 

examples/distributed_tree/distributed_knn.cpp

@@ -8,7 +8,6 @@
  *                                                                          *
  * SPDX-License-Identifier: BSD-3-Clause                                    *
  ****************************************************************************/
-
 #include <ArborX.hpp>
 
 #include <Kokkos_Core.hpp>
@@ -22,19 +21,30 @@
 using ExecutionSpace = Kokkos::DefaultExecutionSpace;
 using MemorySpace = ExecutionSpace::memory_space;
 
-// WARNING: It is not allowed to add any definition into namespace ArborX. We
-// do it out of convenience here to print the results using
-// `std::ostream_iterator`. Please do not copy this code to your own code as
-// ArborX reserves the right to overload `operator<<` for any type defined
-// within its namespace. This means your code could break in the future.
-namespace ArborX
+struct PairIndexRank
 {
-std::ostream &operator<<(std::ostream &os, ArborX::PairIndexRank const &x)
+  unsigned index;
+  int rank;
+};
+
+std::ostream &operator<<(std::ostream &os, PairIndexRank const &x)
 {
   os << "(" << x.index << ", " << x.rank << ")";
   return os;
 }
-} // namespace ArborX
+
+struct CallbackWithRank
+{
+  int _rank;
+
+  template <typename Predicate, typename Value, typename OutputFunctor>
+  KOKKOS_FUNCTION void operator()(Predicate const &,
+                                  ArborX::PairValueIndex<Value> const &value,
+                                  OutputFunctor const &out) const
+  {
+    out({value.index, _rank});
+  }
+};
 
 int main(int argc, char *argv[])
 {
@@ -70,12 +80,14 @@ int main(int argc, char *argv[])
             points.data(), points.size()));
 
     ExecutionSpace exec;
-    ArborX::DistributedTree<MemorySpace> tree(comm, exec, points_device);
+    ArborX::DistributedTree<MemorySpace, ArborX::PairValueIndex<Point>> tree(
+        comm, exec, ArborX::Experimental::attach_indices(points_device));
 
-    Kokkos::View<ArborX::PairIndexRank *, MemorySpace> values("Example::values",
-                                                              0);
+    Kokkos::View<PairIndexRank *, MemorySpace> values("Example::values", 0);
     Kokkos::View<int *, MemorySpace> offsets("Example::offsets", 0);
     tree.query(exec, ArborX::Experimental::make_nearest(points_device, 3),
+               ArborX::Experimental::declare_callback_constrained(
+                   CallbackWithRank{comm_rank}),
                values, offsets);
 
     if (comm_rank == 0)
@@ -94,7 +106,7 @@ int main(int argc, char *argv[])
                 std::ostream_iterator<int>(std::cout, " "));
       std::cout << "\nvalues : ";
       std::copy(values_host.data(), values_host.data() + values.size(),
-                std::ostream_iterator<ArborX::PairIndexRank>(std::cout, " "));
+                std::ostream_iterator<PairIndexRank>(std::cout, " "));
       std::cout << "\n";
     }
   }

src/ArborX_DistributedTree.hpp

@@ -18,7 +18,6 @@
 #include <ArborX_DetailsDistributedTreeSpatial.hpp>
 #include <ArborX_DetailsKokkosExtStdAlgorithms.hpp>
 #include <ArborX_LinearBVH.hpp>
-#include <ArborX_PairIndexRank.hpp>
 #include <ArborX_PairValueIndex.hpp>
 
 #include <Kokkos_Core.hpp>

test/ArborX_BoostRTreeHelpers.hpp

@@ -24,7 +24,7 @@
 #include <ArborX_Predicates.hpp>
 #include <ArborX_Sphere.hpp>
 #ifdef ARBORX_ENABLE_MPI
-#include <ArborX_PairIndexRank.hpp>
+#include "ArborX_PairIndexRank.hpp"
 #endif
 
 #include <boost/range/adaptors.hpp>

src/details/ArborX_PairIndexRank.hpp → test/ArborX_PairIndexRank.hpp

@@ -14,6 +14,10 @@
 
 #include <Kokkos_Macros.hpp>
 
+#include <boost/test/tools/detail/print_helper.hpp>
+
+#include <iostream>
+
 namespace ArborX
 {
 
@@ -38,4 +42,18 @@ struct PairIndexRank
 
 } // namespace ArborX
 
+namespace boost::test_tools::tt_detail
+{
+
+template <>
+struct print_log_value<ArborX::PairIndexRank>
+{
+  void operator()(std::ostream &os, ArborX::PairIndexRank const &p)
+  {
+    os << '(' << p.index << ',' << p.rank << ')';
+  }
+};
+
+} // namespace boost::test_tools::tt_detail
+
 #endif

test/Search_UnitTestHelpers.hpp

@@ -13,7 +13,6 @@
 #define ARBORX_SEARCH_TEST_HELPERS_HPP
 
 // clang-format off
-#include "boost_ext/ArborXPairIndexRankComparison.hpp"
 #include "boost_ext/KokkosPairComparison.hpp"
 #include "boost_ext/TupleComparison.hpp"
 #include "boost_ext/CompressedStorageComparison.hpp"
@@ -22,6 +21,7 @@
 #include "ArborX_EnableViewComparison.hpp"
 #ifdef ARBORX_ENABLE_MPI
 #include "ArborX_BoostRTreeHelpers.hpp"
+#include "ArborX_PairIndexRank.hpp"
 #include <ArborX_DistributedTree.hpp>
 #endif
 #include <ArborX_Box.hpp>