use std::collections::{HashMap, HashSet};
use crate::error::CompileError;
use bb_ir::proto::onnx::ModelProto;
const MODULE_CALL_DOMAIN: &str = "ai.bytesandbrains.module";
pub fn verify_no_dangling_calls(model: &ModelProto) -> Result<HashSet<String>, CompileError> {
let mut by_name: HashMap<&str, usize> = HashMap::new();
for (idx, f) in model.functions.iter().enumerate() {
by_name.insert(f.name.as_str(), idx);
}
let Some(root) = model.functions.first() else {
return Ok(HashSet::new());
};
let mut reached: HashSet<String> = HashSet::new();
reached.insert(root.name.clone());
let mut frontier: Vec<usize> = vec![0];
while let Some(idx) = frontier.pop() {
let f = &model.functions[idx];
for node in &f.node {
if node.domain != MODULE_CALL_DOMAIN {
continue;
}
let callee_name = node.op_type.as_str();
let Some(&callee_idx) = by_name.get(callee_name) else {
return Err(CompileError::Internal {
detail: format!(
"function `{}` emits CALL chasing missing function `{}`",
f.name, callee_name,
),
});
};
if reached.insert(callee_name.to_string()) {
frontier.push(callee_idx);
}
}
}
Ok(reached)
}