Skip to main content

bb_compiler/
resolve_slots.rs

1//! `resolve_slots` + pre-flight. Final check before a ModelProto is
2//! emitted: rejects a Module where the same role domain hosts BOTH
3//! a concrete-type provider AND a generic `(required_trait,
4//! slot_id)` provider — that combination is `AmbiguousRole`.
5//!
6//! The runner maps `CompileError::AmbiguousRole` to
7//! `BuildError::AmbiguousRole` on the public `Module::build()`
8//! surface.
9
10use std::collections::BTreeMap;
11
12use crate::error::CompileError;
13use bb_ir::proto::onnx::FunctionProto;
14
15/// Resolve slots + run pre-flight. Pure.
16pub fn resolve_slots(function: &FunctionProto) -> Result<(), CompileError> {
17    // Per role: collect distinct concrete_type providers + distinct
18    // (required_trait, slot_id) generic providers.
19    let mut role_providers: BTreeMap<String, RoleProviders> = BTreeMap::new();
20
21    for node in &function.node {
22        if !node.domain.starts_with("ai.bytesandbrains.role.") {
23            continue;
24        }
25        let providers = role_providers.entry(node.domain.clone()).or_default();
26        if let Some(concrete) = meta_value(node, "ai.bytesandbrains.concrete_type") {
27            providers.concrete_types.insert(concrete.to_string());
28        }
29        if let (Some(rt), Some(sid)) = (
30            meta_value(node, "ai.bytesandbrains.required_trait"),
31            meta_value(node, "ai.bytesandbrains.slot_id"),
32        ) {
33            if let Ok(id) = sid.parse::<u32>() {
34                providers.generic_slots.insert(id, rt.to_string());
35            }
36        }
37    }
38
39    for (role, providers) in role_providers {
40        if !providers.concrete_types.is_empty() && !providers.generic_slots.is_empty() {
41            // First concrete + first generic slot are surfaced as
42            // canonical witnesses. The pass is deterministic because
43            // BTreeMap iteration is ordered.
44            let concrete_type = providers
45                .concrete_types
46                .iter()
47                .next()
48                .cloned()
49                .unwrap_or_default();
50            let (slot_id, _trait_name) = providers
51                .generic_slots
52                .iter()
53                .next()
54                .map(|(k, v)| (*k, v.clone()))
55                .unwrap_or_default();
56            return Err(CompileError::AmbiguousRole {
57                role,
58                concrete_type,
59                generic_slot_id: slot_id,
60            });
61        }
62    }
63
64    Ok(())
65}
66
67#[derive(Default)]
68struct RoleProviders {
69    concrete_types: std::collections::BTreeSet<String>,
70    generic_slots: BTreeMap<u32, String>,
71}
72
73fn meta_value<'a>(node: &'a bb_ir::proto::onnx::NodeProto, key: &str) -> Option<&'a str> {
74    node.metadata_props
75        .iter()
76        .find(|p| p.key == key)
77        .map(|p| p.value.as_str())
78}
79