bb_compiler/
validate_bootstrap_composition.rs

1//! Structural check for the bootstrap-as-function composition tree.
2//!
3//! Every Module's bootstrap recording emits its own FunctionProto
4//! stamped `MODULE_PHASE_BOOTSTRAP`. A parent Module that composes
5//! children via `ModuleCall::bootstrap(g)` records a CALL NodeProto
6//! into its own bootstrap pointing at the child's
7//! `"<child>__bootstrap"` FunctionProto. This pass walks the call
8//! graph rooted at the target's bootstrap and surfaces:
9//!
10//! - [`CompileError::BootstrapCompositionGap`] when a Call points at a
11//!   bootstrap function name that has no matching FunctionProto in the
12//!   model — the engine's FunctionCall dispatch would refuse to seed a
13//!   CallContext for a missing target.
14//! - [`CompileError::BootstrapCompositionCycle`] when the function-call
15//!   graph is not a DAG. Bootstrap completion is one-shot; a cycle in
16//!   the composition tree wedges the engine in `bootstrap_pending`
17//!   forever.
18//!
19//! Runs at the top level of the compile pipeline — before inlining or
20//! partitioning can rearrange the function table — so the error
21//! surfaces against the recorded composition shape the author wrote.
22
23use std::collections::{HashMap, HashSet};
24
25use bb_ir::keys::{read_function_module_phase, MODULE_PHASE_BOOTSTRAP};
26use bb_ir::proto::onnx::ModelProto;
27
28use crate::error::CompileError;
29
30/// CALL NodeProto domain produced by `Graph::with_function`. Same
31/// constant as `verify_no_dangling_calls` / `inline_for_partition` use;
32/// duplicated here so this pass stays self-contained.
33const MODULE_CALL_DOMAIN: &str = "ai.bytesandbrains.module";
34
35/// Walk the bootstrap-call graph rooted at `<target>__bootstrap` and
36/// verify every reachable Call resolves to an existing FunctionProto.
37///
38/// `target_name` is the root function's name (i.e. the body half of
39/// the top-level Module). Models whose root Module has no bootstrap
40/// recording are no-ops — the trait's default `fn bootstrap(&self,
41/// _g: &mut Graph) {}` emits an empty bootstrap that
42/// `Module::build` drops on the floor.
43pub fn validate_bootstrap_composition(
44    model: &ModelProto,
45    target_name: &str,
46) -> Result<(), CompileError> {
47    let by_name: HashMap<&str, &bb_ir::proto::onnx::FunctionProto> = model
48        .functions
49        .iter()
50        .map(|f| (f.name.as_str(), f))
51        .collect();
52
53    let root_bootstrap = format!("{target_name}__bootstrap");
54    let Some(root_fn) = by_name.get(root_bootstrap.as_str()) else {
55        return Ok(());
56    };
57    if read_function_module_phase(root_fn) != Some(MODULE_PHASE_BOOTSTRAP) {
58        // Name collides with the bootstrap convention but the phase
59        // stamp is missing — treat as not-a-bootstrap rather than
60        // misclassify.
61        return Ok(());
62    }
63
64    // Tri-coloured DFS over the bootstrap call graph. `gray` carries
65    // the active recursion stack so a back-edge surfaces the
66    // participating function names in cycle order.
67    let mut black: HashSet<String> = HashSet::new();
68    let mut gray: Vec<String> = Vec::new();
69    walk(&root_bootstrap, &by_name, &mut gray, &mut black)
70}
71
72fn walk(
73    name: &str,
74    by_name: &HashMap<&str, &bb_ir::proto::onnx::FunctionProto>,
75    gray: &mut Vec<String>,
76    black: &mut HashSet<String>,
77) -> Result<(), CompileError> {
78    if black.contains(name) {
79        return Ok(());
80    }
81    if let Some(pos) = gray.iter().position(|n| n == name) {
82        let mut involves: Vec<String> = gray[pos..].to_vec();
83        involves.push(name.to_string());
84        return Err(CompileError::BootstrapCompositionCycle { involves });
85    }
86
87    let Some(function) = by_name.get(name) else {
88        // Caller is responsible for confirming `name` resolves before
89        // calling `walk`; an unresolved name here means a child Call
90        // pointed at a missing bootstrap.
91        return Err(CompileError::BootstrapCompositionGap {
92            caller: gray.last().cloned().unwrap_or_else(|| name.to_string()),
93            target: name.to_string(),
94        });
95    };
96
97    gray.push(name.to_string());
98    for node in &function.node {
99        if node.domain != MODULE_CALL_DOMAIN {
100            continue;
101        }
102        let target = node.op_type.as_str();
103        if target.is_empty() {
104            continue;
105        }
106        if !by_name.contains_key(target) {
107            return Err(CompileError::BootstrapCompositionGap {
108                caller: name.to_string(),
109                target: target.to_string(),
110            });
111        }
112        // Only descend into callees that themselves carry the
113        // bootstrap-phase stamp. A Call from a bootstrap into a body
114        // function is a recording bug the DSL never emits, but
115        // ignoring non-bootstrap callees here keeps the walk
116        // bootstrap-scoped — body-only call cycles surface in
117        // `verify_no_dangling_calls` / the per-graph cycle pass.
118        let callee = by_name[target];
119        if read_function_module_phase(callee) == Some(MODULE_PHASE_BOOTSTRAP) {
120            walk(target, by_name, gray, black)?;
121        }
122    }
123    gray.pop();
124    black.insert(name.to_string());
125    Ok(())
126}
127
bb_compiler/validate_bootstrap_composition.rs

bb_compiler/
validate_bootstrap_composition.rs
