Skip to main content

noether_engine/lagrange/
mod.rs

1mod ast;
2
3pub use ast::{collect_stage_ids, CompositionGraph, CompositionNode};
4
5use noether_core::stage::StageId;
6use noether_store::StageStore;
7use sha2::{Digest, Sha256};
8
9/// Parse a Lagrange JSON string into a CompositionGraph.
10pub fn parse_graph(json: &str) -> Result<CompositionGraph, serde_json::Error> {
11    serde_json::from_str(json)
12}
13
14/// Errors raised by `resolve_stage_prefixes` when an ID in the graph cannot
15/// be uniquely resolved against the store.
16#[derive(Debug, Clone)]
17pub enum PrefixResolutionError {
18    /// The prefix did not match any stage in the store.
19    NotFound { prefix: String },
20    /// The prefix matched multiple stages — author must use a longer prefix.
21    Ambiguous {
22        prefix: String,
23        matches: Vec<String>,
24    },
25}
26
27impl std::fmt::Display for PrefixResolutionError {
28    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
29        match self {
30            Self::NotFound { prefix } => {
31                write!(f, "no stage in store matches prefix '{prefix}'")
32            }
33            Self::Ambiguous { prefix, matches } => {
34                write!(
35                    f,
36                    "stage prefix '{prefix}' is ambiguous; matches {} stages — \
37                     use a longer prefix. First few: {}",
38                    matches.len(),
39                    matches
40                        .iter()
41                        .take(3)
42                        .map(|s| &s[..16.min(s.len())])
43                        .collect::<Vec<_>>()
44                        .join(", ")
45                )
46            }
47        }
48    }
49}
50
51impl std::error::Error for PrefixResolutionError {}
52
53/// Walk a composition graph and replace any stage IDs that are unique
54/// prefixes of a real stage in the store with their full 64-character IDs.
55///
56/// Exact matches are passed through unchanged. Hand-authored graphs can
57/// therefore use 8-character prefixes (the same form `noether stage list`
58/// prints) without manually looking up the full hash.
59pub fn resolve_stage_prefixes(
60    node: &mut CompositionNode,
61    store: &(impl StageStore + ?Sized),
62) -> Result<(), PrefixResolutionError> {
63    // Snapshot the IDs once — repeated walks would otherwise pay for it per node.
64    let ids: Vec<String> = store.list(None).iter().map(|s| s.id.0.clone()).collect();
65    resolve_in_node(node, &ids)
66}
67
68fn resolve_in_node(
69    node: &mut CompositionNode,
70    all_ids: &[String],
71) -> Result<(), PrefixResolutionError> {
72    match node {
73        CompositionNode::Stage { id, .. } => {
74            // Exact match: nothing to do.
75            if all_ids.iter().any(|i| i == &id.0) {
76                return Ok(());
77            }
78            // Otherwise, look for prefix matches.
79            let matches: Vec<&String> = all_ids.iter().filter(|i| i.starts_with(&id.0)).collect();
80            match matches.len() {
81                0 => Err(PrefixResolutionError::NotFound {
82                    prefix: id.0.clone(),
83                }),
84                1 => {
85                    *id = StageId(matches[0].clone());
86                    Ok(())
87                }
88                _ => Err(PrefixResolutionError::Ambiguous {
89                    prefix: id.0.clone(),
90                    matches: matches.into_iter().cloned().collect(),
91                }),
92            }
93        }
94        CompositionNode::RemoteStage { .. } | CompositionNode::Const { .. } => Ok(()),
95        CompositionNode::Sequential { stages } => {
96            for s in stages {
97                resolve_in_node(s, all_ids)?;
98            }
99            Ok(())
100        }
101        CompositionNode::Parallel { branches } => {
102            for b in branches.values_mut() {
103                resolve_in_node(b, all_ids)?;
104            }
105            Ok(())
106        }
107        CompositionNode::Branch {
108            predicate,
109            if_true,
110            if_false,
111        } => {
112            resolve_in_node(predicate, all_ids)?;
113            resolve_in_node(if_true, all_ids)?;
114            resolve_in_node(if_false, all_ids)
115        }
116        CompositionNode::Fanout { source, targets } => {
117            resolve_in_node(source, all_ids)?;
118            for t in targets {
119                resolve_in_node(t, all_ids)?;
120            }
121            Ok(())
122        }
123        CompositionNode::Merge { sources, target } => {
124            for s in sources {
125                resolve_in_node(s, all_ids)?;
126            }
127            resolve_in_node(target, all_ids)
128        }
129        CompositionNode::Retry { stage, .. } => resolve_in_node(stage, all_ids),
130        CompositionNode::Let { bindings, body } => {
131            for b in bindings.values_mut() {
132                resolve_in_node(b, all_ids)?;
133            }
134            resolve_in_node(body, all_ids)
135        }
136    }
137}
138
139/// Serialize a CompositionGraph to pretty-printed JSON.
140pub fn serialize_graph(graph: &CompositionGraph) -> Result<String, serde_json::Error> {
141    serde_json::to_string_pretty(graph)
142}
143
144/// Compute a deterministic composition ID (SHA-256 of canonical JSON).
145pub fn compute_composition_id(graph: &CompositionGraph) -> Result<String, serde_json::Error> {
146    let bytes = serde_json::to_vec(graph)?;
147    let hash = Sha256::digest(&bytes);
148    Ok(hex::encode(hash))
149}
150
151#[cfg(test)]
152mod tests {
153    use super::*;
154    use crate::lagrange::ast::CompositionNode;
155    use noether_core::stage::StageId;
156
157    #[test]
158    fn parse_and_serialize_round_trip() {
159        let graph = CompositionGraph::new(
160            "test",
161            CompositionNode::Stage {
162                id: StageId("abc".into()),
163                config: None,
164            },
165        );
166        let json = serialize_graph(&graph).unwrap();
167        let parsed = parse_graph(&json).unwrap();
168        assert_eq!(graph, parsed);
169    }
170
171    #[test]
172    fn composition_id_is_deterministic() {
173        let graph = CompositionGraph::new(
174            "test",
175            CompositionNode::Stage {
176                id: StageId("abc".into()),
177                config: None,
178            },
179        );
180        let id1 = compute_composition_id(&graph).unwrap();
181        let id2 = compute_composition_id(&graph).unwrap();
182        assert_eq!(id1, id2);
183        assert_eq!(id1.len(), 64);
184    }
185}