From ca02e1080b714fd3f6adff91ab07a194ca38ff00 Mon Sep 17 00:00:00 2001
From: Arkadiusz Szczepkowicz <arek.szczepkowicz@ng-analytics.com>
Date: Tue, 16 May 2023 15:49:16 +0200
Subject: [PATCH] #2074: Update finding next phase for OfflineLB

---
 .../balance/lb_data_restart_reader.h          | 30 +++++++++++--------
 .../collection/balance/offlinelb/offlinelb.cc |  5 ++--
 2 files changed, 21 insertions(+), 14 deletions(-)

diff --git a/src/vt/vrt/collection/balance/lb_data_restart_reader.h b/src/vt/vrt/collection/balance/lb_data_restart_reader.h
index d0a35a34cd..fd6640b6d3 100644
--- a/src/vt/vrt/collection/balance/lb_data_restart_reader.h
+++ b/src/vt/vrt/collection/balance/lb_data_restart_reader.h
@@ -123,6 +123,24 @@ struct LBDataRestartReader : runtime::component::Component<LBDataRestartReader>
       | num_phases_;
   }
 
+  /**
+   * \brief Find the next phase
+   *
+   * \param phase the current phase
+   *
+   * \return the next phase
+   */
+  PhaseType findNextPhase(PhaseType phase) const {
+    auto next = phase + 1;
+    for(; next < num_phases_; ++next) {
+      if(history_.find(next) != history_.end()) {
+        return next;
+      }
+    }
+    vtAssert(history_.find(next) != history_.end(), "Must have a valid phase");
+    return next;
+  }
+
   /**
    * \brief Get the elements assigned for a given phase
    *
@@ -131,12 +149,6 @@ struct LBDataRestartReader : runtime::component::Component<LBDataRestartReader>
    * \return element assigned to this node
    */
   std::set<ElementIDStruct> const& getDistro(PhaseType phase) {
-    for(; phase < num_phases_; ++phase) {
-      if(history_.find(phase) != history_.end()) {
-        break;
-      }
-    }
-
     auto iter = history_.find(phase);
     vtAssert(iter != history_.end(), "Must have a valid phase");
     return iter->second;
@@ -148,12 +160,6 @@ struct LBDataRestartReader : runtime::component::Component<LBDataRestartReader>
    * \param[in] phase the phase to clear
    */
   void clearDistro(PhaseType phase) {
-    for(; phase < num_phases_; ++phase) {
-      if(history_.find(phase) != history_.end()) {
-        break;
-      }
-    }
-
     auto iter = history_.find(phase);
     if (iter != history_.end()) {
       history_.erase(iter);
diff --git a/src/vt/vrt/collection/balance/offlinelb/offlinelb.cc b/src/vt/vrt/collection/balance/offlinelb/offlinelb.cc
index cb0ba986f3..fad602f208 100644
--- a/src/vt/vrt/collection/balance/offlinelb/offlinelb.cc
+++ b/src/vt/vrt/collection/balance/offlinelb/offlinelb.cc
@@ -55,11 +55,12 @@ void OfflineLB::init(objgroup::proxy::Proxy<OfflineLB> in_proxy) {
 }
 
 void OfflineLB::runLB(TimeType) {
-  auto const& distro = theLBDataReader()->getDistro(phase_ + 1);
+  auto nextPhase = theLBDataReader()->findNextPhase(phase_);
+  auto const& distro = theLBDataReader()->getDistro(nextPhase);
   for (auto&& elm : distro) {
     migrateObjectTo(elm, theContext()->getNode());
   }
-  theLBDataReader()->clearDistro(phase_ + 1);
+  theLBDataReader()->clearDistro(nextPhase);
 }
 
 }}}} /* end namespace vt::vrt::collection::lb */