fix(core): preserve wide counts in indexed end tokens

phlex/core/detail/accumulator_node.hpp

@@ -4,6 +4,7 @@
 #include "phlex/core/fold/send.hpp"
 #include "phlex/core/message.hpp"
 #include "phlex/phlex_core_export.hpp"
+#include "phlex/utilities/signed_size.hpp"
 
 #include "oneapi/tbb/concurrent_hash_map.h"
 #include "oneapi/tbb/concurrent_queue.h"
@@ -12,6 +13,7 @@
 
 #include <atomic>
 #include <cassert>
+#include <cstddef>
 #include <functional>
 #include <memory>
 #include <string>
@@ -121,7 +123,9 @@ namespace phlex::detail::internal {
     struct cached_accumulator {
       std::shared_ptr<accumulator_msg_t> accumulator_msg;
       tbb::concurrent_queue<std::size_t> msg_ids;
-      std::atomic<int> counter;
+      // Signed balance of pending invocations. It may be negative when a flush arrives before
+      // all concurrent invocations are processed; zero means that the partition is complete.
+      std::atomic<signed_size_t> pending_invocations;
       std::atomic_flag flush_received;
       std::size_t original_message_id{};
     };
@@ -270,7 +274,7 @@ namespace phlex::detail::internal {
     entry->original_message_id = msg.msg_id;
     emit_pending_ids(entry);
     // Handle the flush-before-partition case: if the flush token already arrived (and
-    // left counter == 0 because no fold inputs flowed under this partition), emit the
+    // left pending_invocations == 0 because no fold inputs flowed under this partition), emit the
     // initial accumulator value now that we finally have the index and original message
     // ID needed to construct the output.
     cleanup_cache_entry(a);
@@ -283,7 +287,7 @@ namespace phlex::detail::internal {
     accessor a;
     cached_results_.insert(a, index->hash());
     auto* entry = &a->second;
-    entry->counter -= count;
+    entry->pending_invocations -= count;
     std::ignore = entry->flush_received.test_and_set();
     cleanup_cache_entry(a);
   }
@@ -312,10 +316,11 @@ namespace phlex::detail::internal {
     auto* entry = &a->second;
     // The `accumulator_msg` check guards the flush-before-partition case: a zero-count
     // flush can establish the cache entry with `flush_received == true` and
-    // `counter == 0` before the partition message has supplied the index and initial
+    // `pending_invocations == 0` before the partition message has supplied the index and initial
     // value.  In that case we defer emission until `handle_partition_message` re-enters
     // cleanup with a populated `accumulator_msg`.
-    if (entry->flush_received.test() and entry->counter == 0 and entry->accumulator_msg) {
+    if (entry->flush_received.test() and entry->pending_invocations == 0 and
+        entry->accumulator_msg) {
       output_port<0>(repeater_).try_put(entry->accumulator_msg->release_as_message(
         node_name_, output_, entry->original_message_id));
       ++emitted_result_count_;
@@ -330,7 +335,7 @@ namespace phlex::detail::internal {
     if (!cached_results_.find(a, key)) {
       return;
     }
-    ++a->second.counter;
+    ++a->second.pending_invocations;
     cleanup_cache_entry(a);
   }
 }

phlex/core/detail/repeater_node.cpp

@@ -72,10 +72,10 @@ namespace phlex::detail::internal {
     }
   }
 
-  int repeater_node::emit_pending_ids(cached_product* entry)
+  signed_size_t repeater_node::emit_pending_ids(cached_product* entry)
   {
     assert(entry->data_msg);
-    int num_emitted{};
+    signed_size_t num_emitted{};
     std::size_t msg_id{};
     while (entry->msg_ids.try_pop(msg_id)) {
       output_port<0>(repeater_).try_put({.store = entry->data_msg->store, .id = msg_id});
@@ -100,7 +100,7 @@ namespace phlex::detail::internal {
     cached_products_.insert(a, key);
     auto* entry = &a->second;
     entry->data_msg = std::make_shared<message>(msg);
-    entry->counter += emit_pending_ids(entry);
+    entry->pending_invocations += emit_pending_ids(entry);
     return key;
   }
 
@@ -111,7 +111,7 @@ namespace phlex::detail::internal {
     accessor a;
     cached_products_.insert(a, key);
     auto* entry = &a->second;
-    entry->counter -= count;
+    entry->pending_invocations -= count;
     std::ignore = entry->flush_received.test_and_set();
     return key;
   }
@@ -133,7 +133,7 @@ namespace phlex::detail::internal {
         auto* entry = &a->second;
         if (entry->data_msg) {
           output_port<0>(repeater_).try_put(*entry->data_msg);
-          ++entry->counter;
+          ++entry->pending_invocations;
         }
       }
       return key;
@@ -145,7 +145,7 @@ namespace phlex::detail::internal {
     auto* entry = &a->second;
     if (entry->data_msg) {
       output_port<0>(repeater_).try_put({.store = entry->data_msg->store, .id = msg_id});
-      entry->counter += 1 + emit_pending_ids(entry);
+      entry->pending_invocations += 1 + emit_pending_ids(entry);
     } else {
       entry->msg_ids.push(msg_id);
     }
@@ -161,11 +161,11 @@ namespace phlex::detail::internal {
 
     auto* entry = &a->second;
     if (!cache_enabled_) {
-      if (entry->counter == 0 and entry->data_msg) {
+      if (entry->pending_invocations == 0 and entry->data_msg) {
         output_port<0>(repeater_).try_put(*entry->data_msg);
       }
       cached_products_.erase(a);
-    } else if (entry->flush_received.test() and entry->counter == 0) {
+    } else if (entry->flush_received.test() and entry->pending_invocations == 0) {
       cached_products_.erase(a);
     }
   }

phlex/core/detail/repeater_node.hpp

@@ -4,12 +4,14 @@
 #include "phlex/phlex_core_export.hpp"
 
 #include "phlex/core/message.hpp"
+#include "phlex/utilities/signed_size.hpp"
 
 #include "oneapi/tbb/concurrent_hash_map.h"
 #include "oneapi/tbb/concurrent_queue.h"
 #include "oneapi/tbb/flow_graph.h"
 
 #include <atomic>
+#include <cstddef>
 #include <memory>
 #include <string>
 
@@ -46,14 +48,16 @@ namespace phlex::detail::internal {
     struct cached_product {
       std::shared_ptr<message> data_msg;
       tbb::concurrent_queue<std::size_t> msg_ids;
-      std::atomic<int> counter;
+      // Signed balance of pending invocations. It may be negative when a flush arrives before
+      // all concurrent invocations are processed; zero means that the partition is complete.
+      std::atomic<signed_size_t> pending_invocations;
       std::atomic_flag flush_received;
     };
 
     using cache_t = tbb::concurrent_hash_map<std::size_t, cached_product>; // Key is the index hash
     using accessor = cache_t::accessor;
 
-    int emit_pending_ids(cached_product* entry);
+    signed_size_t emit_pending_ids(cached_product* entry);
     std::size_t handle_data_message(message const& msg);
     std::size_t handle_flush_token(indexed_end_token const& token);
     std::size_t handle_index_message(index_message const& msg);

phlex/core/index_router.cpp

@@ -282,7 +282,7 @@ namespace phlex::detail {
       [end_token_entries = std::move(end_token_entries)](flush_gate const& fc) {
         for (auto const& entry : *end_token_entries) {
           auto const count = fc.committed_count_for_layer(entry.counting_layer_hash);
-          entry.flush_port->try_put({.index = fc.index(), .count = static_cast<int>(count)});
+          entry.flush_port->try_put({.index = fc.index(), .count = count});
         }
       });
 
@@ -513,7 +513,7 @@ namespace phlex::detail {
     // be at or below zero from earlier rollup notifications) and erroneously declare the tracker
     // ready.
     if (not is_lowest_layer_hash(child_layer_hash)) {
-      gate.expect_child_rollups(static_cast<std::ptrdiff_t>(count));
+      gate.expect_child_rollups(count);
     }
     gate.update_expected_count(child_layer_hash, count);
   }

phlex/core/message.hpp

@@ -9,6 +9,7 @@
 #include "phlex/model/handle.hpp"
 #include "phlex/model/identifier.hpp"
 #include "phlex/model/product_store.hpp"
+#include "phlex/utilities/signed_size.hpp"
 #include "phlex/utilities/sized_tuple.hpp"
 
 #include "oneapi/tbb/flow_graph.h" // <-- belongs somewhere else
@@ -31,7 +32,11 @@ namespace phlex::detail {
 
   struct indexed_end_token {
     data_cell_index_ptr index;
-    int count;
+    // The count is the number of direct children processed for this index. It uses signed_size_t
+    // because it is subtracted from a pending-invocations counter, which may be negative when the
+    // indexed_end_token arrives before all pending invocations are processed.
+    // (See the pending_invocations members in repeater_node and accumulator_node.)
+    signed_size_t count;
   };
 
   struct message {

phlex/model/flush_gate.cpp

@@ -22,10 +22,10 @@ namespace phlex::detail {
     return std::ranges::fold_left(committed_counts_ | std::views::values, 0uz, std::plus{});
   }
 
-  std::size_t flush_gate::committed_count_for_layer(
+  signed_size_t flush_gate::committed_count_for_layer(
     data_cell_index::hash_type const layer_hash) const
   {
-    return committed_counts_.count(layer_hash);
+    return checked_signed_size(committed_counts_.count(layer_hash));
   }
 
   void flush_gate::update_expected_count(data_cell_index::hash_type const layer_hash,
@@ -43,7 +43,10 @@ namespace phlex::detail {
     --pending_child_rollups_;
   }
 
-  void flush_gate::expect_child_rollups(std::ptrdiff_t const n) { pending_child_rollups_ += n; }
+  void flush_gate::expect_child_rollups(std::size_t const n)
+  {
+    pending_child_rollups_ += checked_signed_size(n);
+  }
 
   void flush_gate::send_flush()
   {

phlex/model/flush_gate.hpp

@@ -34,6 +34,7 @@
 #include "phlex/model/data_cell_counts.hpp"
 #include "phlex/model/data_cell_index.hpp"
 #include "phlex/phlex_model_export.hpp"
+#include "phlex/utilities/signed_size.hpp"
 
 #include <atomic>
 #include <cstddef>
@@ -57,7 +58,7 @@ namespace phlex::detail {
     data_cell_index_ptr index() const { return index_; }
     std::size_t expected_total_count() const;
     std::size_t committed_total_count() const;
-    std::size_t committed_count_for_layer(data_cell_index::hash_type layer_hash) const;
+    signed_size_t committed_count_for_layer(data_cell_index::hash_type layer_hash) const;
     data_cell_counts const& committed_counts() const { return committed_counts_; }
 
     // Merges an expected child count into the accumulated expected counts.  Each call
@@ -74,7 +75,7 @@ namespace phlex::detail {
     // for by the expected-count message that produced them (from the input_node or an
     // unfold).  The pending counter is signed because rollups can be recorded before the
     // corresponding expected-count message has been processed.
-    void expect_child_rollups(std::ptrdiff_t n);
+    void expect_child_rollups(std::size_t n);
 
     void set_flush_callback(flush_callback_t callback) { flush_callback_ = std::move(callback); }
     void send_flush();
@@ -94,7 +95,7 @@ namespace phlex::detail {
     std::size_t expected_flush_count_{0};
     // Signed running balance: (expected non-lowest direct-child rollups) - (rollups received).
     // Commit-ready when this reaches zero (and the expected-count message has arrived).
-    std::atomic<std::ptrdiff_t> pending_child_rollups_{0};
+    std::atomic<signed_size_t> pending_child_rollups_{0};
     flush_callback_t flush_callback_;
   };
 

phlex/utilities/signed_size.hpp

@@ -0,0 +1,22 @@
+#ifndef PHLEX_UTILITIES_SIGNED_SIZE_HPP
+#define PHLEX_UTILITIES_SIGNED_SIZE_HPP
+
+#include <cstddef>
+#include <limits>
+#include <stdexcept>
+#include <type_traits>
+#include <utility>
+
+namespace phlex::detail {
+  using signed_size_t = std::make_signed_t<std::size_t>;
+
+  inline signed_size_t checked_signed_size(std::size_t const value)
+  {
+    if (std::cmp_greater(value, std::numeric_limits<signed_size_t>::max())) {
+      throw std::overflow_error{"Value exceeds signed size range"};
+    }
+    return static_cast<signed_size_t>(value);
+  }
+}
+
+#endif // PHLEX_UTILITIES_SIGNED_SIZE_HPP

test/flush_gate_test.cpp

@@ -18,14 +18,18 @@
 #include "phlex/model/data_cell_index.hpp"
 #include "phlex/model/flush_gate.hpp"
 #include "phlex/model/identifier.hpp"
+#include "phlex/utilities/signed_size.hpp"
 
 #include "catch2/catch_test_macros.hpp"
 #include "oneapi/tbb/concurrent_hash_map.h"
 #include "oneapi/tbb/concurrent_vector.h"
 #include "oneapi/tbb/parallel_for.h"
 
+#include <limits>
 #include <memory>
 #include <ranges>
+#include <stdexcept>
+#include <utility>
 #include <vector>
 
 using namespace phlex;
@@ -102,6 +106,46 @@ TEST_CASE("flush_gate: single-layer hierarchy (job -> runs)", "[flush_gate]")
   CHECK(gates.empty());
 }
 
+TEST_CASE("flush_gate: large count behavior", "[flush_gate]")
+{
+  auto job = data_cell_index::job();
+  auto large_run = job->make_child("large_run", 0);
+
+  // committed_count_for_layer(...) returns signed_size_t, while update_expected_count(...)
+  // accepts std::size_t. Test a value above INT_MAX that remains representable as signed_size_t.
+  auto const count_above_int_max = static_cast<signed_size_t>(std::numeric_limits<int>::max()) + 1;
+  auto gate = make_gate(job, 0);
+
+  gate->update_expected_count(large_run->layer_hash(),
+                              static_cast<std::size_t>(count_above_int_max));
+  REQUIRE(gate->all_children_accounted());
+
+  SECTION("Preserve counts above INT_MAX")
+  {
+    CHECK(gate->committed_count_for_layer(large_run->layer_hash()) == count_above_int_max);
+  }
+  SECTION("Reject counts above signed_size_t range")
+  {
+    // We assume std::size_t's maximum is larger than signed_size_t's maximum, which is reasonable
+    // for the platforms we need to support.
+    static_assert(
+      std::cmp_greater(std::numeric_limits<std::size_t>::max(),
+                       std::numeric_limits<signed_size_t>::max()),
+      "std::size_t max must be larger than signed_size_t max for this test to be valid");
+
+    auto too_large_run = job->make_child("too_large_run", 0);
+    auto const count_above_signed_size_max =
+      static_cast<std::size_t>(std::numeric_limits<signed_size_t>::max()) + 1;
+    auto overflow_gate = make_gate(job, 0);
+
+    overflow_gate->update_expected_count(too_large_run->layer_hash(), count_above_signed_size_max);
+    REQUIRE(overflow_gate->all_children_accounted());
+
+    CHECK_THROWS_AS(overflow_gate->committed_count_for_layer(too_large_run->layer_hash()),
+                    std::overflow_error);
+  }
+}
+
 TEST_CASE("flush_gate: two-layer hierarchy (job -> runs -> spills)", "[flush_gate]")
 {
   constexpr std::size_t n_runs = 3;