Skip to main content

bb_dsl/
graph.rs

1//! Recording context wrapping the in-progress `FunctionProto`. The
2//! proto is the IR — semantic BB attributes ride on proto fields,
3//! not a parallel Rust shadow store. See `docs/IR_AND_DSL.md` §2.
4//!
5//! Rust-side wrapper carries what the proto can't represent:
6//! `instance_for_pointer` (pointer-identity dedup for generic
7//! placeholders) and `site_counter` (output-name minting cache).
8//!
9//! `Module::build()` constructs `Graph` automatically; `Graph::new()`
10//! is for acceptance tests.
11
12use std::any::TypeId;
13use std::collections::HashMap;
14
15use crate::output::Output;
16use bb_ir::proto::onnx::tensor_proto::DataType as DT;
17use bb_ir::proto::onnx::{
18    attribute_proto, type_proto, AttributeProto, FunctionProto, NodeProto, StringStringEntryProto,
19    TensorShapeProto, TypeProto, ValueInfoProto,
20};
21use bb_ir::types::TypeNode;
22
23use crate::recorded::RecordedModule;
24
25/// Composition-hierarchy chain stamped by [`Graph::with_function`].
26/// Read by the compiler's partition naming.
27const MODULE_INSTANCE_KEY: &str = "ai.bytesandbrains.module_instance";
28
29fn upsert_metadata(props: &mut Vec<StringStringEntryProto>, key: &str, value: &str) {
30    if let Some(entry) = props.iter_mut().find(|p| p.key == key) {
31        entry.value = value.to_string();
32    } else {
33        props.push(StringStringEntryProto {
34            key: key.to_string(),
35            value: value.to_string(),
36        });
37    }
38}
39
40/// Recording context every DSL method writes into.
41pub struct Graph {
42    /// The IR body.
43    function: FunctionProto,
44
45    /// Output-name counter.
46    site_counter: u64,
47
48    /// Pointer-identity dedup. `TypeId` discriminator avoids
49    /// collapsing distinct ZST placeholders (all ZSTs share an
50    /// address); same-type ZSTs still alias — documented in
51    /// `docs/DEPLOYMENT.md`.
52    instance_for_pointer: HashMap<(TypeId, *const ()), u32>,
53    next_instance_id: u32,
54
55    /// Active `with_function` scopes; joined chain stamped onto
56    /// each NodeProto's `MODULE_INSTANCE_KEY`. Empty stack → default
57    /// to `@default`.
58    module_scope: Vec<String>,
59
60    /// Nested-`with_function` FunctionProtos. The top-level
61    /// `Module::op` wrap folds into the root `function` instead of
62    /// creating an entry here.
63    sub_functions: Vec<FunctionProto>,
64
65    /// Recording-target stack. `None`/empty → root `function`.
66    recording_target: Vec<Option<usize>>,
67
68    /// `true` once any `with_function` has fired. The top-level
69    /// wrap-at-depth-1 check uses this to keep the body in
70    /// `function[0]` rather than synthesizing a CALL.
71    has_seen_function: bool,
72
73    /// Typed recorder errors; `Module::build` surfaces the first.
74    pending_errors: Vec<crate::module::BuildError>,
75
76    /// Recording-mode stack. `with_function` pushes `Sealed` so
77    /// inner `input()` calls don't leak formals into the outer
78    /// FunctionProto. Empty → `Open`.
79    mode_stack: Vec<RecordingMode>,
80
81    /// `formal_name → actual_handle` bindings pre-loaded by
82    /// `ModuleCall::input`. `g.input("name")` consults the top of
83    /// the stack first.
84    formal_binding_stack: Vec<HashMap<String, Output>>,
85
86    /// `(target_idx, name) → (handle, TypeNode)` for ports
87    /// registered via `output()`. Idempotent: at most one
88    /// PassThrough producer per name. `usize::MAX` indexes the
89    /// root function.
90    named_output_types: HashMap<(usize, String), (Output, &'static TypeNode)>,
91}
92
93/// Open top-level vs. sealed nested-function recording.
94#[derive(Clone, Copy, Debug, PartialEq, Eq)]
95pub enum RecordingMode {
96    /// `input()` propagates to root + every active sub-function.
97    Open,
98    /// `input()` lands only on the immediate sub-function.
99    Sealed,
100}
101
102impl Graph {
103    /// Empty `Graph`. `Module::build()` wraps the body in
104    /// `with_function(self.name(), ...)` automatically.
105    pub fn new() -> Self {
106        Self {
107            function: FunctionProto::default(),
108            site_counter: 0,
109            instance_for_pointer: HashMap::new(),
110            next_instance_id: 0,
111            module_scope: Vec::new(),
112            sub_functions: Vec::new(),
113            recording_target: Vec::new(),
114            has_seen_function: false,
115            pending_errors: Vec::new(),
116            mode_stack: Vec::new(),
117            named_output_types: HashMap::new(),
118            formal_binding_stack: Vec::new(),
119        }
120    }
121
122    /// Current recording mode; `Open` when stack empty.
123    fn current_mode(&self) -> RecordingMode {
124        self.mode_stack
125            .last()
126            .copied()
127            .unwrap_or(RecordingMode::Open)
128    }
129
130    /// Drain accumulated recorder errors.
131    pub fn take_pending_errors(&mut self) -> Vec<crate::module::BuildError> {
132        std::mem::take(&mut self.pending_errors)
133    }
134
135    /// Register a named output port. Idempotent — a second call for
136    /// the same `(target_idx, name)` returns the prior handle.
137    pub fn output(&mut self, name: &str, handle: Output) {
138        let target_idx = self
139            .recording_target
140            .last()
141            .and_then(|t| *t)
142            .unwrap_or(usize::MAX);
143        let key = (target_idx, name.to_string());
144        if self.named_output_types.contains_key(&key) {
145            return;
146        }
147        let type_node = handle.type_node;
148
149        // PassThrough renames the producer's value to the port name
150        // so it appears as a NodeProto output for downstream passes.
151        self.push_node(NodeProto {
152            op_type: bb_ir::syscall_ids::OP_PASS_THROUGH.into(),
153            domain: bb_ir::syscall_ids::SYSCALL_DOMAIN.into(),
154            input: vec![handle.name.clone()],
155            output: vec![name.to_string()],
156            ..Default::default()
157        });
158
159        let function: &mut FunctionProto = match target_idx {
160            usize::MAX => &mut self.function,
161            idx => &mut self.sub_functions[idx],
162        };
163        if function.output.iter().all(|n| n != name) {
164            function.output.push(name.to_string());
165            function
166                .value_info
167                .push(type_meta_to_value_info(name, type_node));
168        }
169        let registered = Output::new(name.to_string(), type_node);
170        self.named_output_types.insert(key, (registered, type_node));
171    }
172
173    /// Emit a `wire.Send` to `peers` (a `Vec<PeerId>` at dispatch)
174    /// and register `name` as a network output. The compiler's
175    /// `partition_by_wire_ops` cuts here; `synthesize_wire_recvs`
176    /// materializes the matching `wire.Recv`.
177    pub fn net_out(&mut self, name: &str, peers: Output, value: Output) {
178        let value_type = value.type_node;
179        let port_name = name.to_string();
180        let handle_name = self.next_site_name();
181
182        let target_idx = self
183            .recording_target
184            .last()
185            .and_then(|t| *t)
186            .unwrap_or(usize::MAX);
187        let key = (target_idx, port_name.clone());
188        let already_registered = self.named_output_types.contains_key(&key);
189
190        self.push_node(NodeProto {
191            op_type: bb_ir::syscall_ids::OP_WIRE_SEND.into(),
192            domain: bb_ir::syscall_ids::WIRE_DOMAIN.into(),
193            input: vec![value.name.clone(), peers.name],
194            output: vec![port_name.clone(), handle_name.clone()],
195            ..Default::default()
196        });
197        self.declare_value_info(&port_name, value_type);
198        self.declare_value_info(&handle_name, &bb_ir::types::TYPE_WIRE_REQ_ID);
199
200        if !already_registered {
201            let function: &mut FunctionProto = match target_idx {
202                usize::MAX => &mut self.function,
203                idx => &mut self.sub_functions[idx],
204            };
205            if function.output.iter().all(|n| n != &port_name) {
206                function.output.push(port_name.clone());
207            }
208            let handle = Output::new(port_name.clone(), value_type);
209            self.named_output_types.insert(key, (handle, value_type));
210        }
211    }
212
213    /// Pack N typed Outputs into one composite. Receiver pairs
214    /// with [`Self::unbundle`]. Composite envelope is
215    /// [`bb_ir::types::TYPE_COMPOSITE`]. Panics on empty `parts`.
216    pub fn bundle(&mut self, parts: &[Output]) -> Output {
217        assert!(
218            !parts.is_empty(),
219            "Graph::bundle: parts slice is empty; need >= 1 child Output",
220        );
221        let bundle_name = self.next_site_name();
222        let inputs: Vec<String> = parts.iter().map(|p| p.name.clone()).collect();
223
224        let child_count = parts.len();
225        let child_types = parts
226            .iter()
227            .map(|p| p.type_node.denotation)
228            .collect::<Vec<_>>()
229            .join(",");
230
231        self.push_node(NodeProto {
232            op_type: "Bundle".into(),
233            domain: "ai.bytesandbrains.composite".into(),
234            input: inputs,
235            output: vec![bundle_name.clone()],
236            attribute: vec![
237                attr_int(
238                    "ai.bytesandbrains.composite.child_count",
239                    child_count as i64,
240                ),
241                attr_string("ai.bytesandbrains.composite.child_types", &child_types),
242            ],
243            ..Default::default()
244        });
245        self.declare_value_info(&bundle_name, &bb_ir::types::TYPE_COMPOSITE);
246        Output::new(bundle_name, &bb_ir::types::TYPE_COMPOSITE)
247    }
248
249    /// Extract a composite Output back into its N child Outputs.
250    /// `part_types` declares the expected child TypeNodes positionally;
251    /// the runtime op validates the envelope's child count against the
252    /// declared length and emits one `BytesValue`-shaped output per
253    /// child, named `child_{i}` and typed against `part_types[i]` via
254    /// the stamped `ValueInfoProto.denotation`. Downstream consumers
255    /// decode against that denotation, matching the wire.Recv pattern.
256    ///
257    /// Panics with a recording-time error if `part_types` is empty —
258    /// the matching `g.bundle` cannot have produced a zero-child
259    /// envelope.
260    pub fn unbundle(&mut self, composite: Output, part_types: &[&'static TypeNode]) -> Vec<Output> {
261        assert!(
262            !part_types.is_empty(),
263            "Graph::unbundle: part_types slice is empty; need >= 1 declared child type",
264        );
265        let child_count = part_types.len();
266        let port_names: Vec<String> = (0..child_count).map(|_| self.next_site_name()).collect();
267        let child_types = part_types
268            .iter()
269            .map(|t| t.denotation)
270            .collect::<Vec<_>>()
271            .join(",");
272
273        self.push_node(NodeProto {
274            op_type: "Unbundle".into(),
275            domain: "ai.bytesandbrains.composite".into(),
276            input: vec![composite.name],
277            output: port_names.clone(),
278            attribute: vec![
279                attr_int(
280                    "ai.bytesandbrains.composite.child_count",
281                    child_count as i64,
282                ),
283                attr_string("ai.bytesandbrains.composite.child_types", &child_types),
284            ],
285            ..Default::default()
286        });
287        for (port_name, type_node) in port_names.iter().zip(part_types.iter()) {
288            self.declare_value_info(port_name, type_node);
289        }
290        port_names
291            .into_iter()
292            .zip(part_types.iter())
293            .map(|(name, t)| Output::new(name, t))
294            .collect()
295    }
296
297    /// Look up a previously-registered output port by name on the
298    /// current recording target. Returns `None` when neither
299    /// `output(name, ...)` nor an enclosing scope has registered
300    /// the port — callers report `BuildError::MissingOutputPort`.
301    pub fn lookup_output(&self, name: &str) -> Option<Output> {
302        let target_idx = self
303            .recording_target
304            .last()
305            .and_then(|t| *t)
306            .unwrap_or(usize::MAX);
307        self.named_output_types
308            .get(&(target_idx, name.to_string()))
309            .map(|(h, _)| h.clone())
310    }
311
312    /// Push a [`crate::module::BuildError`] onto the recorder's
313    /// pending-errors queue. Used by methods that must keep their
314    /// existing return shape (e.g. `Graph::wire` returns the typed
315    /// output triple) but want a typed-error escape from a panic.
316    pub fn record_build_error(&mut self, err: crate::module::BuildError) {
317        self.pending_errors.push(err);
318    }
319
320    /// Mutable view of whichever FunctionProto the recorder is
321    /// currently writing into. Either the root `function` or one of
322    /// the `sub_functions` (per ).
323    fn current_function_mut(&mut self) -> &mut FunctionProto {
324        match self.recording_target.last() {
325            Some(Some(idx)) => &mut self.sub_functions[*idx],
326            _ => &mut self.function,
327        }
328    }
329
330    /// Extract the recorded function body for the compiler to
331    /// consume. Called by `Module::build()` after `module.op()`
332    /// returns. The chosen-path install constructs concrete
333    /// instances via the inventory's `construct_fn` at install
334    /// time; the IR carries no instance state.
335    pub fn finish(self) -> RecordedModule {
336        RecordedModule {
337            function: self.function,
338            sub_functions: self.sub_functions,
339        }
340    }
341
342    /// Pointer-identity-keyed slot allocation for generic
343    /// placeholders. Appends `"__slot_<slot_id>"` to
344    /// `FunctionProto.attribute` on first encounter.
345    pub fn register_generic<T: 'static>(
346        &mut self,
347        instance: &T,
348        _required_trait: &'static str,
349    ) -> u32 {
350        let key = (TypeId::of::<T>(), (instance as *const T).cast::<()>());
351        if let Some(&id) = self.instance_for_pointer.get(&key) {
352            return id;
353        }
354        let id = self.next_instance_id;
355        self.next_instance_id += 1;
356        self.instance_for_pointer.insert(key, id);
357        self.current_function_mut()
358            .attribute
359            .push(format!("__slot_{id}"));
360        id
361    }
362
363    /// Declare a Module input by name. Lands with
364    /// [`bb_ir::types::TYPE_BYTES`] sentinel; the TypeSolver
365    /// narrows it later. Propagates up the recording-target chain
366    /// in `Open` mode so enclosing CALL NodeProtos stay
367    /// referenceable.
368    pub fn input(&mut self, name: &str) -> Output {
369        // Promote the actual's type into the formal's value_info
370        // when the fluent builder pre-loaded a binding.
371        let bound_type = self
372            .formal_binding_stack
373            .last()
374            .and_then(|m| m.get(name))
375            .map(|h| h.type_node);
376
377        let build_vi = |name: &str| match bound_type {
378            Some(type_node) => type_meta_to_value_info(name, type_node),
379            None => opaque_value_info(name),
380        };
381
382        // Sealed → write only the immediate sub-function. Open →
383        // root + every active sub-function.
384        let active_targets: Vec<Option<usize>> = match self.current_mode() {
385            RecordingMode::Sealed => match self.recording_target.last() {
386                Some(slot) => vec![*slot],
387                None => Vec::new(),
388            },
389            RecordingMode::Open => self.recording_target.to_vec(),
390        };
391
392        let mut seen_root = false;
393        let touch_root = matches!(self.current_mode(), RecordingMode::Open);
394        for target in active_targets
395            .iter()
396            .chain(std::iter::once(&None).take(if touch_root { 1 } else { 0 }))
397        {
398            let function: &mut FunctionProto = match target {
399                Some(idx) => &mut self.sub_functions[*idx],
400                None => {
401                    if seen_root {
402                        continue;
403                    }
404                    seen_root = true;
405                    &mut self.function
406                }
407            };
408            if function.input.iter().all(|n| n != name) {
409                function.input.push(name.to_string());
410                function.value_info.push(build_vi(name));
411            }
412        }
413
414        Output::new(name.to_string(), &bb_ir::types::TYPE_BYTES)
415    }
416
417    /// Allocate a fresh value-name. Monotonic counter; format
418    /// `"v<n>"`.
419    pub fn next_site_name(&mut self) -> String {
420        let n = self.site_counter;
421        self.site_counter += 1;
422        format!("v{n}")
423    }
424
425    /// Stamp a `ValueInfoProto` for `name` on the current target.
426    /// Idempotent. Recorders call this on every minted output.
427    pub fn declare_value_info(&mut self, name: &str, type_node: &'static bb_ir::types::TypeNode) {
428        let function = self.current_function_mut();
429        if function.value_info.iter().any(|v| v.name == name) {
430            return;
431        }
432        function
433            .value_info
434            .push(type_meta_to_value_info(name, type_node));
435    }
436
437    /// Push a NodeProto into the active target. Stamps
438    /// `MODULE_INSTANCE_KEY` with the joined `with_function` chain;
439    /// an existing stamp is prefixed (preserves replayed hierarchy).
440    pub fn push_node(&mut self, mut node: NodeProto) {
441        if !self.module_scope.is_empty() {
442            let prefix = self.module_scope.join("_");
443            let existing = node
444                .metadata_props
445                .iter()
446                .find(|p| p.key == MODULE_INSTANCE_KEY)
447                .map(|p| p.value.clone());
448            let combined = match existing {
449                Some(inner) if !inner.is_empty() => format!("{prefix}_{inner}"),
450                _ => prefix,
451            };
452            upsert_metadata(&mut node.metadata_props, MODULE_INSTANCE_KEY, &combined);
453        }
454        self.current_function_mut().node.push(node);
455    }
456
457    /// Record `body` into a sub-FunctionProto named `name` and emit
458    /// a CALL in the outer target. Top-level wraps fold the body
459    /// into `function[0]` instead of synthesizing a CALL.
460    ///
461    /// `bindings` pre-loads formal→actual handles so
462    /// `g.input(formal)` inside `body` returns the actual.
463    ///
464    /// Returns `(child_port_name, parent_call_output_name)` pairs
465    /// for non-top-level wraps; empty for top-level (`g.output`
466    /// registers directly in the parent scope).
467    pub fn with_function<F>(
468        &mut self,
469        name: &str,
470        bindings: &[(String, Output)],
471        body: F,
472    ) -> Vec<(String, String)>
473    where
474        F: FnOnce(&mut Graph),
475    {
476        // Top-level wrap iff first call AND root function untouched.
477        let is_top_level_wrap = !self.has_seen_function
478            && self.recording_target.is_empty()
479            && self.function.node.is_empty()
480            && self.function.input.is_empty()
481            && self.function.attribute_proto.is_empty();
482
483        self.has_seen_function = true;
484
485        if is_top_level_wrap {
486            // Body becomes the entry; no CALL emitted.
487            self.function.name = name.to_string();
488            self.module_scope.push(name.to_string());
489            let depth = self.module_scope.len();
490            body(self);
491            debug_assert_eq!(
492                self.module_scope.len(),
493                depth,
494                "with_function body must not mutate the scope stack",
495            );
496            self.module_scope.pop();
497            return Vec::new();
498        }
499
500        // Existing same-name sub-function → record into a scratch
501        // slot and discard, so the canonical FunctionProto stays
502        // shared across parents.
503        let target_idx = if let Some(idx) = self.sub_functions.iter().position(|f| f.name == name) {
504            idx
505        } else {
506            let new_idx = self.sub_functions.len();
507            self.sub_functions.push(FunctionProto {
508                name: name.to_string(),
509                ..Default::default()
510            });
511            new_idx
512        };
513
514        let is_duplicate = target_idx + 1 != self.sub_functions.len();
515        let recording_idx = if is_duplicate {
516            let scratch_idx = self.sub_functions.len();
517            self.sub_functions.push(FunctionProto::default());
518            scratch_idx
519        } else {
520            target_idx
521        };
522
523        // Bind formals so `g.input(formal)` returns the actual.
524        let binding_map: HashMap<String, Output> = bindings
525            .iter()
526            .map(|(name, h)| (name.clone(), h.clone()))
527            .collect();
528        self.formal_binding_stack.push(binding_map);
529
530        self.recording_target.push(Some(recording_idx));
531        self.module_scope.push(name.to_string());
532        // seal the recorder while the body runs so any
533        // `input()` calls land only on this sub-function and don't
534        // leak out into the root function above.
535        self.mode_stack.push(RecordingMode::Sealed);
536        let depth = self.module_scope.len();
537        body(self);
538        debug_assert_eq!(
539            self.module_scope.len(),
540            depth,
541            "with_function body must not mutate the scope stack",
542        );
543        self.mode_stack.pop();
544        self.module_scope.pop();
545        self.recording_target.pop();
546        self.formal_binding_stack.pop();
547
548        // The body's declared outputs (via `g.output(name, value)` /
549        // `g.net_out(name, peers, value)`) already populated the
550        // sub-function's `output[]` + `value_info[]` lists. Snapshot
551        // the recorded output names so the CALL NodeProto's
552        // positional output slots match the sub-function's
553        // declarations.
554        let recorded_outputs: Vec<String> = self.sub_functions[recording_idx].output.clone();
555
556        if is_duplicate {
557            self.sub_functions.pop();
558        }
559
560        // Emit the CALL NodeProto in the parent scope. Input list is
561        // the actuals' names (parent-scope), positionally aligned
562        // with the sub-function's declared input ports; output list
563        // is freshly minted outer-scope names — one per
564        // sub-function output.
565        let final_name = self.sub_functions[target_idx].name.clone();
566        let call_inputs: Vec<String> = bindings.iter().map(|(_, h)| h.name.clone()).collect();
567        let call_outputs: Vec<String> = (0..recorded_outputs.len())
568            .map(|_| self.next_site_name())
569            .collect();
570        let call = NodeProto {
571            op_type: final_name,
572            domain: "ai.bytesandbrains.module".into(),
573            input: call_inputs,
574            output: call_outputs.clone(),
575            ..Default::default()
576        };
577        self.push_node(call);
578
579        recorded_outputs.into_iter().zip(call_outputs).collect()
580    }
581
582    /// Read-only view of the recorded `FunctionProto`. 's
583    /// compiler + the acceptance tests read everything from here -
584    /// the proto is the single source of truth.
585    pub fn function(&self) -> &FunctionProto {
586        &self.function
587    }
588
589    /// Read-only view of the nested `with_function` sub-functions
590    /// (). Test-only accessor; the
591    /// canonical hand-off is via `Graph::finish() -> RecordedModule`.
592    #[cfg(test)]
593    pub(crate) fn sub_functions_for_test(&self) -> &[FunctionProto] {
594        &self.sub_functions
595    }
596}
597
598impl Default for Graph {
599    fn default() -> Self {
600        Self::new()
601    }
602}
603
604/// Stage-5 canonical mapping `TypeNode → ValueInfoProto.type`. Used
605/// by [`Graph::input`] so the compiler's `validate` Rule 5 finds a
606/// type declaration on every Module input.
607///
608/// Tensor-typed denotations (`ai.bytesandbrains.tensor.*`) map to a
609/// `TensorType` carrying the matching ONNX `DataType` elem_type;
610/// everything else maps to an `OpaqueType` keyed by denotation +
611/// the `ai.bytesandbrains` domain. 's `bb.wire v1` constellation
612/// may refine the mapping further.
613/// Build a `ValueInfoProto` whose `TypeProto` is the canonical
614/// opaque placeholder. Used by the single-arg [`Graph::input`]
615/// API where authors declare ports by name only and the
616/// compiler's TypeSolver narrows the type from connected ops.
617fn opaque_value_info(name: &str) -> bb_ir::proto::onnx::ValueInfoProto {
618    type_meta_to_value_info(name, &bb_ir::types::TYPE_BYTES)
619}
620
621fn type_meta_to_value_info(
622    name: &str,
623    type_node: &'static TypeNode,
624) -> bb_ir::proto::onnx::ValueInfoProto {
625    let value = if let Some(elem_type) = tensor_elem_from_denotation(type_node.denotation) {
626        // Without per-shape metadata on `TypeNode`, every recorded
627        // tensor declares an unconstrained shape — downstream
628        // type-checking treats absence as "broadcastable", and
629        // per-instance shape is the host's responsibility.
630        type_proto::Value::TensorType(type_proto::Tensor {
631            elem_type,
632            shape: Some(TensorShapeProto::default()),
633        })
634    } else {
635        type_proto::Value::OpaqueType(type_proto::Opaque {
636            domain: "ai.bytesandbrains".into(),
637            name: type_node.denotation.into(),
638        })
639    };
640
641    ValueInfoProto {
642        name: name.to_string(),
643        r#type: Some(TypeProto {
644            value: Some(value),
645            denotation: type_node.denotation.into(),
646        }),
647        ..Default::default()
648    }
649}
650
651fn tensor_elem_from_denotation(denotation: &str) -> Option<i32> {
652    Some(match denotation {
653        "ai.bytesandbrains.tensor.f32" => DT::Float as i32,
654        "ai.bytesandbrains.tensor.f64" => DT::Double as i32,
655        "ai.bytesandbrains.tensor.i32" => DT::Int32 as i32,
656        "ai.bytesandbrains.tensor.i64" => DT::Int64 as i32,
657        "ai.bytesandbrains.tensor.bool" => DT::Bool as i32,
658        _ if denotation.starts_with("ai.bytesandbrains.tensor.") => DT::Undefined as i32,
659        _ => return None,
660    })
661}
662
663/// Construct a `StringStringEntryProto` for `metadata_props` or
664/// `attribute_proto.metadata_props`. Used by every DSL method body.
665pub fn kv(key: &str, value: &str) -> StringStringEntryProto {
666    StringStringEntryProto {
667        key: key.to_string(),
668        value: value.to_string(),
669    }
670}
671
672/// Construct an `AttributeProto` of type `INT` for `NodeProto.attribute`.
673/// Used by DSL methods that pass scalar `i64` config (`axis`, `group`,
674/// `to`, etc.).
675pub fn attr_int(name: &str, value: i64) -> AttributeProto {
676    AttributeProto {
677        name: name.to_string(),
678        r#type: attribute_proto::AttributeType::Int as i32,
679        i: value,
680        ..Default::default()
681    }
682}
683
684/// Construct an `AttributeProto` of type `FLOAT`. Used for `epsilon`,
685/// `alpha`, `momentum`, etc.
686pub fn attr_float(name: &str, value: f32) -> AttributeProto {
687    AttributeProto {
688        name: name.to_string(),
689        r#type: attribute_proto::AttributeType::Float as i32,
690        f: value,
691        ..Default::default()
692    }
693}
694
695/// Construct an `AttributeProto` of type `INTS`. Used for shape /
696/// axes / strides / kernel_shape / pads / dilations / perm vectors.
697pub fn attr_ints(name: &str, values: Vec<i64>) -> AttributeProto {
698    AttributeProto {
699        name: name.to_string(),
700        r#type: attribute_proto::AttributeType::Ints as i32,
701        ints: values,
702        ..Default::default()
703    }
704}
705
706/// Construct an `AttributeProto` of type `GRAPH`. ONNX carries
707/// `If` / `Loop` body sub-graphs in `AttributeProto.g`.
708pub fn attr_graph(name: &str, value: bb_ir::proto::onnx::GraphProto) -> AttributeProto {
709    AttributeProto {
710        name: name.to_string(),
711        r#type: attribute_proto::AttributeType::Graph as i32,
712        g: Some(value),
713        ..Default::default()
714    }
715}
716
717/// Construct an `AttributeProto` of type `STRING`. The proto stores
718/// strings as `s: Vec<u8>`; this helper hides the bytes encoding.
719/// Used by ops carrying structured-string metadata
720/// (e.g. comma-separated TypeNode denotation lists on `composite`
721/// Bundle / Unbundle).
722pub fn attr_string(name: &str, value: &str) -> AttributeProto {
723    AttributeProto {
724        name: name.to_string(),
725        r#type: attribute_proto::AttributeType::String as i32,
726        s: value.as_bytes().to_vec(),
727        ..Default::default()
728    }
729}
730
731/// Construct an `AttributeProto` of type `TENSOR`. Used by `Constant`
732/// for embedded literal payloads.
733pub fn attr_tensor(name: &str, value: bb_ir::proto::onnx::TensorProto) -> AttributeProto {
734    AttributeProto {
735        name: name.to_string(),
736        r#type: attribute_proto::AttributeType::Tensor as i32,
737        t: Some(value),
738        ..Default::default()
739    }
740}
741