From f4d35696b4181db60e6754521a4e5ae399bb35cf Mon Sep 17 00:00:00 2001
From: Trevor Steil <steil1@llnl.gov>
Date: Wed, 31 Jul 2024 20:15:40 -0500
Subject: [PATCH] Fixes bug in block_partitioner

---
 include/ygm/container/array.hpp                    | 14 +++++++-------
 include/ygm/container/detail/block_partitioner.hpp |  4 ++--
 2 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/include/ygm/container/array.hpp b/include/ygm/container/array.hpp
index 14e11bae..49b19538 100644
--- a/include/ygm/container/array.hpp
+++ b/include/ygm/container/array.hpp
@@ -355,8 +355,10 @@ class array
     std::vector<std::pair<const key_type, const mapped_type>> tmp_values;
     tmp_values.reserve(local_size());
     local_for_all(
-        [&tmp_values](const key_type& index, const mapped_type& value) {
-          tmp_values.push_back(std::make_pair(index, value));
+        [&tmp_values, size](const key_type& index, const mapped_type& value) {
+          if (index < size) {
+            tmp_values.push_back(std::make_pair(index, value));
+          }
         });
 
     m_global_size = size;
@@ -368,14 +370,14 @@ class array
 
     // Repopulate array values
     for (const auto& [index, value] : tmp_values) {
-      if (index < size) {
-        async_set(index, value);
-      }
+      async_set(index, value);
     }
 
     m_comm.barrier();
   }
 
+  void resize(const size_type size) { resize(size, m_default_value); }
+
   size_t local_size() { return partitioner.local_size(); }
 
   size_t size() const {
@@ -383,8 +385,6 @@ class array
     return m_global_size;
   }
 
-  void resize(const size_type size) { resize(size, m_default_value); }
-
   void local_clear() { resize(0); }
 
   void local_swap(self_type& other) {
diff --git a/include/ygm/container/detail/block_partitioner.hpp b/include/ygm/container/detail/block_partitioner.hpp
index 32ac86cf..830088e3 100644
--- a/include/ygm/container/detail/block_partitioner.hpp
+++ b/include/ygm/container/detail/block_partitioner.hpp
@@ -21,7 +21,7 @@ struct block_partitioner {
         m_partitioned_size(partitioned_size) {
     m_small_block_size = partitioned_size / m_comm_size;
     m_large_block_size =
-        m_small_block_size + ((partitioned_size / m_comm_size) > 0);
+        m_small_block_size + ((partitioned_size % m_comm_size) > 0);
 
     if (m_comm_rank < (partitioned_size % m_comm_size)) {
       m_local_start_index = m_comm_rank * m_large_block_size;
@@ -62,7 +62,7 @@ struct block_partitioner {
 
   index_type local_index(const index_type &global_index) {
     index_type to_return = global_index - m_local_start_index;
-    ASSERT_RELEASE((to_return >= 0) && (to_return <= m_small_block_size));
+    ASSERT_RELEASE((to_return >= 0) && (to_return < m_local_size));
     return to_return;
   }