Prefer tree.query() instead of query(tree, ...)

aprokop · Jan 5, 2024 · b7bf542 · b7bf542
1 parent 3906767
commit b7bf542
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 11 deletions.
diff --git a/src/details/ArborX_DetailsDistributedTreeNearest.hpp b/src/details/ArborX_DetailsDistributedTreeNearest.hpp
@@ -183,8 +183,8 @@ void DistributedTreeImpl::deviseStrategy(ExecutionSpace const &space,
   auto const &bottom_tree_sizes = tree._bottom_tree_sizes;
 
   // Find the k nearest local trees.
-  query(top_tree, space, predicates, LegacyDefaultCallback{}, nearest_ranks,
-        offset);
+  top_tree.query(space, predicates, LegacyDefaultCallback{}, nearest_ranks,
+                 offset);
 
   // Accumulate total leave count in the local trees until it reaches k which
   // is the number of neighbors queried for.  Stop if local trees get
@@ -267,10 +267,11 @@ void DistributedTreeImpl::reassessStrategy(
       WithinDistanceFromPredicates<Predicates, decltype(farthest_distances)>{
           queries, farthest_distances});
 
-  query(top_tree, space,
-        WithinDistanceFromPredicates<Predicates, decltype(farthest_distances)>{
-            queries, farthest_distances},
-        LegacyDefaultCallback{}, nearest_ranks, offset);
+  top_tree.query(
+      space,
+      WithinDistanceFromPredicates<Predicates, decltype(farthest_distances)>{
+          queries, farthest_distances},
+      LegacyDefaultCallback{}, nearest_ranks, offset);
   // NOTE: in principle, we could perform radius searches on the bottom_tree
   // rather than nearest queries.
 }
@@ -345,8 +346,8 @@ DistributedTreeImpl::queryDispatchImpl(NearestPredicateTag, Tree const &tree,
       // Perform queries that have been received
       Kokkos::View<PairIndexDistance *, MemorySpace> out(
           "ArborX::DistributedTree::query::pairs_index_distance", 0);
-      query(bottom_tree, space, fwd_queries, callback_with_distance, out,
-            offset);
+      bottom_tree.query(space, fwd_queries, callback_with_distance, out,
+                        offset);
 
       // Unzip
       auto const n = out.extent(0);

diff --git a/src/details/ArborX_DetailsDistributedTreeSpatial.hpp b/src/details/ArborX_DetailsDistributedTreeSpatial.hpp
@@ -55,8 +55,8 @@ DistributedTreeImpl::queryDispatch(SpatialPredicateTag, Tree const &tree,
 
   Kokkos::View<int *, MemorySpace> intersected_ranks(
       "ArborX::DistributedTree::query::spatial::intersected_ranks", 0);
-  query(top_tree, space, predicates, LegacyDefaultCallback{}, intersected_ranks,
-        offset);
+  top_tree.query(space, predicates, LegacyDefaultCallback{}, intersected_ranks,
+                 offset);
 
   Kokkos::View<int *, MemorySpace> ranks(
       "ArborX::DistributedTree::query::spatial::ranks", 0);
@@ -78,7 +78,7 @@ DistributedTreeImpl::queryDispatch(SpatialPredicateTag, Tree const &tree,
                    fwd_predicates, ids, ranks);
 
     // Perform predicates that have been received
-    query(bottom_tree, space, fwd_predicates, callback, out, offset);
+    bottom_tree.query(space, fwd_predicates, callback, out, offset);
 
     // Communicate results back
     communicateResultsBack(comm, space, out, offset, ranks, ids);