From f144f34552fb03739beb64c061d8689c82fda7d8 Mon Sep 17 00:00:00 2001
From: Jonathan Lifflander <jliffla@sandia.gov>
Date: Thu, 1 Jul 2021 16:03:27 -0700
Subject: [PATCH] #1490: vrt_coll: finish generalizing code

---
 src/vt/vrt/collection/manager.h      | 16 ++++++++++------
 src/vt/vrt/collection/manager.impl.h | 15 ++++++++++-----
 2 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/src/vt/vrt/collection/manager.h b/src/vt/vrt/collection/manager.h
index 0eb7385662..1679ac468c 100644
--- a/src/vt/vrt/collection/manager.h
+++ b/src/vt/vrt/collection/manager.h
@@ -1871,9 +1871,12 @@ struct CollectionManager
     typename ColT::IndexType range, std::string const& file_base
   );
 
-private:
-  template <typename ColT, typename IdxT = typename ColT::IndexType>
-  struct RestoreMigrateMsg : vt::Message {
+  template <
+    typename ColT,
+    typename MsgT = vt::Message,
+    typename IdxT = typename ColT::IndexType
+  >
+  struct RestoreMigrateMsg : MsgT {
     RestoreMigrateMsg() = default;
     RestoreMigrateMsg(NodeType in_to_node, IdxT in_idx, CollectionProxyWrapType<ColT> in_proxy)
       : to_node_(in_to_node),
@@ -1887,17 +1890,18 @@ struct CollectionManager
   };
 
   template <typename ColT, typename IdxT = typename ColT::IndexType>
-  struct RestoreMigrateColMsg : RestoreMigrateMsg<ColT, IdxT>, CollectionMessage<ColT> {
+  struct RestoreMigrateColMsg
+    : RestoreMigrateMsg<ColT, CollectionMessage<ColT>, IdxT>
+  {
     RestoreMigrateColMsg() = default;
     RestoreMigrateColMsg(NodeType in_to_node, IdxT in_idx, CollectionProxyWrapType<ColT> in_proxy)
-      : RestoreMigrateMsg<ColT>(in_to_node, in_idx, in_proxy)
+      : RestoreMigrateMsg<ColT, CollectionMessage<ColT>, IdxT>(in_to_node, in_idx, in_proxy)
     { }
   };
 
   template <typename ColT>
   static void restoreHandler(RestoreMigrateMsg<ColT>* msg);
 
-public:
   /**
    * \brief Restore the collection (collective) from file on top of an existing
    * collection. Migrates collection elements to the rank saved from the
diff --git a/src/vt/vrt/collection/manager.impl.h b/src/vt/vrt/collection/manager.impl.h
index 943e3e32c0..db12124387 100644
--- a/src/vt/vrt/collection/manager.impl.h
+++ b/src/vt/vrt/collection/manager.impl.h
@@ -3162,7 +3162,9 @@ void CollectionManager::checkpointToFile(
 
 namespace detail {
 template <typename ColT>
-inline void restoreOffHomeElement(CollectionManager::RestoreMigrateColMsg<ColT>* msg, ColT*) {
+inline void restoreOffHomeElement(
+  CollectionManager::RestoreMigrateColMsg<ColT>* msg, ColT*
+) {
   auto idx = msg->idx_;
   auto node = msg->to_node_;
   auto proxy = msg->proxy_;
@@ -3219,10 +3221,13 @@ void CollectionManager::restoreFromFileInPlace(
         vtAssertExpr(mapped_node != uninitialized_destination);
         auto this_node = theContext()->getNode();
 
-        auto msg = makeMessage<RestoreMigrateMsg<ColT>>(this_node, idx, proxy);
-        theMsg()->sendMsg<
-          RestoreMigrateMsg<ColT>, restoreHandler<ColT>
-        >(mapped_node, msg);
+        using MsgType = RestoreMigrateMsg<ColT>;
+        auto msg = makeMessage<MsgType>(this_node, idx, proxy);
+        if (mapped_node != this_node) {
+          theMsg()->sendMsg<MsgType, restoreHandler<ColT>>(mapped_node, msg);
+        } else {
+          restoreHandler<ColT>(msg.get());
+        }
       }
     }
   });