From 377f4716147bbebb5d42769fb4b92dc79c3dc120 Mon Sep 17 00:00:00 2001
From: Phil Miller <pbmille@sandia.gov>
Date: Wed, 16 Sep 2020 12:45:19 -0400
Subject: [PATCH] #1055: Preserve ability to look up object -> proxy mapping
 outside LB (i.e. during model testing)

---
 src/vt/vrt/collection/balance/node_stats.cc      | 16 ++++++++++------
 .../test_model_per_collection.extended.cc        |  4 ++--
 2 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/src/vt/vrt/collection/balance/node_stats.cc b/src/vt/vrt/collection/balance/node_stats.cc
index 2d007d452f..0a443a4877 100644
--- a/src/vt/vrt/collection/balance/node_stats.cc
+++ b/src/vt/vrt/collection/balance/node_stats.cc
@@ -133,14 +133,19 @@ void NodeStats::startIterCleanup(PhaseType phase, int look_back) {
   // Convert the temp ID node_data_ for the last iteration into perm ID for
   // stats output
   auto const prev_data = std::move(node_data_[phase]);
-  std::unordered_map<ElementIDType,TimeType> new_data;
+  std::unordered_map<ElementIDType,TimeType> new_node_data;
+  std::unordered_map<ElementIDType,VirtualProxyType> new_collection_lookup;
   for (auto& elm : prev_data) {
-    auto iter = node_temp_to_perm_.find(elm.first);
+    auto temp_id = elm.first;
+    auto iter = node_temp_to_perm_.find(temp_id);
     vtAssert(iter != node_temp_to_perm_.end(), "Temp ID must exist");
     auto perm_id = iter->second;
-    new_data[perm_id] = elm.second;
+
+    new_node_data[perm_id] = elm.second;
+    new_collection_lookup[perm_id] = node_collection_lookup_[temp_id];
   }
-  node_data_[phase] = std::move(new_data);
+  node_data_[phase] = std::move(new_node_data);
+  node_collection_lookup_ = std::move(new_collection_lookup);
 
   if (phase - look_back >= 0) {
     node_data_.erase(phase - look_back);
@@ -152,7 +157,6 @@ void NodeStats::startIterCleanup(PhaseType phase, int look_back) {
   NodeStats::node_migrate_.clear();
   NodeStats::node_temp_to_perm_.clear();
   NodeStats::node_perm_to_temp_.clear();
-  node_collection_lookup_.clear();
 }
 
 ElementIDType NodeStats::getNextElm() {
@@ -289,7 +293,7 @@ ElementIDType NodeStats::addNodeStats(
 
   vt_debug_print(
     lb, node,
-    "NodeStats::addNodeStats: temp_id={}, perm_id={}, phase={}, subphases={}, load={}\n",
+    "NodeStats::addNodeStats: temp_id={:x}, perm_id={:x}, phase={}, subphases={}, load={}\n",
     temp_id, perm_id, phase, subphase_time.size(), time
   );
 
diff --git a/tests/unit/collection/test_model_per_collection.extended.cc b/tests/unit/collection/test_model_per_collection.extended.cc
index 9d5180248d..03aa88cf00 100644
--- a/tests/unit/collection/test_model_per_collection.extended.cc
+++ b/tests/unit/collection/test_model_per_collection.extended.cc
@@ -90,9 +90,9 @@ std::unordered_map<ElementIDType, VirtualProxyType> id_proxy_map;
 
 template <typename ColT>
 void colHandler(MyMsg<ColT>*, ColT* col) {
-  // do nothing, except setting up our map using the temp ID, which will hit
+  // do nothing, except setting up our map using the perm ID, which will hit
   // every node
-  id_proxy_map[col->getTempID()] = col->getProxy();
+  id_proxy_map[col->getElmID()] = col->getProxy();
 }
 
 TEST_F(TestModelPerCollection, test_model_per_collection_1) {