bb_compiler/rx_chain.rs
1//! RX-chain head bookkeeping shared by the RX gate-insertion passes.
2//!
3//! Each synthesized `wire::Recv` carries a `RX_CHAIN_HEAD_KEY`
4//! metadata entry naming the currently-effective output that
5//! downstream consumers read from. Initial value is the Recv's own
6//! `output[0]`; each RX gate pass replaces it with the gate's
7//! output. Successive passes consult this to chain themselves in
8//! order: `Recv → DedupGateRx → PeerHealthGateRx → BackoffGateRx`.
9
10use bb_ir::proto::onnx::{NodeProto, StringStringEntryProto};
11
12/// Metadata key on a `wire::Recv` carrying its current RX-chain head.
13pub const RX_CHAIN_HEAD_KEY: &str = "ai.bytesandbrains.rx_chain_head";
14
15/// Current RX-chain head for `recv` - the output name downstream
16/// consumers read from. Defaults to `recv.output[0]` when no gate
17/// has yet been inserted.
18pub fn rx_chain_head(recv: &NodeProto) -> String {
19 recv.metadata_props
20 .iter()
21 .find(|p| p.key == RX_CHAIN_HEAD_KEY)
22 .map(|p| p.value.clone())
23 .unwrap_or_else(|| recv.output.first().cloned().unwrap_or_default())
24}
25
26/// Set the RX-chain head metadata on `recv` to `new_head`.
27pub fn set_rx_chain_head(recv: &mut NodeProto, new_head: &str) {
28 if let Some(existing) = recv
29 .metadata_props
30 .iter_mut()
31 .find(|p| p.key == RX_CHAIN_HEAD_KEY)
32 {
33 existing.value = new_head.to_string();
34 return;
35 }
36 recv.metadata_props.push(StringStringEntryProto {
37 key: RX_CHAIN_HEAD_KEY.to_string(),
38 value: new_head.to_string(),
39 });
40}