use crate::pipeline::{read_stage_signature, PipelineError, ResolvedInput};
use cuttlefish_abi::Ty;
use cuttlefish_core::graph::{Branches, InputExpr, NodeGraph};
use std::collections::{BTreeMap, HashMap, HashSet};
use wasmtime::Engine;
#[derive(Clone)]
pub struct CheckedNode {
pub name: String,
pub kind: crate::catalog::ArtifactKind,
pub resolved: Option<String>,
pub module_bytes: Vec<u8>,
pub signature: cuttlefish_abi::Signature,
pub input: Option<InputExpr>,
pub repeat_until: Option<String>,
pub max_iterations: Option<u32>,
pub script: Option<String>,
pub over: Option<std::path::PathBuf>,
pub item_output: Option<cuttlefish_abi::Ty>,
pub accept: Vec<cuttlefish_core::graph::AcceptCheck>,
pub on_fail: Vec<cuttlefish_core::graph::Rung>,
}
pub fn fanout_collection_ty() -> Ty {
Ty::Record(BTreeMap::from([
("results_path".to_string(), Ty::Text),
("failures_path".to_string(), Ty::Text),
("succeeded".to_string(), Ty::Json),
("failed".to_string(), Ty::Json),
]))
}
pub struct CheckedGraph {
pub nodes: Vec<CheckedNode>,
pub exclusive_to: HashMap<String, BranchExclusivity>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BranchExclusivity {
pub decision: String,
pub label: String,
}
#[derive(Debug, thiserror::Error)]
pub enum DagError {
#[error(transparent)]
Pipeline(#[from] PipelineError),
#[error("node `{node}` references unknown node `{referenced}`")]
UnknownReference {
node: String,
referenced: String,
},
#[error(
"node `{node}` has an inbound edge that would form a cycle with no \
repeat_until marker — add one (with max_iterations) if this loop is intentional"
)]
UnmarkedCycle {
node: String,
},
#[error(
"node `{node}`'s input mixes output from branch label `{label_a}` and \
label `{label_b}` of the same `branches.{decision}` decision — a node \
cannot depend on more than one mutually-exclusive branch outcome at once"
)]
ConflictingBranchFanIn {
node: String,
decision: String,
label_a: String,
label_b: String,
},
#[error("node `{consumer}` needs {expected}, but `{producer}` produces {produced}")]
SeamMismatch {
producer: String,
produced: String,
consumer: String,
expected: String,
},
#[error("a graph needs at least one node")]
Empty,
}
pub fn check_graph(
engine: &Engine,
graph: &NodeGraph,
branches: &Branches,
resolved: &HashMap<String, ResolvedInput>,
) -> Result<CheckedGraph, DagError> {
if graph.nodes.is_empty() {
return Err(DagError::Empty);
}
let mut signatures = HashMap::new();
let mut item_outputs: HashMap<String, cuttlefish_abi::Ty> = HashMap::new();
for (name, node) in &graph.nodes {
let input = resolved.get(name).expect("caller resolved every node");
let mut signature = read_stage_signature(engine, input)?;
if node.over.is_some() {
item_outputs.insert(name.clone(), signature.output.clone());
signature.output = fanout_collection_ty();
}
signatures.insert(name.clone(), signature);
}
let order = topological_order(graph)?;
let exclusive_to = compute_branch_exclusivity(graph, branches)?;
let mut nodes = Vec::with_capacity(order.len());
for name in &order {
let node = graph
.get(name)
.expect("topological_order only returns known nodes");
let input_resolved = resolved.get(name).expect("caller resolved every node");
let signature = signatures.get(name).unwrap().clone();
if let Some(expr) = &node.input {
let produced = evaluate_expr_ty(expr, &signatures, graph)?;
if !produced.assignable_to(&signature.input) {
let (producer, produced_str) = describe_expr(expr, &signatures);
return Err(DagError::SeamMismatch {
producer,
produced: produced_str,
consumer: name.clone(),
expected: signature.input.to_string(),
});
}
}
nodes.push(CheckedNode {
name: name.clone(),
kind: input_resolved.kind,
resolved: input_resolved.resolved.clone(),
module_bytes: input_resolved.bytes.clone(),
signature,
input: node.input.clone(),
repeat_until: node.repeat_until.clone(),
max_iterations: node.max_iterations,
script: input_resolved.script.clone(),
over: node.over.clone(),
item_output: item_outputs.get(name).cloned(),
accept: node.accept.clone(),
on_fail: node.on_fail.clone(),
});
}
Ok(CheckedGraph {
nodes,
exclusive_to,
})
}
pub fn graph_fingerprint(nodes: &[CheckedNode]) -> String {
use sha2::{Digest, Sha256};
let mut hasher = Sha256::new();
for node in nodes {
hash_length_prefixed(&mut hasher, node.name.as_bytes());
hash_length_prefixed(&mut hasher, node.signature.to_string().as_bytes());
hash_length_prefixed(
&mut hasher,
node.over
.as_ref()
.map(|p| p.as_os_str().as_encoded_bytes())
.unwrap_or(b""),
);
hash_length_prefixed(&mut hasher, policy_repr(node).as_bytes());
}
crate::hex::encode(hasher.finalize())
}
fn policy_repr(node: &CheckedNode) -> String {
use cuttlefish_core::graph::{AcceptCheck, Rung};
let mut out = String::new();
for check in &node.accept {
match check {
AcceptCheck::Schema(path) => {
out.push_str("schema:");
out.push_str(&path.to_string_lossy());
}
AcceptCheck::Judge { model, prompt } => {
out.push_str("judge:");
if let Some(m) = model {
out.push_str(&m.to_string());
}
out.push(':');
out.push_str(prompt);
}
}
out.push('\n');
}
for rung in &node.on_fail {
match rung {
Rung::Retry(n) => out.push_str(&format!("retry:{n}")),
Rung::Reroute(m) => out.push_str(&format!("reroute:{m}")),
Rung::Escalate => out.push_str("escalate"),
}
out.push('\n');
}
out
}
fn hash_length_prefixed(hasher: &mut sha2::Sha256, bytes: &[u8]) {
use sha2::Digest;
hasher.update((bytes.len() as u64).to_be_bytes());
hasher.update(bytes);
}
fn topological_order(graph: &NodeGraph) -> Result<Vec<String>, DagError> {
let mut deps: HashMap<&str, HashSet<&str>> = HashMap::new();
for (name, node) in &graph.nodes {
deps.entry(name).or_default();
if let Some(expr) = &node.input {
for referenced in referenced_nodes(expr) {
if graph.get(referenced).is_none() {
return Err(DagError::UnknownReference {
node: name.clone(),
referenced: referenced.to_string(),
});
}
let is_marked_self_loop = referenced == name && node.repeat_until.is_some();
if !is_marked_self_loop {
deps.entry(name).or_default().insert(referenced);
}
}
}
}
let mut order = Vec::new();
let mut remaining: HashMap<&str, HashSet<&str>> = deps.clone();
loop {
let ready: Vec<&str> = remaining
.iter()
.filter(|(_, d)| d.is_empty())
.map(|(n, _)| *n)
.collect();
if ready.is_empty() {
break;
}
let mut ready = ready;
ready.sort(); for n in &ready {
order.push(n.to_string());
remaining.remove(n);
}
for deps in remaining.values_mut() {
for n in &ready {
deps.remove(n);
}
}
}
if order.len() != graph.nodes.len() {
let stuck = graph
.nodes
.iter()
.map(|(n, _)| n.as_str())
.find(|n| !order.contains(&n.to_string()))
.unwrap();
return Err(DagError::UnmarkedCycle {
node: stuck.to_string(),
});
}
Ok(order)
}
fn referenced_nodes(expr: &InputExpr) -> Vec<&str> {
match expr {
InputExpr::FromNode(n) => vec![n.as_str()],
InputExpr::Record(fields) => fields.values().flat_map(referenced_nodes).collect(),
InputExpr::List(items) => items.iter().flat_map(referenced_nodes).collect(),
}
}
#[allow(clippy::only_used_in_recursion)]
fn evaluate_expr_ty(
expr: &InputExpr,
signatures: &HashMap<String, cuttlefish_abi::Signature>,
graph: &NodeGraph,
) -> Result<Ty, DagError> {
match expr {
InputExpr::FromNode(n) => Ok(signatures
.get(n)
.ok_or_else(|| DagError::UnknownReference {
node: "?".into(),
referenced: n.clone(),
})?
.output
.clone()),
InputExpr::Record(fields) => {
let mut out = BTreeMap::new();
for (k, v) in fields {
out.insert(k.clone(), evaluate_expr_ty(v, signatures, graph)?);
}
Ok(Ty::Record(out))
}
InputExpr::List(items) => {
let first = items.first().ok_or_else(|| DagError::UnknownReference {
node: "?".into(),
referenced: "<empty list>".into(),
})?;
Ok(Ty::List(Box::new(evaluate_expr_ty(
first, signatures, graph,
)?)))
}
}
}
fn describe_expr(
expr: &InputExpr,
signatures: &HashMap<String, cuttlefish_abi::Signature>,
) -> (String, String) {
match expr {
InputExpr::FromNode(n) => (
n.clone(),
signatures
.get(n)
.map(|s| s.output.to_string())
.unwrap_or_default(),
),
InputExpr::Record(fields) => (
"<composite>".to_string(),
match same_node_repeated_across_every_field(fields) {
Some(node) => format!(
"{expr:?} -- every field here maps to `{node}.out`; if you meant to pass \
`{node}`'s whole output through unchanged, write `in = {node}.out;` with no \
braces instead of wrapping it field by field"
),
None => format!("{expr:?}"),
},
),
other => ("<composite>".to_string(), format!("{other:?}")),
}
}
fn same_node_repeated_across_every_field(fields: &BTreeMap<String, InputExpr>) -> Option<&str> {
let mut names = fields.values().map(|v| match v {
InputExpr::FromNode(n) => Some(n.as_str()),
_ => None,
});
let first = names.next()??;
names.all(|n| n == Some(first)).then_some(first)
}
fn compute_branch_exclusivity(
graph: &NodeGraph,
branches: &Branches,
) -> Result<HashMap<String, BranchExclusivity>, DagError> {
let mut exclusive_to: HashMap<String, BranchExclusivity> = HashMap::new();
for (decision, labels) in &branches.decisions {
for (label, target) in labels {
if graph.get(target).is_none() {
return Err(DagError::UnknownReference {
node: decision.clone(),
referenced: target.clone(),
});
}
exclusive_to.insert(
target.clone(),
BranchExclusivity {
decision: decision.clone(),
label: label.clone(),
},
);
}
}
loop {
let mut changed = false;
for (name, node) in &graph.nodes {
let Some(expr) = &node.input else { continue };
let mut found: Option<BranchExclusivity> = None;
for referenced in referenced_nodes(expr) {
if let Some(ex) = exclusive_to.get(referenced) {
match &found {
None => found = Some(ex.clone()),
Some(existing)
if existing.decision == ex.decision && existing.label != ex.label =>
{
return Err(DagError::ConflictingBranchFanIn {
node: name.clone(),
decision: ex.decision.clone(),
label_a: existing.label.clone(),
label_b: ex.label.clone(),
});
}
_ => {}
}
}
}
if let Some(ex) = found {
if exclusive_to.insert(name.clone(), ex).is_none() {
changed = true;
}
}
}
if !changed {
break;
}
}
Ok(exclusive_to)
}