#1497: vrt_coll: implement function to get local indices

DARMA-tasking · Jul 6, 2021 · 749470b · 749470b
1 parent b59bc88
commit 749470b
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 0 deletions.
diff --git a/src/vt/vrt/collection/manager.h b/src/vt/vrt/collection/manager.h
@@ -1800,6 +1800,16 @@ struct CollectionManager
   template <typename ColT, typename IndexT = typename ColT::IndexType>
   IndexT getRange(VirtualProxyType proxy);
 
+  /**
+   * \brief Get the local indices that are currently on this node
+   *
+   * \param[in] proxy the proxy of the collection
+   *
+   * \return the set of indices on this node
+   */
+  template <typename ColT, typename IndexT = typename ColT::IndexType>
+  std::set<IndexT> getLocalIndices(CollectionProxyWrapType<ColT> proxy);
+
   /**
    * \brief Make the filename for checkpoint/restore
    *

diff --git a/src/vt/vrt/collection/manager.impl.h b/src/vt/vrt/collection/manager.impl.h
@@ -3057,6 +3057,18 @@ IndexT CollectionManager::getRange(VirtualProxyType proxy) {
   return col_holder->max_idx;
 }
 
+template <typename ColT, typename IndexT>
+std::set<IndexT> CollectionManager::getLocalIndices(
+  CollectionProxyWrapType<ColT> proxy
+) {
+  auto elm_holder = findElmHolder<ColT>(proxy);
+  std::set<IndexT> local;
+  elm_holder->foreach([&](IndexT const& idx, CollectionBase<ColT,IndexT>*) {
+    local.insert(idx);
+  });
+  return local;
+}
+
 template <typename IndexT>
 std::string CollectionManager::makeMetaFilename(
   std::string file_base, bool make_sub_dirs