Skip to main content

ipfrs_tensorlogic/
dependency_graph.rs

1//! Tensor Dependency Graph — tracks data dependencies between tensors and rules,
2//! enabling incremental recomputation when tensors are updated.
3
4use std::collections::{HashMap, HashSet, VecDeque};
5
6// ---------------------------------------------------------------------------
7// DependencyKind
8// ---------------------------------------------------------------------------
9
10/// The semantic relationship represented by a dependency edge.
11#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
12pub enum DependencyKind {
13    /// The source tensor is an input to the destination rule.
14    TensorInput,
15    /// The destination tensor is produced by the source rule.
16    TensorOutput,
17    /// The source rule implies the destination rule must re-run.
18    RuleImplication,
19    /// Two rules share a fact dependency.
20    SharedFact,
21}
22
23// ---------------------------------------------------------------------------
24// DependencyEdge
25// ---------------------------------------------------------------------------
26
27/// A directed dependency edge between two graph nodes.
28#[derive(Clone, Debug, PartialEq, Eq)]
29pub struct DependencyEdge {
30    /// Source node ID.
31    pub from_id: u64,
32    /// Target node ID.
33    pub to_id: u64,
34    /// The semantic kind of the dependency.
35    pub kind: DependencyKind,
36    /// Edge weight: 1 = normal, higher = stronger dependency.
37    pub weight: u32,
38}
39
40// ---------------------------------------------------------------------------
41// DirtySet
42// ---------------------------------------------------------------------------
43
44/// Set of node IDs that require recomputation.
45#[derive(Clone, Debug, Default)]
46pub struct DirtySet {
47    pub dirty: HashSet<u64>,
48}
49
50impl DirtySet {
51    /// Create a new, empty `DirtySet`.
52    pub fn new() -> Self {
53        Self {
54            dirty: HashSet::new(),
55        }
56    }
57
58    /// Mark `id` as needing recomputation.
59    pub fn mark_dirty(&mut self, id: u64) {
60        self.dirty.insert(id);
61    }
62
63    /// Returns `true` if `id` is currently dirty.
64    pub fn is_dirty(&self, id: u64) -> bool {
65        self.dirty.contains(&id)
66    }
67
68    /// Clear the dirty flag for `id`.
69    pub fn clear_dirty(&mut self, id: u64) {
70        self.dirty.remove(&id);
71    }
72
73    /// Return all dirty IDs, sorted ascending.
74    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// ---------------------------------------------------------------------------
82// GraphStats
83// ---------------------------------------------------------------------------
84
85/// Statistics snapshot for a [`TensorDependencyGraph`].
86#[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
95// ---------------------------------------------------------------------------
96// TensorDependencyGraph
97// ---------------------------------------------------------------------------
98
99/// Tracks data dependencies between tensors and rules, enabling incremental
100/// recomputation when tensors are updated.
101pub struct TensorDependencyGraph {
102    /// All registered node IDs.
103    pub nodes: HashSet<u64>,
104    /// All dependency edges.
105    pub edges: Vec<DependencyEdge>,
106    /// Set of dirty (stale) node IDs.
107    pub dirty: DirtySet,
108}
109
110impl TensorDependencyGraph {
111    /// Create a new, empty graph.
112    pub fn new() -> Self {
113        Self {
114            nodes: HashSet::new(),
115            edges: Vec::new(),
116            dirty: DirtySet::new(),
117        }
118    }
119
120    /// Register a node — idempotent.
121    pub fn add_node(&mut self, id: u64) {
122        self.nodes.insert(id);
123    }
124
125    /// Add a dependency edge.  Also registers the `from` and `to` nodes.
126    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    /// Remove a node and all edges that reference it.  Clears any dirty flag.
133    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    /// Mark `id` dirty and transitively propagate to all reachable dependents
140    /// via `TensorOutput`, `RuleImplication`, and `SharedFact` edges where
141    /// `from_id == id`.
142    pub fn mark_dirty(&mut self, id: u64) {
143        // BFS over outgoing edges with the three propagating kinds.
144        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            // Collect successors first to avoid borrow issues.
150            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    /// Topological sort (Kahn's algorithm) of dirty nodes only.
175    /// Returns nodes in processing order (dependencies before dependents).
176    /// On cycle detection, appends remaining nodes in arbitrary order.
177    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        // Build in-degree map considering only edges within the dirty subgraph.
184        let mut in_degree: HashMap<u64, usize> = dirty.iter().map(|&n| (n, 0)).collect();
185        // Adjacency list restricted to dirty nodes.
186        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)| deg == 0)
198            .map(|(&n, _)| n)
199            .collect();
200
201        // Deterministic ordering within equal in-degree tiers.
202        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        // Handle cycle — append remaining nodes in sorted order.
226        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    /// Direct successors of `id` (edges where `from_id == id`), sorted ascending.
241    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    /// Direct predecessors of `id` (edges where `to_id == id`), sorted ascending.
254    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    /// Compute and return a statistics snapshot.
267    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    /// Clear the entire dirty set.
298    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// ---------------------------------------------------------------------------
310// Tests
311// ---------------------------------------------------------------------------
312
313#[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    // 1. Empty graph stats
336    #[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    // 2. add_node idempotent
348    #[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    // 3. add_edge registers nodes
358    #[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    // 4. remove_node cleans edges
368    #[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    // 5. remove_node clears dirty
379    #[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    // 6. mark_dirty propagates transitively
390    #[test]
391    fn test_mark_dirty_propagates_transitively() {
392        let mut g = TensorDependencyGraph::new();
393        // chain: 1 ->TensorOutput-> 2 ->RuleImplication-> 3 ->SharedFact-> 4
394        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    // 7. mark_dirty single node no propagation if no outgoing edges
405    #[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    // 8. mark_dirty does NOT propagate over TensorInput edges
415    #[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    // 9. recompute_order empty dirty
425    #[test]
426    fn test_recompute_order_empty_dirty() {
427        let g = TensorDependencyGraph::new();
428        assert!(g.recompute_order().is_empty());
429    }
430
431    // 10. recompute_order single dirty node
432    #[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    // 11. recompute_order respects topo order (dependency before dependent)
442    #[test]
443    fn test_recompute_order_topo() {
444        let mut g = TensorDependencyGraph::new();
445        // A -> B -> C
446        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    // 12. recompute_order handles diamond dependency
461    #[test]
462    fn test_recompute_order_diamond() {
463        let mut g = TensorDependencyGraph::new();
464        //        1
465        //       / \
466        //      2   3
467        //       \ /
468        //        4
469        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    // 13. dependents_of correct
489    #[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    // 14. dependencies_of correct
500    #[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    // 15. DirtySet mark/check/clear
511    #[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    // 16. DirtySet all_dirty sorted
522    #[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    // 17. GraphStats node_count / edge_count
533    #[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    // 18. max_in_degree computed correctly
546    #[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    // 19. max_out_degree computed correctly
557    #[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    // 20. clear_all_dirty empties dirty set
568    #[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    // 21. cycle handling doesn't panic
581    #[test]
582    fn test_cycle_no_panic() {
583        let mut g = TensorDependencyGraph::new();
584        // A -> B -> A (cycle)
585        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        // Both nodes should be in the result, no panic.
590        assert_eq!(order.len(), 2);
591        assert!(order.contains(&1));
592        assert!(order.contains(&2));
593    }
594
595    // 22. DependencyKind variants accessible
596    #[test]
597    fn test_dependency_kind_variants() {
598        let kinds = [
599            DependencyKind::TensorInput,
600            DependencyKind::TensorOutput,
601            DependencyKind::RuleImplication,
602            DependencyKind::SharedFact,
603        ];
604        // Each variant should be copy-able and comparable.
605        for k in kinds {
606            let k2 = k;
607            assert_eq!(k, k2);
608        }
609    }
610
611    // 23. weight preserved in edge
612    #[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    // 24. stats dirty_count reflects current dirty state
623    #[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    // 25. remove_node leaves unrelated edges intact
636    #[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    // 26. mark_dirty does not re-enqueue already-dirty nodes (termination)
647    #[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)); // duplicate edge
652        g.mark_dirty(1);
653        assert!(g.dirty.is_dirty(2));
654        // Should not infinite-loop.
655    }
656}