Skip to main content

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}