1use crate::core::types::{RiskTier, ValidationTier};
9use serde::{Deserialize, Serialize};
10use std::collections::{HashMap, HashSet};
11
12#[derive(Debug, Clone)]
21pub struct PlanItemInput {
22 pub component_id: String,
24 pub current_version: String,
26 pub proposed_version: String,
28 pub validation_tier: ValidationTier,
30 pub selected: bool,
32 pub rationale: String,
34 pub dependencies: Vec<String>,
36 pub isolation_safe: bool,
38}
39
40#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
42pub struct PlanItem {
43 pub component_id: String,
45 pub current_version: String,
47 pub proposed_version: String,
49 pub validation_tier: ValidationTier,
51 pub selected: bool,
53 pub rationale: String,
55 #[serde(default)]
57 pub dependencies: Vec<String>,
58 pub isolation_safe: bool,
60 pub risk_tier: RiskTier,
62 #[serde(default)]
66 pub exclusive_group: String,
67}
68
69impl PlanItem {
70 pub fn new(input: PlanItemInput) -> Self {
72 Self {
73 component_id: input.component_id,
74 current_version: input.current_version,
75 proposed_version: input.proposed_version,
76 validation_tier: input.validation_tier,
77 selected: input.selected,
78 rationale: input.rationale,
79 dependencies: input.dependencies,
80 isolation_safe: input.isolation_safe,
81 risk_tier: RiskTier::Medium, exclusive_group: String::new(), }
84 }
85}
86
87pub fn topological_sort(items: Vec<PlanItem>) -> Vec<PlanItem> {
105 let mut remaining_deps: HashMap<String, Vec<String>> = HashMap::new();
107 let mut dependents: HashMap<String, Vec<String>> = HashMap::new();
109 let mut all_nodes: HashSet<String> = HashSet::new();
110
111 for item in &items {
113 all_nodes.insert(item.component_id.clone());
114 remaining_deps.insert(item.component_id.clone(), item.dependencies.clone());
116 for dep in &item.dependencies {
118 dependents
119 .entry(dep.clone())
120 .or_default()
121 .push(item.component_id.clone());
122 }
123 }
124
125 let mut sorted_nodes = Vec::new();
127
128 let mut queue: Vec<String> = all_nodes
130 .iter()
131 .filter(|id| remaining_deps.get(*id).is_none_or(|d| d.is_empty()))
132 .cloned()
133 .collect();
134
135 while let Some(node) = queue.pop() {
136 sorted_nodes.push(node.clone());
137
138 if let Some(dep_on_node) = dependents.get(&node) {
140 for dependent in dep_on_node.clone() {
141 if let Some(deps) = remaining_deps.get_mut(&dependent) {
143 deps.retain(|d| d != &node);
144 if deps.is_empty() {
146 queue.push(dependent);
147 }
148 }
149 }
150 }
151 }
152
153 if sorted_nodes.len() != all_nodes.len() {
154 panic!("Dependency cycle detected in plan items");
155 }
156
157 let mut item_map: HashMap<String, PlanItem> = items
159 .into_iter()
160 .map(|item| (item.component_id.clone(), item))
161 .collect();
162
163 sorted_nodes
164 .into_iter()
165 .filter_map(|id| item_map.remove(&id))
166 .collect()
167}
168
169#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
175pub struct Plan {
176 pub items: Vec<PlanItem>,
178 pub summary: PlanSummary,
180}
181
182#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
184pub struct PlanSummary {
185 pub total: usize,
186 pub selected: usize,
187 pub safe: usize,
188 pub guarded: usize,
189 pub candidate: usize,
190 pub experimental: usize,
191 pub blocked: usize,
192}
193
194impl Plan {
195 pub fn new(items: Vec<PlanItem>) -> Self {
197 let items = topological_sort(items);
198 let summary = PlanSummary::from_items(&items);
199 Self { items, summary }
200 }
201}
202
203impl PlanSummary {
204 pub fn from_items(items: &[PlanItem]) -> Self {
206 Self {
207 total: items.len(),
208 selected: items.iter().filter(|i| i.selected).count(),
209 safe: items
210 .iter()
211 .filter(|i| i.risk_tier == RiskTier::Low)
212 .count(),
213 guarded: items
214 .iter()
215 .filter(|i| i.risk_tier == RiskTier::Medium)
216 .count(),
217 candidate: items
218 .iter()
219 .filter(|i| i.risk_tier == RiskTier::High)
220 .count(),
221 experimental: items
222 .iter()
223 .filter(|i| i.validation_tier == ValidationTier::Experimental)
224 .count(),
225 blocked: 0, }
227 }
228}
229
230#[cfg(test)]
235mod tests {
236 use super::*;
237 use crate::core::types::{RiskTier, ValidationTier};
238
239 #[test]
244 fn test_topological_sort_simple_dependency() {
245 let items = vec![
246 PlanItem::new(PlanItemInput {
247 component_id: "pytorch".into(),
248 current_version: "2.4.0".into(),
249 proposed_version: "2.5.0".into(),
250 validation_tier: ValidationTier::Validated,
251 selected: true,
252 rationale: "update pytorch".into(),
253 dependencies: vec!["rocm".to_string()],
254 isolation_safe: true,
255 }),
256 PlanItem::new(PlanItemInput {
257 component_id: "rocm".into(),
258 current_version: "7.2.0".into(),
259 proposed_version: "7.3.0".into(),
260 validation_tier: ValidationTier::Validated,
261 selected: true,
262 rationale: "update rocm".into(),
263 dependencies: vec![],
264 isolation_safe: true,
265 }),
266 ];
267
268 let sorted = topological_sort(items);
269 assert_eq!(sorted[0].component_id, "rocm");
270 assert_eq!(sorted[1].component_id, "pytorch");
271 }
272
273 #[test]
274 fn test_topological_sort_complex_dependency_chain() {
275 let items = vec![
276 PlanItem::new(PlanItemInput {
277 component_id: "pytorch".into(),
278 current_version: "2.4.0".into(),
279 proposed_version: "2.5.0".into(),
280 validation_tier: ValidationTier::Validated,
281 selected: true,
282 rationale: "update pytorch".into(),
283 dependencies: vec!["rocm".to_string()],
284 isolation_safe: true,
285 }),
286 PlanItem::new(PlanItemInput {
287 component_id: "triton".into(),
288 current_version: "3.0.0".into(),
289 proposed_version: "3.1.0".into(),
290 validation_tier: ValidationTier::Validated,
291 selected: true,
292 rationale: "update triton".into(),
293 dependencies: vec!["pytorch".to_string()],
294 isolation_safe: true,
295 }),
296 PlanItem::new(PlanItemInput {
297 component_id: "rocm".into(),
298 current_version: "7.2.0".into(),
299 proposed_version: "7.3.0".into(),
300 validation_tier: ValidationTier::Validated,
301 selected: true,
302 rationale: "update rocm".into(),
303 dependencies: vec![],
304 isolation_safe: true,
305 }),
306 ];
307
308 let sorted = topological_sort(items);
309 assert_eq!(sorted[0].component_id, "rocm");
310 assert_eq!(sorted[1].component_id, "pytorch");
311 assert_eq!(sorted[2].component_id, "triton");
312 }
313
314 #[test]
315 fn test_topological_sort_empty_dependencies() {
316 let items = vec![
317 PlanItem::new(PlanItemInput {
318 component_id: "pytorch".into(),
319 current_version: "2.4.0".into(),
320 proposed_version: "2.5.0".into(),
321 validation_tier: ValidationTier::Validated,
322 selected: true,
323 rationale: "update pytorch".into(),
324 dependencies: vec![],
325 isolation_safe: true,
326 }),
327 PlanItem::new(PlanItemInput {
328 component_id: "rocm".into(),
329 current_version: "7.2.0".into(),
330 proposed_version: "7.3.0".into(),
331 validation_tier: ValidationTier::Validated,
332 selected: true,
333 rationale: "update rocm".into(),
334 dependencies: vec![],
335 isolation_safe: true,
336 }),
337 ];
338
339 let sorted = topological_sort(items);
340 let ids: HashSet<_> = sorted.iter().map(|i| i.component_id.clone()).collect();
342 assert!(ids.contains("pytorch"));
343 assert!(ids.contains("rocm"));
344 }
345
346 #[test]
351 #[should_panic(expected = "Dependency cycle detected")]
352 fn test_topological_sort_detects_cycle() {
353 let items = vec![
354 PlanItem::new(PlanItemInput {
355 component_id: "pytorch".into(),
356 current_version: "2.4.0".into(),
357 proposed_version: "2.5.0".into(),
358 validation_tier: ValidationTier::Validated,
359 selected: true,
360 rationale: "update pytorch".into(),
361 dependencies: vec!["triton".to_string()],
362 isolation_safe: true,
363 }),
364 PlanItem::new(PlanItemInput {
365 component_id: "triton".into(),
366 current_version: "3.0.0".into(),
367 proposed_version: "3.1.0".into(),
368 validation_tier: ValidationTier::Validated,
369 selected: true,
370 rationale: "update triton".into(),
371 dependencies: vec!["pytorch".to_string()],
372 isolation_safe: true,
373 }),
374 ];
375
376 topological_sort(items);
377 }
378
379 #[test]
384 fn test_plan_creation_applies_topological_sort() {
385 let items = vec![
386 PlanItem::new(PlanItemInput {
387 component_id: "pytorch".into(),
388 current_version: "2.4.0".into(),
389 proposed_version: "2.5.0".into(),
390 validation_tier: ValidationTier::Validated,
391 selected: true,
392 rationale: "update pytorch".into(),
393 dependencies: vec!["rocm".to_string()],
394 isolation_safe: true,
395 }),
396 PlanItem::new(PlanItemInput {
397 component_id: "rocm".into(),
398 current_version: "7.2.0".into(),
399 proposed_version: "7.3.0".into(),
400 validation_tier: ValidationTier::Validated,
401 selected: true,
402 rationale: "update rocm".into(),
403 dependencies: vec![],
404 isolation_safe: true,
405 }),
406 ];
407
408 let plan = Plan::new(items);
409 assert_eq!(plan.items[0].component_id, "rocm");
410 assert_eq!(plan.items[1].component_id, "pytorch");
411 }
412
413 #[test]
418 fn test_plan_summary_counts() {
419 let items = vec![
420 PlanItem {
421 component_id: "rocm".to_string(),
422 current_version: "7.2.0".to_string(),
423 proposed_version: "7.3.0".to_string(),
424 validation_tier: ValidationTier::Validated,
425 selected: true,
426 rationale: "update rocm".to_string(),
427 dependencies: vec![],
428 isolation_safe: true,
429 risk_tier: RiskTier::Low,
430 exclusive_group: String::new(),
431 },
432 PlanItem {
433 component_id: "pytorch".to_string(),
434 current_version: "2.4.0".to_string(),
435 proposed_version: "2.5.0".to_string(),
436 validation_tier: ValidationTier::Validated,
437 selected: true,
438 rationale: "update pytorch".to_string(),
439 dependencies: vec!["rocm".to_string()],
440 isolation_safe: true,
441 risk_tier: RiskTier::Medium,
442 exclusive_group: String::new(),
443 },
444 PlanItem {
445 component_id: "deepspeed".to_string(),
446 current_version: "".to_string(),
447 proposed_version: "0.15.0".to_string(),
448 validation_tier: ValidationTier::Experimental,
449 selected: false,
450 rationale: "new experimental".to_string(),
451 dependencies: vec!["pytorch".to_string()],
452 isolation_safe: true,
453 risk_tier: RiskTier::High,
454 exclusive_group: String::new(),
455 },
456 ];
457
458 let plan = Plan::new(items);
459 assert_eq!(plan.summary.total, 3);
460 assert_eq!(plan.summary.selected, 2);
461 assert_eq!(plan.summary.safe, 1);
462 assert_eq!(plan.summary.guarded, 1);
463 assert_eq!(plan.summary.candidate, 1);
464 assert_eq!(plan.summary.experimental, 1);
465 assert_eq!(plan.summary.blocked, 0);
466 }
467
468 #[test]
473 fn test_plan_item_serde_roundtrip() {
474 let item = PlanItem {
475 component_id: "pytorch".to_string(),
476 current_version: "2.4.0".to_string(),
477 proposed_version: "2.5.0".to_string(),
478 validation_tier: ValidationTier::Validated,
479 selected: true,
480 rationale: "update pytorch".to_string(),
481 dependencies: vec!["rocm".to_string()],
482 isolation_safe: true,
483 risk_tier: RiskTier::Medium,
484 exclusive_group: String::new(),
485 };
486
487 let json = serde_json::to_string(&item).unwrap();
488 let back: PlanItem = serde_json::from_str(&json).unwrap();
489 assert_eq!(item, back);
490 }
491
492 #[test]
497 fn test_plan_serde_roundtrip() {
498 let items = vec![
499 PlanItem::new(PlanItemInput {
500 component_id: "rocm".into(),
501 current_version: "7.2.0".into(),
502 proposed_version: "7.3.0".into(),
503 validation_tier: ValidationTier::Validated,
504 selected: true,
505 rationale: "update rocm".into(),
506 dependencies: vec![],
507 isolation_safe: true,
508 }),
509 PlanItem::new(PlanItemInput {
510 component_id: "pytorch".into(),
511 current_version: "2.4.0".into(),
512 proposed_version: "2.5.0".into(),
513 validation_tier: ValidationTier::Validated,
514 selected: true,
515 rationale: "update pytorch".into(),
516 dependencies: vec!["rocm".to_string()],
517 isolation_safe: true,
518 }),
519 ];
520
521 let plan = Plan::new(items);
522 let json = serde_json::to_string(&plan).unwrap();
523 let back: Plan = serde_json::from_str(&json).unwrap();
524 assert_eq!(plan.items, back.items);
525 assert_eq!(plan.summary, back.summary);
526 }
527
528 #[test]
533 fn test_topological_sort_with_llama_cpp_dependencies() {
534 let items = vec![
535 PlanItem::new(PlanItemInput {
536 component_id: "llama-cpp".into(),
537 current_version: "".into(),
538 proposed_version: "latest".into(),
539 validation_tier: ValidationTier::Experimental,
540 selected: true,
541 rationale: "install llama-cpp".into(),
542 dependencies: vec!["rocm".to_string()],
543 isolation_safe: true,
544 }),
545 PlanItem::new(PlanItemInput {
546 component_id: "rocm".into(),
547 current_version: "7.2.0".into(),
548 proposed_version: "7.3.0".into(),
549 validation_tier: ValidationTier::Validated,
550 selected: true,
551 rationale: "update rocm".into(),
552 dependencies: vec![],
553 isolation_safe: true,
554 }),
555 ];
556
557 let sorted = topological_sort(items);
558 assert_eq!(sorted[0].component_id, "rocm");
559 assert_eq!(sorted[1].component_id, "llama-cpp");
560 }
561}