fix(core): unify multi-layer join routing and completion accounting

phlex/core/CMakeLists.txt

@@ -61,6 +61,7 @@ install(
     glue.hpp
     graph_proxy.hpp
     input_arguments.hpp
+    join_registration.hpp
     message.hpp
     multilayer_join_node.hpp
     index_router.hpp

phlex/core/declared_unfold.cpp

@@ -9,17 +9,11 @@
 #include "phlex/model/products.hpp"
 #include "phlex/utilities/hashing.hpp"
 
-#include <fmt/format.h>
-// IWYU wants to remove the include below, but it is needed for the fmt::format call.
-#include <fmt/ranges.h> // IWYU pragma: keep
 #include <gsl/pointers>
 #include <oneapi/tbb/flow_graph.h>
 
 #include <cstddef>
 #include <memory>
-#include <ranges>
-#include <set>
-#include <stdexcept>
 #include <string>
 #include <utility>
 #include <vector>
@@ -60,14 +54,6 @@ namespace phlex::detail {
                       require_layers::always},
     child_layer_{std::move(child_layer)}
   {
-    // Forbid multi-layer inputs until Phlex v0.5.
-    std::set collapsed_layers{std::from_range, layers()};
-    if (collapsed_layers.size() > 1u) {
-      throw std::invalid_argument(
-        fmt::format("Unfolds cannot yet accept multi-layer inputs (node: {}, layers: {})",
-                    this->name().to_string(),
-                    collapsed_layers));
-    }
   }
 
   declared_unfold::~declared_unfold() = default;

phlex/core/detail/repeater_node.cpp

@@ -191,7 +191,10 @@ namespace phlex::detail::internal {
       assert(entry->data_msg);
       output_port<0>(repeater_).try_put(*entry->data_msg);
       cached_products_.erase(a);
-    } else if (entry->flush_received.test() and entry->pending_invocations == 0) {
+    } else if (entry->flush_received.test() and entry->pending_invocations == 0 and
+               entry->data_msg) {
+      // A zero-count flush can precede the product. Keep its completion state until the
+      // product arrives, otherwise that late product would be cached without a future flush.
       cached_products_.erase(a);
     }
   }

phlex/core/fold_join_node.hpp

@@ -3,11 +3,12 @@
 
 #include "phlex/core/detail/accumulator_node.hpp"
 #include "phlex/core/detail/repeater_node.hpp"
+#include "phlex/core/join_registration.hpp"
 #include "phlex/core/message.hpp"
 
+#include <gsl/assert>
 #include <oneapi/tbb/flow_graph.h>
 
-#include <cassert>
 #include <utility>
 #include <vector>
 
@@ -36,28 +37,30 @@ namespace phlex::detail {
   //             ⋮                                        ├──► join_node ──► (accum, data[0..N-1])
   //   data[N-1] ──► repeater[N-1] ────► port<N>(join) ───┘
 
-  template <std::size_t NInputs>
+  template <std::size_t NProducts>
   using partition_index_with_messages =
-    decltype(std::tuple_cat(std::tuple<index_message>{}, message_tuple<NInputs>{}));
+    decltype(std::tuple_cat(std::tuple<index_message>{}, message_tuple<NProducts>{}));
 
-  template <typename FoldResult, std::size_t NInputs>
+  template <typename FoldResult, std::size_t NProducts>
   using accumulator_with_messages = decltype(std::tuple_cat(
-    std::tuple<internal::accumulator_message<FoldResult>>{}, message_tuple<NInputs>{}));
+    std::tuple<internal::accumulator_message<FoldResult>>{}, message_tuple<NProducts>{}));
 
-  template <typename FoldResult, std::size_t NInputs>
+  template <typename FoldResult, std::size_t NProducts>
   using fold_join_node_base_t =
-    tbb::flow::composite_node<partition_index_with_messages<NInputs>,
-                              std::tuple<accumulator_with_messages<FoldResult, NInputs>>>;
-
-  template <typename FoldResult, std::size_t NInputs>
-  class fold_join_node : public fold_join_node_base_t<FoldResult, NInputs> {
-    using args_t = partition_index_with_messages<NInputs>;
-    using base_t = fold_join_node_base_t<FoldResult, NInputs>;
-    using join_args_t = accumulator_with_messages<FoldResult, NInputs>;
+    tbb::flow::composite_node<partition_index_with_messages<NProducts>,
+                              std::tuple<accumulator_with_messages<FoldResult, NProducts>>>;
+
+  template <typename FoldResult, std::size_t NProducts>
+  class fold_join_node : public fold_join_node_base_t<FoldResult, NProducts> {
+    using args_t = partition_index_with_messages<NProducts>;
+    using base_t = fold_join_node_base_t<FoldResult, NProducts>;
+    using join_args_t = accumulator_with_messages<FoldResult, NProducts>;
     using result_initializer_t = std::function<std::unique_ptr<FoldResult>(data_cell_index const&)>;
     using input_t = base_t::input_ports_type;
     using output_t = base_t::output_ports_type;
 
+    static_assert(NProducts >= 1, "fold_join_node requires at least one data product.");
+
     template <std::size_t... Is>
     static auto make_join(tbb::flow::graph& g, std::index_sequence<Is...>)
     {
@@ -86,17 +89,17 @@ namespace phlex::detail {
                        partition_layer_name,
                        std::move(output),
                        std::move(result_initializer)},
-      join_{make_join(g, std::make_index_sequence<NInputs>{})},
+      join_{make_join(g, std::make_index_sequence<NProducts>{})},
       name_{node_name},
       partition_layer_{partition_layer_name},
       layers_{std::move(layer_names)}
     {
-      assert(NInputs == layers_.size());
+      Expects(NProducts == layers_.size());
 
       // We do not worry about collapsing layers for the fold-join node as the accumulator
       // repeater will always belong to a higher layer than (at least one of) the data products
       // that serve as input to the fold operation.
-      repeaters_.reserve(NInputs);
+      repeaters_.reserve(NProducts);
       for (auto const& layer : layers_) {
         repeaters_.push_back(std::make_unique<internal::repeater_node>(g, name_, layer));
       }
@@ -110,7 +113,7 @@ namespace phlex::detail {
         (make_edge(*repeaters_[Is], input_port<Is + 1>(join_)), ...);
       };
 
-      set_ports(std::make_index_sequence<NInputs>{});
+      set_ports(std::make_index_sequence<NProducts>{});
     }
 
     tbb::flow::receiver<index_message>& partition_port()
@@ -126,36 +129,18 @@ namespace phlex::detail {
       return input_port<3>(result_repeater_);
     }
 
-    // Returns one named_index_port per slot the index router must wire up.
-    //
-    // The first entry is the *partition slot* — it routes the partition index to the
-    // accumulator's index_port and delivers the partition-level flush token to the
-    // accumulator's flush_port.  Its routing layer is the partition layer (e.g. "job"),
-    // but its counting layer is the fold's most-derived input data layer (e.g. "event"),
-    // because the accumulator is incremented once per fold call (= once per cell at the
-    // most-derived input layer) and the flush count must balance against that.
-    //
-    // The remaining entries are the data-input slots, one per repeater.  For these the
-    // routing layer and counting layer are the same: the repeater both routes and counts
-    // at its own input layer.
-    //
-    // FIXME: For multi-input folds, the "most-derived" input layer is the deepest of
-    // `layers_` in the data hierarchy.  Until the router exposes a depth comparison, we
-    // pick `layers_[0]`; all current tests have single-input folds so this is correct in
-    // practice.
-    std::vector<named_index_port> index_ports()
+    // The partition constrains the product-driven invocation domain. Its initialization
+    // connection remains separate from the accumulator's per-invocation index receiver.
+    // Ancestor data repeaters and the accumulator complete against the same invocations.
+    join_registration index_ports()
     {
-      std::vector<named_index_port> result;
-      result.reserve(1 + repeaters_.size()); // +1 for the result repeater
-      phlex::experimental::identifier const counting_layer_for_partition =
-        layers_.empty() ? partition_layer_ : layers_[0];
-      result.emplace_back(partition_layer_,
-                          counting_layer_for_partition,
-                          &result_repeater_.flush_port(),
-                          &result_repeater_.index_port());
+      join_registration result{.input_layers = layers_, .participants = {}};
+      result.participants.reserve(1 + repeaters_.size()); // +1 for the result repeater
+      result.participants.emplace_back(
+        partition_layer_, &result_repeater_.index_port(), &result_repeater_.flush_port());
       for (std::size_t i = 0; i != repeaters_.size(); ++i) {
-        result.emplace_back(
-          layers_[i], layers_[i], &repeaters_[i]->flush_port(), &repeaters_[i]->index_port());
+        result.participants.emplace_back(
+          layers_[i], &repeaters_[i]->index_port(), &repeaters_[i]->flush_port());
       }
       return result;
     }

phlex/core/fold_node.hpp

@@ -8,6 +8,7 @@
 #include "phlex/core/fold_join_node.hpp"
 #include "phlex/core/fwd.hpp"
 #include "phlex/core/input_arguments.hpp"
+#include "phlex/core/join_registration.hpp"
 #include "phlex/core/message.hpp"
 #include "phlex/core/node_builder.hpp"
 #include "phlex/core/product_selector.hpp"
@@ -129,7 +130,7 @@ namespace phlex::detail {
     tbb::flow::sender<message>& output_port() override { return join_.output_port(); }
     phlex::experimental::product_specifications const& output() const override { return output_; }
 
-    named_index_ports index_ports() final { return join_.index_ports(); }
+    join_registration index_ports() final { return join_.index_ports(); }
     std::size_t num_calls() const final { return calls_.load(); }
 
     std::size_t apply_fold(

phlex/core/framework_graph.cpp

@@ -24,10 +24,10 @@
 #include <format>
 #include <map>
 #include <ranges>
-#include <set>
 #include <stdexcept>
 #include <string>
 #include <utility>
+#include <vector>
 
 namespace phlex::detail {
   namespace {
@@ -65,30 +65,23 @@ namespace phlex::detail {
       return result;
     }
 
-    // Collects (input layer, child layer) pairs and per-input-layer unfold counts from all
-    // registered unfold nodes.  One pair is emitted per (unfold node, input layer): a
-    // multi-input unfold contributes the same child layer under each of its input layers,
-    // which is over-approximate but harmless — only the most-derived input actually parents
-    // children at runtime, and the extra synthetic paths are never reached by
-    // counting_layer_hashes_under for any real fold.  The per-input-layer count lets
-    // flush_gates wait for a flush message from every unfold that consumes a given layer
-    // before evaluating done().
+    // Keep each unfold's inputs together: only their most-derived cell parents the children
+    // and receives the unfold flush. Finalization resolves that parent on each applicable branch.
     index_router::unfold_data unfold_layers(declared_unfolds const& unfolds)
     {
       index_router::unfold_data result;
-      for (auto const& n : unfolds | std::views::values) {
-        phlex::experimental::identifier const child_layer{n->child_layer()};
-        std::set<phlex::experimental::identifier> input_layers_for_unfold;
-        for (auto const& input : n->input()) {
-          // Existence of layer already validated
-          auto const& input_layer =
-            static_cast<phlex::experimental::identifier const&>(input.layer);
-          if (not input_layers_for_unfold.insert(input_layer).second) {
-            continue;
-          }
-          ++result.count_per_input_layer[input_layer];
-          result.layer_pairs.push_back({.input = input_layer, .output = child_layer});
-        }
+      result.reserve(unfolds.size());
+
+      using phlex::experimental::identifier;
+      for (auto const& [name, n] : unfolds) {
+        index_router::unfold_layer_spec spec{
+          .name = name,
+          .input_layers =
+            std::views::transform(
+              n->input(), [](auto const& input) -> identifier const& { return input.layer; }) |
+            std::ranges::to<std::vector>(),
+          .output_layer = identifier{n->child_layer()}};
+        result.push_back(std::move(spec));
       }
       return result;
     }
@@ -256,12 +249,6 @@ namespace phlex::detail {
     auto [provider_input_ports, multilayer_join_index_ports] =
       make_computational_edges(nodes_, filters_, graph_, stage_);
 
-    if (provider_input_ports.empty()) {
-      assert(multilayer_join_index_ports.empty());
-      // No algorithms downstream of source.
-      return;
-    }
-
     // Index-router finalization makes edges between the index-set nodes and the provider nodes.
     index_router_.finalize(graph_,
                            fixed_hierarchy_.layer_paths(),