1use noether_core::stage::{StageId, StageLifecycle};
18use noether_store::StageStore;
19
20use super::ast::CompositionNode;
21
22pub const MAX_DEPRECATION_HOPS: usize = 10;
27
28#[derive(Debug, Clone, PartialEq, Eq)]
30pub struct DeprecationRewrite {
31 pub from: StageId,
32 pub to: StageId,
33}
34
35#[derive(Debug, Clone, PartialEq, Eq)]
38pub enum ChainEvent {
39 CycleDetected { stage: StageId },
43 MaxHopsExceeded { stage: StageId },
46}
47
48#[derive(Debug, Default, Clone)]
50pub struct DeprecationReport {
51 pub rewrites: Vec<DeprecationRewrite>,
54 pub events: Vec<ChainEvent>,
58}
59
60pub fn resolve_deprecated_stages(
69 node: &mut CompositionNode,
70 store: &dyn StageStore,
71) -> DeprecationReport {
72 let mut report = DeprecationReport::default();
73 walk(node, store, &mut report);
74 report
75}
76
77fn walk(node: &mut CompositionNode, store: &dyn StageStore, report: &mut DeprecationReport) {
78 match node {
79 CompositionNode::Stage { id, .. } => follow_chain(id, store, report),
80 CompositionNode::Sequential { stages } => {
81 for s in stages {
82 walk(s, store, report);
83 }
84 }
85 CompositionNode::Parallel { branches } => {
86 for (_, branch) in branches.iter_mut() {
87 walk(branch, store, report);
88 }
89 }
90 CompositionNode::Branch {
91 predicate,
92 if_true,
93 if_false,
94 } => {
95 walk(predicate, store, report);
96 walk(if_true, store, report);
97 walk(if_false, store, report);
98 }
99 CompositionNode::Retry { stage, .. } => walk(stage, store, report),
100 CompositionNode::Fanout { source, targets } => {
101 walk(source, store, report);
102 for t in targets {
103 walk(t, store, report);
104 }
105 }
106 CompositionNode::Merge { sources, target } => {
107 for s in sources {
108 walk(s, store, report);
109 }
110 walk(target, store, report);
111 }
112 CompositionNode::Const { .. } | CompositionNode::RemoteStage { .. } => {}
113 CompositionNode::Let { bindings, body } => {
114 for b in bindings.values_mut() {
115 walk(b, store, report);
116 }
117 walk(body, store, report);
118 }
119 }
120}
121
122fn follow_chain(id: &mut StageId, store: &dyn StageStore, report: &mut DeprecationReport) {
123 let mut visited: std::collections::HashSet<StageId> = std::collections::HashSet::new();
124 visited.insert(id.clone());
125 let mut current = id.clone();
126 let mut hops = 0usize;
127
128 while let Ok(Some(stage)) = store.get(¤t) {
129 let successor = match &stage.lifecycle {
130 StageLifecycle::Deprecated { successor_id } => successor_id.clone(),
131 _ => break,
132 };
133
134 if !visited.insert(successor.clone()) {
135 report.events.push(ChainEvent::CycleDetected {
139 stage: successor.clone(),
140 });
141 break;
142 }
143
144 hops += 1;
145 if hops > MAX_DEPRECATION_HOPS {
146 report.events.push(ChainEvent::MaxHopsExceeded {
147 stage: successor.clone(),
148 });
149 break;
150 }
151
152 report.rewrites.push(DeprecationRewrite {
153 from: current.clone(),
154 to: successor.clone(),
155 });
156 current = successor;
157 }
158
159 if current != *id {
160 *id = current;
161 }
162}
163
164#[cfg(test)]
165mod tests {
166 use super::*;
167 use crate::lagrange::ast::Pinning;
168 use noether_core::effects::EffectSet;
169 use noether_core::stage::{
170 compute_stage_id, CostEstimate, Stage, StageLifecycle, StageSignature,
171 };
172 use noether_core::types::NType;
173 use noether_store::MemoryStore;
174 use std::collections::BTreeSet;
175
176 fn stage(name: &str, impl_hash: &str, lifecycle: StageLifecycle) -> Stage {
177 let signature = StageSignature {
178 input: NType::Text,
179 output: NType::Text,
180 effects: EffectSet::pure(),
181 implementation_hash: impl_hash.into(),
182 };
183 let id = compute_stage_id(name, &signature).unwrap();
184 Stage {
185 id,
186 signature_id: None,
187 signature,
188 capabilities: BTreeSet::new(),
189 cost: CostEstimate {
190 time_ms_p50: None,
191 tokens_est: None,
192 memory_mb: None,
193 },
194 description: "t".into(),
195 examples: vec![],
196 lifecycle,
197 ed25519_signature: None,
198 signer_public_key: None,
199 implementation_code: None,
200 implementation_language: None,
201 ui_style: None,
202 tags: vec![],
203 aliases: vec![],
204 name: Some(name.into()),
205 properties: Vec::new(),
206 }
207 }
208
209 fn leaf(id: &StageId) -> CompositionNode {
210 CompositionNode::Stage {
211 id: id.clone(),
212 pinning: Pinning::Both,
213 config: None,
214 }
215 }
216
217 #[test]
218 fn noop_on_active_stage() {
219 let mut store = MemoryStore::new();
220 let active = stage("a", "ha", StageLifecycle::Active);
221 let id = active.id.clone();
222 store.put(active).unwrap();
223
224 let mut root = leaf(&id);
225 let report = resolve_deprecated_stages(&mut root, &store);
226 assert!(report.rewrites.is_empty());
227 assert!(report.events.is_empty());
228 }
229
230 #[test]
231 fn single_hop_rewrites() {
232 let mut store = MemoryStore::new();
233 let new_stage = stage("new", "hn", StageLifecycle::Active);
234 let new_id = new_stage.id.clone();
235 store.put(new_stage).unwrap();
236
237 let old_active = stage("old", "ho", StageLifecycle::Active);
240 let old_id = old_active.id.clone();
241 store.put(old_active).unwrap();
242 store
243 .update_lifecycle(
244 &old_id,
245 StageLifecycle::Deprecated {
246 successor_id: new_id.clone(),
247 },
248 )
249 .unwrap();
250
251 let mut root = leaf(&old_id);
252 let report = resolve_deprecated_stages(&mut root, &store);
253 assert_eq!(
254 report.rewrites,
255 vec![DeprecationRewrite {
256 from: old_id,
257 to: new_id.clone(),
258 }]
259 );
260 match root {
261 CompositionNode::Stage { id, .. } => assert_eq!(id, new_id),
262 _ => unreachable!(),
263 }
264 assert!(report.events.is_empty());
265 }
266
267 #[test]
268 fn cycle_detected_and_surfaced() {
269 use noether_store::{StageStore, StoreError, StoreStats};
275 use std::collections::HashMap;
276
277 struct CyclicStore {
278 stages: HashMap<String, Stage>,
279 }
280
281 impl StageStore for CyclicStore {
282 fn put(&mut self, _s: Stage) -> Result<StageId, StoreError> {
283 unimplemented!()
284 }
285 fn upsert(&mut self, _s: Stage) -> Result<StageId, StoreError> {
286 unimplemented!()
287 }
288 fn remove(&mut self, _id: &StageId) -> Result<(), StoreError> {
289 unimplemented!()
290 }
291 fn get(&self, id: &StageId) -> Result<Option<&Stage>, StoreError> {
292 Ok(self.stages.get(&id.0))
293 }
294 fn contains(&self, id: &StageId) -> bool {
295 self.stages.contains_key(&id.0)
296 }
297 fn list(&self, _lc: Option<&StageLifecycle>) -> Vec<&Stage> {
298 self.stages.values().collect()
299 }
300 fn update_lifecycle(
301 &mut self,
302 _id: &StageId,
303 _lc: StageLifecycle,
304 ) -> Result<(), StoreError> {
305 unimplemented!()
306 }
307 fn stats(&self) -> StoreStats {
308 StoreStats {
309 total: self.stages.len(),
310 by_lifecycle: Default::default(),
311 by_effect: Default::default(),
312 }
313 }
314 }
315
316 let a = stage("a", "ha", StageLifecycle::Active);
317 let b = stage("b", "hb", StageLifecycle::Active);
318 let a_id = a.id.clone();
319 let b_id = b.id.clone();
320
321 let a_dep = Stage {
322 lifecycle: StageLifecycle::Deprecated {
323 successor_id: b_id.clone(),
324 },
325 ..a
326 };
327 let b_dep = Stage {
328 lifecycle: StageLifecycle::Deprecated {
329 successor_id: a_id.clone(),
330 },
331 ..b
332 };
333 let mut stages = HashMap::new();
334 stages.insert(a_id.0.clone(), a_dep);
335 stages.insert(b_id.0.clone(), b_dep);
336 let store = CyclicStore { stages };
337
338 let mut root = leaf(&a_id);
339 let report = resolve_deprecated_stages(&mut root, &store);
340
341 assert!(
343 report
344 .events
345 .iter()
346 .any(|e| matches!(e, ChainEvent::CycleDetected { .. })),
347 "expected CycleDetected event, got {:?}",
348 report.events
349 );
350 }
351}