bb_compiler/
resolve_slots.rs1use std::collections::BTreeMap;
11
12use crate::error::CompileError;
13use bb_ir::proto::onnx::FunctionProto;
14
15pub fn resolve_slots(function: &FunctionProto) -> Result<(), CompileError> {
17 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 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