Skip to main content

bb_dsl/
syscalls.rs

1//! DSL-side syscall helpers — record canonical `NodeProto`s into a
2//! `Graph`. The runtime-side dispatch impls live in `bb-ops`; the
3//! two sides agree on stable `(domain, op_type)` string constants
4//! re-exported from `bb_ir::syscall_ids`.
5
6use crate::graph::{attr_tensor, kv, Graph};
7use crate::output::Output;
8use bb_ir::proto::onnx::{tensor_proto, NodeProto, TensorProto};
9
10/// Canonical `(domain, op_type)` string constants the DSL helpers
11/// stamp onto recorded `NodeProto`s. The strings live in
12/// `bb_ir::syscall_ids` so the DSL + compiler + runtime cite one
13/// declaration; this module re-exports them under shorter local
14/// aliases for the helper bodies below.
15pub mod ids {
16    pub use bb_ir::syscall_ids::{
17        OP_CONSTANT as CONSTANT_OP, OP_GATE_DISPATCH as GATE_DISPATCH_OP,
18        OP_PASS_THROUGH as PASS_THROUGH_OP, SYSCALL_DOMAIN,
19    };
20}
21
22/// `(domain, op_type)` registration key for the `ai.bytesandbrains.address_book`
23/// custom-op family. Matches the strings registered in
24/// `bb-ops/src/syscalls/peers/{insert,insert_many,lookup}.rs` via
25/// `register_op!`.
26const ADDRESS_BOOK_DOMAIN: &str = "ai.bytesandbrains.address_book";
27const INSERT_MANY_OP: &str = "InsertMany";
28const LOOKUP_OP: &str = "Lookup";
29
30/// `(domain, op_type)` registration key for the `GlobalRegistryClient`
31/// `Announce` op. Mirrors the literal string registered in
32/// `bb-ops/src/protocols/global_registry/mod.rs` (the constant
33/// `GLOBAL_REGISTRY_DOMAIN`); duplicated here as a string literal
34/// because `bb-dsl` cannot depend on `bb-ops` (`bb-ops` re-exports
35/// `bb-dsl`'s authoring surface).
36const GLOBAL_REGISTRY_DOMAIN: &str = "ai.bytesandbrains.protocol.global_registry";
37const ANNOUNCE_OP: &str = "Announce";
38
39/// `metadata_props` key for the bootstrap-seed label stamped on
40/// `constant` NodeProtos for diagnostics.
41const BOOTSTRAP_SEED_KEY: &str = "ai.bytesandbrains.bootstrap.seed";
42
43/// `metadata_props` key prefix for input-name remapping on the
44/// `Announce` NodeProto. The runtime's `ProtocolRuntime` dispatch
45/// reads these to recover the logical input role (here:
46/// `server_peer`).
47const ANNOUNCE_SERVER_PEER_KEY: &str = "ai.bytesandbrains.input.server_peer";
48
49/// Record a `PassThrough` syscall NodeProto into a `Graph`.
50///
51/// The framework's structural identity op - threads a value through
52/// a partition without doing any compute. Authors reach for it when
53/// a partition needs a non-wire node (e.g. a receiver class that
54/// only forwards values it receives over the wire). The recorded
55/// NodeProto's home class is inferred by the compiler from the
56/// input's home.
57pub fn pass_through(g: &mut Graph, input: Output) -> Output {
58    let out_name = g.next_site_name();
59    g.push_node(NodeProto {
60        op_type: ids::PASS_THROUGH_OP.into(),
61        domain: ids::SYSCALL_DOMAIN.into(),
62        input: vec![input.name],
63        output: vec![out_name.clone()],
64        ..Default::default()
65    });
66    g.declare_value_info(&out_name, input.type_node);
67    Output::new(out_name, input.type_node)
68}
69
70/// Record an `AddressBook::InsertMany(peer, addresses)` custom-op
71/// NodeProto into a `Graph`. New peer creates an entry with
72/// `ref_count = 1`; known peer dedupe-appends every address without
73/// touching `ref_count`. Empty `addresses` vec surfaces as a
74/// dispatch-time `OpError`.
75pub fn address_book_insert_many(g: &mut Graph, peer: Output, addresses: Output) -> Output {
76    let out_name = g.next_site_name();
77    g.push_node(NodeProto {
78        op_type: INSERT_MANY_OP.into(),
79        domain: ADDRESS_BOOK_DOMAIN.into(),
80        input: vec![peer.name, addresses.name],
81        output: vec![out_name.clone()],
82        ..Default::default()
83    });
84    g.declare_value_info(&out_name, &bb_ir::types::TYPE_TRIGGER);
85    Output::new(out_name, &bb_ir::types::TYPE_TRIGGER)
86}
87
88/// Record an `AddressBook::Lookup(peer)` custom-op NodeProto into a
89/// `Graph`. Output carries the full ordered `TYPE_ADDRESS_VEC`;
90/// callers that need a single address pick one downstream. Unknown
91/// or empty-address peer surfaces as a dispatch-time `OpError`.
92pub fn address_book_lookup(g: &mut Graph, peer: Output) -> Output {
93    let out_name = g.next_site_name();
94    g.push_node(NodeProto {
95        op_type: LOOKUP_OP.into(),
96        domain: ADDRESS_BOOK_DOMAIN.into(),
97        input: vec![peer.name],
98        output: vec![out_name.clone()],
99        ..Default::default()
100    });
101    g.declare_value_info(&out_name, &bb_ir::types::TYPE_ADDRESS_VEC);
102    Output::new(out_name, &bb_ir::types::TYPE_ADDRESS_VEC)
103}
104
105/// Record a `GateDispatch` syscall NodeProto into a `Graph` - a
106/// multi-edge synchronization barrier.
107pub fn gate_dispatch(g: &mut Graph, inputs: &[Output]) -> Output {
108    let out_name = g.next_site_name();
109    g.push_node(NodeProto {
110        op_type: ids::GATE_DISPATCH_OP.into(),
111        domain: ids::SYSCALL_DOMAIN.into(),
112        input: inputs.iter().map(|o| o.name.clone()).collect(),
113        output: vec![out_name.clone()],
114        ..Default::default()
115    });
116    g.declare_value_info(&out_name, &bb_ir::types::TYPE_BYTES);
117    Output::new(out_name, &bb_ir::types::TYPE_BYTES)
118}
119
120/// Record a typed `Constant` syscall NodeProto. Bootstrap-stage
121/// constants seed the AddressBook + GlobalRegistry ops with the
122/// server's PeerId and dial bag at install time. The compiler's
123/// `expand_constant` pass requires `value: TensorProto`; this
124/// helper satisfies that contract with an empty tensor sized for
125/// the declared scalar kind.
126///
127/// `label` rides on `metadata_props` under
128/// `ai.bytesandbrains.bootstrap.seed` for diagnostics. `output_type`
129/// is the recorded `&'static TypeNode` consumers downcast on
130/// (`TYPE_PEER_ID`, `TYPE_ADDRESS_VEC`, …).
131pub fn constant(
132    g: &mut Graph,
133    label: &'static str,
134    output_type: &'static bb_ir::types::TypeNode,
135    data_type: tensor_proto::DataType,
136) -> Output {
137    let out_name = g.next_site_name();
138    let tensor = TensorProto {
139        data_type: data_type as i32,
140        dims: vec![1],
141        ..Default::default()
142    };
143    g.push_node(NodeProto {
144        op_type: ids::CONSTANT_OP.into(),
145        domain: ids::SYSCALL_DOMAIN.into(),
146        input: vec![],
147        output: vec![out_name.clone()],
148        attribute: vec![attr_tensor("value", tensor)],
149        metadata_props: vec![kv(BOOTSTRAP_SEED_KEY, label)],
150        ..Default::default()
151    });
152    g.declare_value_info(&out_name, output_type);
153    Output::new(out_name, output_type)
154}
155
156/// Record a `GlobalRegistryClient::Announce` NodeProto. The client
157/// reads `ctx.local_addresses()` automatically, throttles
158/// sub-interval calls to the server's last advertised heartbeat
159/// interval, and merges the server's address bag from the Handshake
160/// reply.
161///
162/// `server_peer` is the `PeerId` `Output` the announcing client
163/// ships its envelope toward (typically a `Constant` of the server's
164/// stable id, recorded with `constant`). The output is typed as
165/// `TYPE_TRIGGER` — downstream nodes can chain on the wakeup.
166pub fn announce(g: &mut Graph, server_peer: Output) -> Output {
167    let out_name = g.next_site_name();
168    g.push_node(NodeProto {
169        op_type: ANNOUNCE_OP.into(),
170        domain: GLOBAL_REGISTRY_DOMAIN.into(),
171        input: vec![server_peer.name.clone()],
172        output: vec![out_name.clone()],
173        metadata_props: vec![kv(ANNOUNCE_SERVER_PEER_KEY, &server_peer.name)],
174        ..Default::default()
175    });
176    g.declare_value_info(&out_name, &bb_ir::types::TYPE_TRIGGER);
177    Output::new(out_name, &bb_ir::types::TYPE_TRIGGER)
178}