1use std::collections::{HashMap, HashSet, VecDeque};
5
6#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
12pub enum DependencyKind {
13 TensorInput,
15 TensorOutput,
17 RuleImplication,
19 SharedFact,
21}
22
23#[derive(Clone, Debug, PartialEq, Eq)]
29pub struct DependencyEdge {
30 pub from_id: u64,
32 pub to_id: u64,
34 pub kind: DependencyKind,
36 pub weight: u32,
38}
39
40#[derive(Clone, Debug, Default)]
46pub struct DirtySet {
47 pub dirty: HashSet<u64>,
48}
49
50impl DirtySet {
51 pub fn new() -> Self {
53 Self {
54 dirty: HashSet::new(),
55 }
56 }
57
58 pub fn mark_dirty(&mut self, id: u64) {
60 self.dirty.insert(id);
61 }
62
63 pub fn is_dirty(&self, id: u64) -> bool {
65 self.dirty.contains(&id)
66 }
67
68 pub fn clear_dirty(&mut self, id: u64) {
70 self.dirty.remove(&id);
71 }
72
73 pub fn all_dirty(&self) -> Vec<u64> {
75 let mut v: Vec<u64> = self.dirty.iter().copied().collect();
76 v.sort_unstable();
77 v
78 }
79}
80
81#[derive(Clone, Debug, PartialEq, Eq)]
87pub struct GraphStats {
88 pub node_count: usize,
89 pub edge_count: usize,
90 pub dirty_count: usize,
91 pub max_in_degree: usize,
92 pub max_out_degree: usize,
93}
94
95pub struct TensorDependencyGraph {
102 pub nodes: HashSet<u64>,
104 pub edges: Vec<DependencyEdge>,
106 pub dirty: DirtySet,
108}
109
110impl TensorDependencyGraph {
111 pub fn new() -> Self {
113 Self {
114 nodes: HashSet::new(),
115 edges: Vec::new(),
116 dirty: DirtySet::new(),
117 }
118 }
119
120 pub fn add_node(&mut self, id: u64) {
122 self.nodes.insert(id);
123 }
124
125 pub fn add_edge(&mut self, edge: DependencyEdge) {
127 self.nodes.insert(edge.from_id);
128 self.nodes.insert(edge.to_id);
129 self.edges.push(edge);
130 }
131
132 pub fn remove_node(&mut self, id: u64) {
134 self.nodes.remove(&id);
135 self.edges.retain(|e| e.from_id != id && e.to_id != id);
136 self.dirty.clear_dirty(id);
137 }
138
139 pub fn mark_dirty(&mut self, id: u64) {
143 let mut queue: VecDeque<u64> = VecDeque::new();
145 self.dirty.mark_dirty(id);
146 queue.push_back(id);
147
148 while let Some(current) = queue.pop_front() {
149 let successors: Vec<u64> = self
151 .edges
152 .iter()
153 .filter(|e| {
154 e.from_id == current
155 && matches!(
156 e.kind,
157 DependencyKind::TensorOutput
158 | DependencyKind::RuleImplication
159 | DependencyKind::SharedFact
160 )
161 })
162 .map(|e| e.to_id)
163 .collect();
164
165 for succ in successors {
166 if !self.dirty.is_dirty(succ) {
167 self.dirty.mark_dirty(succ);
168 queue.push_back(succ);
169 }
170 }
171 }
172 }
173
174 pub fn recompute_order(&self) -> Vec<u64> {
178 let dirty: HashSet<u64> = self.dirty.dirty.clone();
179 if dirty.is_empty() {
180 return Vec::new();
181 }
182
183 let mut in_degree: HashMap<u64, usize> = dirty.iter().map(|&n| (n, 0)).collect();
185 let mut adj: HashMap<u64, Vec<u64>> = dirty.iter().map(|&n| (n, Vec::new())).collect();
187
188 for edge in &self.edges {
189 if dirty.contains(&edge.from_id) && dirty.contains(&edge.to_id) {
190 adj.entry(edge.from_id).or_default().push(edge.to_id);
191 *in_degree.entry(edge.to_id).or_insert(0) += 1;
192 }
193 }
194
195 let mut queue: VecDeque<u64> = in_degree
196 .iter()
197 .filter(|(_, °)| deg == 0)
198 .map(|(&n, _)| n)
199 .collect();
200
201 let mut sorted_queue: Vec<u64> = queue.drain(..).collect();
203 sorted_queue.sort_unstable();
204 queue.extend(sorted_queue);
205
206 let mut result: Vec<u64> = Vec::with_capacity(dirty.len());
207
208 while let Some(node) = queue.pop_front() {
209 result.push(node);
210 if let Some(neighbors) = adj.get(&node) {
211 let mut next_batch: Vec<u64> = Vec::new();
212 for &neighbor in neighbors {
213 if let Some(deg) = in_degree.get_mut(&neighbor) {
214 *deg = deg.saturating_sub(1);
215 if *deg == 0 {
216 next_batch.push(neighbor);
217 }
218 }
219 }
220 next_batch.sort_unstable();
221 queue.extend(next_batch);
222 }
223 }
224
225 if result.len() < dirty.len() {
227 let processed: HashSet<u64> = result.iter().copied().collect();
228 let mut remaining: Vec<u64> = dirty
229 .iter()
230 .filter(|n| !processed.contains(n))
231 .copied()
232 .collect();
233 remaining.sort_unstable();
234 result.extend(remaining);
235 }
236
237 result
238 }
239
240 pub fn dependents_of(&self, id: u64) -> Vec<u64> {
242 let mut out: Vec<u64> = self
243 .edges
244 .iter()
245 .filter(|e| e.from_id == id)
246 .map(|e| e.to_id)
247 .collect();
248 out.sort_unstable();
249 out.dedup();
250 out
251 }
252
253 pub fn dependencies_of(&self, id: u64) -> Vec<u64> {
255 let mut out: Vec<u64> = self
256 .edges
257 .iter()
258 .filter(|e| e.to_id == id)
259 .map(|e| e.from_id)
260 .collect();
261 out.sort_unstable();
262 out.dedup();
263 out
264 }
265
266 pub fn stats(&self) -> GraphStats {
268 let node_count = self.nodes.len();
269 let edge_count = self.edges.len();
270 let dirty_count = self.dirty.dirty.len();
271
272 let mut in_degree: HashMap<u64, usize> = HashMap::new();
273 let mut out_degree: HashMap<u64, usize> = HashMap::new();
274
275 for node in &self.nodes {
276 in_degree.entry(*node).or_insert(0);
277 out_degree.entry(*node).or_insert(0);
278 }
279
280 for edge in &self.edges {
281 *out_degree.entry(edge.from_id).or_insert(0) += 1;
282 *in_degree.entry(edge.to_id).or_insert(0) += 1;
283 }
284
285 let max_in_degree = in_degree.values().copied().max().unwrap_or(0);
286 let max_out_degree = out_degree.values().copied().max().unwrap_or(0);
287
288 GraphStats {
289 node_count,
290 edge_count,
291 dirty_count,
292 max_in_degree,
293 max_out_degree,
294 }
295 }
296
297 pub fn clear_all_dirty(&mut self) {
299 self.dirty.dirty.clear();
300 }
301}
302
303impl Default for TensorDependencyGraph {
304 fn default() -> Self {
305 Self::new()
306 }
307}
308
309#[cfg(test)]
314mod tests {
315 use super::*;
316
317 fn edge(from: u64, to: u64, kind: DependencyKind) -> DependencyEdge {
318 DependencyEdge {
319 from_id: from,
320 to_id: to,
321 kind,
322 weight: 1,
323 }
324 }
325
326 fn edge_w(from: u64, to: u64, kind: DependencyKind, weight: u32) -> DependencyEdge {
327 DependencyEdge {
328 from_id: from,
329 to_id: to,
330 kind,
331 weight,
332 }
333 }
334
335 #[test]
337 fn test_empty_graph_stats() {
338 let g = TensorDependencyGraph::new();
339 let s = g.stats();
340 assert_eq!(s.node_count, 0);
341 assert_eq!(s.edge_count, 0);
342 assert_eq!(s.dirty_count, 0);
343 assert_eq!(s.max_in_degree, 0);
344 assert_eq!(s.max_out_degree, 0);
345 }
346
347 #[test]
349 fn test_add_node_idempotent() {
350 let mut g = TensorDependencyGraph::new();
351 g.add_node(1);
352 g.add_node(1);
353 g.add_node(1);
354 assert_eq!(g.stats().node_count, 1);
355 }
356
357 #[test]
359 fn test_add_edge_registers_nodes() {
360 let mut g = TensorDependencyGraph::new();
361 g.add_edge(edge(10, 20, DependencyKind::TensorInput));
362 assert!(g.nodes.contains(&10));
363 assert!(g.nodes.contains(&20));
364 assert_eq!(g.stats().node_count, 2);
365 }
366
367 #[test]
369 fn test_remove_node_cleans_edges() {
370 let mut g = TensorDependencyGraph::new();
371 g.add_edge(edge(1, 2, DependencyKind::TensorOutput));
372 g.add_edge(edge(2, 3, DependencyKind::RuleImplication));
373 g.remove_node(2);
374 assert_eq!(g.edges.len(), 0);
375 assert!(!g.nodes.contains(&2));
376 }
377
378 #[test]
380 fn test_remove_node_clears_dirty() {
381 let mut g = TensorDependencyGraph::new();
382 g.add_node(5);
383 g.dirty.mark_dirty(5);
384 assert!(g.dirty.is_dirty(5));
385 g.remove_node(5);
386 assert!(!g.dirty.is_dirty(5));
387 }
388
389 #[test]
391 fn test_mark_dirty_propagates_transitively() {
392 let mut g = TensorDependencyGraph::new();
393 g.add_edge(edge(1, 2, DependencyKind::TensorOutput));
395 g.add_edge(edge(2, 3, DependencyKind::RuleImplication));
396 g.add_edge(edge(3, 4, DependencyKind::SharedFact));
397 g.mark_dirty(1);
398 assert!(g.dirty.is_dirty(1));
399 assert!(g.dirty.is_dirty(2));
400 assert!(g.dirty.is_dirty(3));
401 assert!(g.dirty.is_dirty(4));
402 }
403
404 #[test]
406 fn test_mark_dirty_no_outgoing() {
407 let mut g = TensorDependencyGraph::new();
408 g.add_node(42);
409 g.mark_dirty(42);
410 assert!(g.dirty.is_dirty(42));
411 assert_eq!(g.dirty.dirty.len(), 1);
412 }
413
414 #[test]
416 fn test_mark_dirty_no_propagate_tensor_input() {
417 let mut g = TensorDependencyGraph::new();
418 g.add_edge(edge(1, 2, DependencyKind::TensorInput));
419 g.mark_dirty(1);
420 assert!(g.dirty.is_dirty(1));
421 assert!(!g.dirty.is_dirty(2));
422 }
423
424 #[test]
426 fn test_recompute_order_empty_dirty() {
427 let g = TensorDependencyGraph::new();
428 assert!(g.recompute_order().is_empty());
429 }
430
431 #[test]
433 fn test_recompute_order_single_dirty() {
434 let mut g = TensorDependencyGraph::new();
435 g.add_node(7);
436 g.mark_dirty(7);
437 let order = g.recompute_order();
438 assert_eq!(order, vec![7]);
439 }
440
441 #[test]
443 fn test_recompute_order_topo() {
444 let mut g = TensorDependencyGraph::new();
445 g.add_edge(edge(1, 2, DependencyKind::TensorOutput));
447 g.add_edge(edge(2, 3, DependencyKind::TensorOutput));
448 g.mark_dirty(1);
449 let order = g.recompute_order();
450 let pos = |n: u64| {
451 order
452 .iter()
453 .position(|&x| x == n)
454 .expect("test: should succeed")
455 };
456 assert!(pos(1) < pos(2));
457 assert!(pos(2) < pos(3));
458 }
459
460 #[test]
462 fn test_recompute_order_diamond() {
463 let mut g = TensorDependencyGraph::new();
464 g.add_edge(edge(1, 2, DependencyKind::TensorOutput));
470 g.add_edge(edge(1, 3, DependencyKind::TensorOutput));
471 g.add_edge(edge(2, 4, DependencyKind::TensorOutput));
472 g.add_edge(edge(3, 4, DependencyKind::TensorOutput));
473 g.mark_dirty(1);
474 let order = g.recompute_order();
475 let pos = |n: u64| {
476 order
477 .iter()
478 .position(|&x| x == n)
479 .expect("test: should succeed")
480 };
481 assert!(pos(1) < pos(2));
482 assert!(pos(1) < pos(3));
483 assert!(pos(2) < pos(4));
484 assert!(pos(3) < pos(4));
485 assert_eq!(order.len(), 4);
486 }
487
488 #[test]
490 fn test_dependents_of() {
491 let mut g = TensorDependencyGraph::new();
492 g.add_edge(edge(1, 3, DependencyKind::TensorOutput));
493 g.add_edge(edge(1, 2, DependencyKind::RuleImplication));
494 g.add_edge(edge(5, 1, DependencyKind::SharedFact));
495 let deps = g.dependents_of(1);
496 assert_eq!(deps, vec![2, 3]);
497 }
498
499 #[test]
501 fn test_dependencies_of() {
502 let mut g = TensorDependencyGraph::new();
503 g.add_edge(edge(2, 1, DependencyKind::TensorOutput));
504 g.add_edge(edge(3, 1, DependencyKind::RuleImplication));
505 g.add_edge(edge(1, 5, DependencyKind::SharedFact));
506 let deps = g.dependencies_of(1);
507 assert_eq!(deps, vec![2, 3]);
508 }
509
510 #[test]
512 fn test_dirtyset_mark_check_clear() {
513 let mut ds = DirtySet::new();
514 assert!(!ds.is_dirty(1));
515 ds.mark_dirty(1);
516 assert!(ds.is_dirty(1));
517 ds.clear_dirty(1);
518 assert!(!ds.is_dirty(1));
519 }
520
521 #[test]
523 fn test_dirtyset_all_dirty_sorted() {
524 let mut ds = DirtySet::new();
525 ds.mark_dirty(5);
526 ds.mark_dirty(2);
527 ds.mark_dirty(8);
528 ds.mark_dirty(1);
529 assert_eq!(ds.all_dirty(), vec![1, 2, 5, 8]);
530 }
531
532 #[test]
534 fn test_stats_node_edge_count() {
535 let mut g = TensorDependencyGraph::new();
536 g.add_node(1);
537 g.add_node(2);
538 g.add_edge(edge(1, 2, DependencyKind::TensorInput));
539 g.add_edge(edge(1, 2, DependencyKind::TensorOutput));
540 let s = g.stats();
541 assert_eq!(s.node_count, 2);
542 assert_eq!(s.edge_count, 2);
543 }
544
545 #[test]
547 fn test_max_in_degree() {
548 let mut g = TensorDependencyGraph::new();
549 g.add_edge(edge(1, 3, DependencyKind::TensorOutput));
550 g.add_edge(edge(2, 3, DependencyKind::TensorOutput));
551 g.add_edge(edge(4, 3, DependencyKind::TensorOutput));
552 let s = g.stats();
553 assert_eq!(s.max_in_degree, 3);
554 }
555
556 #[test]
558 fn test_max_out_degree() {
559 let mut g = TensorDependencyGraph::new();
560 g.add_edge(edge(1, 2, DependencyKind::TensorOutput));
561 g.add_edge(edge(1, 3, DependencyKind::TensorOutput));
562 g.add_edge(edge(1, 4, DependencyKind::TensorOutput));
563 let s = g.stats();
564 assert_eq!(s.max_out_degree, 3);
565 }
566
567 #[test]
569 fn test_clear_all_dirty() {
570 let mut g = TensorDependencyGraph::new();
571 g.add_node(1);
572 g.add_node(2);
573 g.mark_dirty(1);
574 g.mark_dirty(2);
575 assert_eq!(g.dirty.dirty.len(), 2);
576 g.clear_all_dirty();
577 assert!(g.dirty.dirty.is_empty());
578 }
579
580 #[test]
582 fn test_cycle_no_panic() {
583 let mut g = TensorDependencyGraph::new();
584 g.add_edge(edge(1, 2, DependencyKind::RuleImplication));
586 g.add_edge(edge(2, 1, DependencyKind::RuleImplication));
587 g.mark_dirty(1);
588 let order = g.recompute_order();
589 assert_eq!(order.len(), 2);
591 assert!(order.contains(&1));
592 assert!(order.contains(&2));
593 }
594
595 #[test]
597 fn test_dependency_kind_variants() {
598 let kinds = [
599 DependencyKind::TensorInput,
600 DependencyKind::TensorOutput,
601 DependencyKind::RuleImplication,
602 DependencyKind::SharedFact,
603 ];
604 for k in kinds {
606 let k2 = k;
607 assert_eq!(k, k2);
608 }
609 }
610
611 #[test]
613 fn test_weight_preserved() {
614 let mut g = TensorDependencyGraph::new();
615 g.add_edge(edge_w(1, 2, DependencyKind::TensorOutput, 42));
616 let e = &g.edges[0];
617 assert_eq!(e.weight, 42);
618 assert_eq!(e.from_id, 1);
619 assert_eq!(e.to_id, 2);
620 }
621
622 #[test]
624 fn test_stats_dirty_count() {
625 let mut g = TensorDependencyGraph::new();
626 g.add_node(1);
627 g.add_node(2);
628 g.add_node(3);
629 g.mark_dirty(1);
630 g.mark_dirty(3);
631 let s = g.stats();
632 assert_eq!(s.dirty_count, 2);
633 }
634
635 #[test]
637 fn test_remove_node_leaves_other_edges() {
638 let mut g = TensorDependencyGraph::new();
639 g.add_edge(edge(1, 2, DependencyKind::TensorOutput));
640 g.add_edge(edge(3, 4, DependencyKind::TensorOutput));
641 g.remove_node(1);
642 assert_eq!(g.edges.len(), 1);
643 assert_eq!(g.edges[0].from_id, 3);
644 }
645
646 #[test]
648 fn test_mark_dirty_terminates_with_shared_targets() {
649 let mut g = TensorDependencyGraph::new();
650 g.add_edge(edge(1, 2, DependencyKind::TensorOutput));
651 g.add_edge(edge(1, 2, DependencyKind::TensorOutput)); g.mark_dirty(1);
653 assert!(g.dirty.is_dirty(2));
654 }
656}