From 8be2e96c6234d2d6b0c20cb46b16e321771eaa2b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cezary=20Skrzy=C5=84ski?=
 <cezary.skrzynski@ng-analytics.com>
Date: Fri, 6 May 2022 16:36:35 +0200
Subject: [PATCH] #1672: use `theNodeLBData` where necessary

---
 src/vt/vrt/collection/balance/baselb/baselb.cc        | 2 +-
 src/vt/vrt/collection/balance/lb_invoke/lb_manager.cc | 4 ++--
 tests/unit/collection/test_lb.extended.cc             | 5 +++--
 3 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/src/vt/vrt/collection/balance/baselb/baselb.cc b/src/vt/vrt/collection/balance/baselb/baselb.cc
index 40f004554b..36eae88ec7 100644
--- a/src/vt/vrt/collection/balance/baselb/baselb.cc
+++ b/src/vt/vrt/collection/balance/baselb/baselb.cc
@@ -263,7 +263,7 @@ void BaseLB::finalize(CountMsg* msg) {
 
 void BaseLB::recvSharedEdges(CommMsg* msg) {
   auto phase = thePhase()->getCurrentPhase();
-  auto comm_map = theNodeStats()->getNodeComm(phase);
+  auto comm_map = theNodeLBData()->getNodeComm(phase);
 
   if (comm_map != nullptr) {
     auto& comm = msg->comm_;
diff --git a/src/vt/vrt/collection/balance/lb_invoke/lb_manager.cc b/src/vt/vrt/collection/balance/lb_invoke/lb_manager.cc
index 41a17f9c13..e9a0eacbc9 100644
--- a/src/vt/vrt/collection/balance/lb_invoke/lb_manager.cc
+++ b/src/vt/vrt/collection/balance/lb_invoke/lb_manager.cc
@@ -732,8 +732,8 @@ void makeGraphSymmetric(
   PhaseType phase, objgroup::proxy::Proxy<lb::BaseLB> proxy
 ) {
   auto const this_node = theContext()->getNode();
-  auto iter = theNodeStats()->getNodeComm()->find(phase);
-  if (iter == theNodeStats()->getNodeComm()->end()) {
+  auto iter = theNodeLBData()->getNodeComm()->find(phase);
+  if (iter == theNodeLBData()->getNodeComm()->end()) {
     return;
   }
 
diff --git a/tests/unit/collection/test_lb.extended.cc b/tests/unit/collection/test_lb.extended.cc
index 7682ca7721..74ea5e7740 100644
--- a/tests/unit/collection/test_lb.extended.cc
+++ b/tests/unit/collection/test_lb.extended.cc
@@ -49,6 +49,7 @@
 
 #include "vt/vrt/collection/manager.h"
 #include "vt/vrt/collection/balance/lb_data_holder.h"
+#include "vt/vrt/collection/balance/node_lb_data.h"
 #include "vt/vrt/collection/balance/lb_invoke/lb_manager.h"
 #include "vt/utils/json/json_reader.h"
 #include "vt/utils/json/json_appender.h"
@@ -163,8 +164,8 @@ TEST_F(TestLoadBalancerOther, test_make_graph_symmetric) {
   // );
 
   auto phase = num_phases - 1;
-  auto iter = theNodeStats()->getNodeComm()->find(phase);
-  ASSERT_NE(iter, theNodeStats()->getNodeComm()->end());
+  auto iter = theNodeLBData()->getNodeComm()->find(phase);
+  ASSERT_NE(iter, theNodeLBData()->getNodeComm()->end());
 
   elm::CommMapType const& comm_data = iter->second;