teeny_core/utils/
dag.rs

1/*
2 * Copyright (c) 2026 Teenygrad.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *   http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17use alloc::vec::Vec;
18
19/// A node in the DAG. Carries a boxed trait object and a list of children.
20pub struct Node<V> {
21    /// This node's payload.
22    pub value: V,
23    /// Indices of this node's children in the owning `Dag`.
24    pub children: Vec<usize>, // indices into Dag.nodes
25    /// Indices of this node's parents in the owning `Dag`.
26    pub parents: Vec<usize>, // indices into Dag.nodes (needed for topological sort, etc)
27}
28
29/// A directed acyclic graph of `V`-valued nodes.
30pub struct Dag<V> {
31    nodes: Vec<Node<V>>,
32}
33
34impl<V> Default for Dag<V> {
35    fn default() -> Self {
36        Self { nodes: Vec::new() }
37    }
38}
39
40impl<V> Dag<V> {
41    /// Create an empty DAG.
42    pub fn new() -> Self {
43        Self::default()
44    }
45
46    /// Adds a node to the graph, returns index of the node.
47    pub fn add_node(&mut self, value: V) -> usize {
48        let idx = self.nodes.len();
49        self.nodes.push(Node {
50            value,
51            children: Vec::new(),
52            parents: Vec::new(),
53        });
54        idx
55    }
56
57    /// Adds a directed edge from `from` to `to`.
58    /// Panics if indices are out of bounds or would cause a self-loop.
59    /// Ignores duplicate edges.
60    pub fn add_edge(&mut self, from: usize, to: usize) {
61        assert!(from != to, "self-loop detected");
62        assert!(
63            from < self.nodes.len() && to < self.nodes.len(),
64            "invalid node index"
65        );
66        // Check for duplicate edges
67        if !self.nodes[from].children.contains(&to) {
68            self.nodes[from].children.push(to);
69        }
70        if !self.nodes[to].parents.contains(&from) {
71            self.nodes[to].parents.push(from);
72        }
73        // NOT checking for cycles for no_std simplicity
74    }
75
76    /// Returns a reference to the node at the given index.
77    pub fn node(&self, idx: usize) -> &Node<V> {
78        &self.nodes[idx]
79    }
80
81    /// Returns a mutable reference to the node at the given index.
82    pub fn node_mut(&mut self, idx: usize) -> &mut Node<V> {
83        &mut self.nodes[idx]
84    }
85
86    /// Returns the number of nodes in the graph.
87    pub fn len(&self) -> usize {
88        self.nodes.len()
89    }
90
91    /// Returns true if the graph contains no nodes.
92    pub fn is_empty(&self) -> bool {
93        self.nodes.is_empty()
94    }
95
96    /// Perform a topological sort. Returns a list of node indices.
97    /// Panics if the graph contains a cycle.
98    pub fn topological_sort(&self) -> Vec<usize> {
99        let mut in_degree = alloc::vec![0; self.nodes.len()];
100        for node in &self.nodes {
101            for &child in &node.children {
102                in_degree[child] += 1;
103            }
104        }
105
106        let mut stack = Vec::new();
107        for (i, &deg) in in_degree.iter().enumerate() {
108            if deg == 0 {
109                stack.push(i);
110            }
111        }
112
113        let mut order = Vec::with_capacity(self.nodes.len());
114        let mut remaining_in_degree = in_degree.clone();
115
116        while let Some(n) = stack.pop() {
117            order.push(n);
118            for &child in &self.nodes[n].children {
119                remaining_in_degree[child] -= 1;
120                if remaining_in_degree[child] == 0 {
121                    stack.push(child);
122                }
123            }
124        }
125
126        assert!(
127            order.len() == self.nodes.len(),
128            "cycle detected in DAG (not acyclic)"
129        );
130        order
131    }
132}
133
134impl<V> IntoIterator for Dag<V> {
135    type Item = Node<V>;
136    type IntoIter = alloc::vec::IntoIter<Node<V>>;
137
138    fn into_iter(self) -> Self::IntoIter {
139        self.nodes.into_iter()
140    }
141}
142
143#[cfg(test)]
144mod tests {
145    use alloc::vec;
146
147    use super::*;
148
149    #[test]
150    fn test_empty_dag() {
151        let dag: Dag<i32> = Dag { nodes: Vec::new() };
152        assert_eq!(dag.len(), 0);
153        assert!(dag.is_empty());
154        assert_eq!(dag.topological_sort(), Vec::<usize>::new());
155    }
156
157    #[test]
158    fn test_single_node() {
159        let node = Node {
160            value: 42,
161            parents: Vec::new(),
162            children: Vec::new(),
163        };
164        let dag = Dag { nodes: vec![node] };
165        assert_eq!(dag.len(), 1);
166        assert!(!dag.is_empty());
167        assert_eq!(dag.topological_sort(), vec![0]);
168    }
169
170    #[test]
171    fn test_chain() {
172        // 0 -> 1 -> 2
173        let n2 = Node {
174            value: 2,
175            children: Vec::new(),
176            parents: Vec::new(),
177        };
178        let n1 = Node {
179            value: 1,
180            children: vec![2],
181            parents: Vec::new(),
182        };
183        let n0 = Node {
184            value: 0,
185            children: vec![1],
186            parents: Vec::new(),
187        };
188        let dag = Dag {
189            nodes: vec![n0, n1, n2],
190        };
191        let order = dag.topological_sort();
192        // The only valid topo order is [0,1,2]
193        assert_eq!(order, vec![0, 1, 2]);
194    }
195
196    #[test]
197    fn test_branch() {
198        //   0
199        //  / \
200        // 1   2
201        //  \ /
202        //   3
203        let n3 = Node {
204            value: 3,
205            children: Vec::new(),
206            parents: Vec::new(),
207        };
208        let n2 = Node {
209            value: 2,
210            children: vec![3],
211            parents: Vec::new(),
212        };
213        let n1 = Node {
214            value: 1,
215            children: vec![3],
216            parents: Vec::new(),
217        };
218        let n0 = Node {
219            value: 0,
220            children: vec![1, 2],
221            parents: Vec::new(),
222        };
223        let dag = Dag {
224            nodes: vec![n0, n1, n2, n3],
225        };
226        let order = dag.topological_sort();
227        // 0 before 1 and 2, 1/2 before 3
228        let pos = |n| order.iter().position(|&x| x == n).unwrap();
229        assert!(pos(0) < pos(1) && pos(0) < pos(2));
230        assert!(pos(1) < pos(3));
231        assert!(pos(2) < pos(3));
232    }
233
234    #[test]
235    #[should_panic(expected = "cycle detected in DAG")]
236    fn test_cycle_panics() {
237        // 0 -> 1 -> 2 -> 0 (cycle)
238        let n0 = Node {
239            value: 0,
240            children: vec![1],
241            parents: Vec::new(),
242        };
243        let n1 = Node {
244            value: 1,
245            children: vec![2],
246            parents: Vec::new(),
247        };
248        let n2 = Node {
249            value: 2,
250            children: vec![0],
251            parents: Vec::new(),
252        };
253        let dag = Dag {
254            nodes: vec![n0, n1, n2],
255        };
256        let _ = dag.topological_sort(); // Should panic
257    }
258}
teeny_core/utils/dag.rs

teeny_core/utils/
dag.rs
