rustc_data_structures/graph/linked_graph/
mod.rs1use std::fmt::Debug;
24
25use rustc_index::bit_set::DenseBitSet;
26use tracing::debug;
27
28#[cfg(test)]
29mod tests;
30
31pub struct LinkedGraph<N, E> {
48    nodes: Vec<Node<N>>,
49    edges: Vec<Edge<E>>,
50}
51
52pub struct Node<N> {
53    first_edge: [EdgeIndex; 2], pub data: N,
55}
56
57#[derive(Debug)]
58pub struct Edge<E> {
59    next_edge: [EdgeIndex; 2], source: NodeIndex,
61    target: NodeIndex,
62    pub data: E,
63}
64
65#[derive(Copy, Clone, PartialEq, Debug)]
66pub struct NodeIndex(pub usize);
67
68#[derive(Copy, Clone, PartialEq, Debug)]
69pub struct EdgeIndex(pub usize);
70
71pub const INVALID_EDGE_INDEX: EdgeIndex = EdgeIndex(usize::MAX);
72
73#[derive(Copy, Clone, Debug, PartialEq)]
75pub struct Direction {
76    repr: usize,
77}
78
79pub const OUTGOING: Direction = Direction { repr: 0 };
80
81pub const INCOMING: Direction = Direction { repr: 1 };
82
83impl NodeIndex {
84    pub fn node_id(self) -> usize {
86        self.0
87    }
88}
89
90impl<N: Debug, E: Debug> LinkedGraph<N, E> {
91    pub fn new() -> Self {
92        Self { nodes: Vec::new(), edges: Vec::new() }
93    }
94
95    pub fn with_capacity(nodes: usize, edges: usize) -> Self {
96        Self { nodes: Vec::with_capacity(nodes), edges: Vec::with_capacity(edges) }
97    }
98
99    #[inline]
102    pub fn all_nodes(&self) -> &[Node<N>] {
103        &self.nodes
104    }
105
106    #[inline]
107    pub fn len_nodes(&self) -> usize {
108        self.nodes.len()
109    }
110
111    #[inline]
112    pub fn all_edges(&self) -> &[Edge<E>] {
113        &self.edges
114    }
115
116    #[inline]
117    pub fn len_edges(&self) -> usize {
118        self.edges.len()
119    }
120
121    pub fn next_node_index(&self) -> NodeIndex {
124        NodeIndex(self.nodes.len())
125    }
126
127    pub fn add_node(&mut self, data: N) -> NodeIndex {
128        let idx = self.next_node_index();
129        self.nodes.push(Node { first_edge: [INVALID_EDGE_INDEX, INVALID_EDGE_INDEX], data });
130        idx
131    }
132
133    pub fn mut_node_data(&mut self, idx: NodeIndex) -> &mut N {
134        &mut self.nodes[idx.0].data
135    }
136
137    pub fn node_data(&self, idx: NodeIndex) -> &N {
138        &self.nodes[idx.0].data
139    }
140
141    pub fn node(&self, idx: NodeIndex) -> &Node<N> {
142        &self.nodes[idx.0]
143    }
144
145    pub fn next_edge_index(&self) -> EdgeIndex {
148        EdgeIndex(self.edges.len())
149    }
150
151    pub fn add_edge(&mut self, source: NodeIndex, target: NodeIndex, data: E) -> EdgeIndex {
152        debug!("graph: add_edge({:?}, {:?}, {:?})", source, target, data);
153
154        let idx = self.next_edge_index();
155
156        let source_first = self.nodes[source.0].first_edge[OUTGOING.repr];
158        let target_first = self.nodes[target.0].first_edge[INCOMING.repr];
159
160        self.edges.push(Edge { next_edge: [source_first, target_first], source, target, data });
163
164        self.nodes[source.0].first_edge[OUTGOING.repr] = idx;
166        self.nodes[target.0].first_edge[INCOMING.repr] = idx;
167
168        idx
169    }
170
171    pub fn edge(&self, idx: EdgeIndex) -> &Edge<E> {
172        &self.edges[idx.0]
173    }
174
175    pub fn enumerated_nodes(&self) -> impl Iterator<Item = (NodeIndex, &Node<N>)> {
178        self.nodes.iter().enumerate().map(|(idx, n)| (NodeIndex(idx), n))
179    }
180
181    pub fn enumerated_edges(&self) -> impl Iterator<Item = (EdgeIndex, &Edge<E>)> {
182        self.edges.iter().enumerate().map(|(idx, e)| (EdgeIndex(idx), e))
183    }
184
185    pub fn each_node<'a>(&'a self, mut f: impl FnMut(NodeIndex, &'a Node<N>) -> bool) -> bool {
186        self.enumerated_nodes().all(|(node_idx, node)| f(node_idx, node))
188    }
189
190    pub fn each_edge<'a>(&'a self, mut f: impl FnMut(EdgeIndex, &'a Edge<E>) -> bool) -> bool {
191        self.enumerated_edges().all(|(edge_idx, edge)| f(edge_idx, edge))
193    }
194
195    pub fn outgoing_edges(&self, source: NodeIndex) -> AdjacentEdges<'_, N, E> {
196        self.adjacent_edges(source, OUTGOING)
197    }
198
199    pub fn incoming_edges(&self, source: NodeIndex) -> AdjacentEdges<'_, N, E> {
200        self.adjacent_edges(source, INCOMING)
201    }
202
203    pub fn adjacent_edges(
204        &self,
205        source: NodeIndex,
206        direction: Direction,
207    ) -> AdjacentEdges<'_, N, E> {
208        let first_edge = self.node(source).first_edge[direction.repr];
209        AdjacentEdges { graph: self, direction, next: first_edge }
210    }
211
212    pub fn successor_nodes(&self, source: NodeIndex) -> impl Iterator<Item = NodeIndex> {
213        self.outgoing_edges(source).targets()
214    }
215
216    pub fn predecessor_nodes(&self, target: NodeIndex) -> impl Iterator<Item = NodeIndex> {
217        self.incoming_edges(target).sources()
218    }
219
220    pub fn depth_traverse(
221        &self,
222        start: NodeIndex,
223        direction: Direction,
224    ) -> DepthFirstTraversal<'_, N, E> {
225        DepthFirstTraversal::with_start_node(self, start, direction)
226    }
227
228    pub fn nodes_in_postorder(
229        &self,
230        direction: Direction,
231        entry_node: NodeIndex,
232    ) -> Vec<NodeIndex> {
233        let mut visited = DenseBitSet::new_empty(self.len_nodes());
234        let mut stack = vec![];
235        let mut result = Vec::with_capacity(self.len_nodes());
236        let mut push_node = |stack: &mut Vec<_>, node: NodeIndex| {
237            if visited.insert(node.0) {
238                stack.push((node, self.adjacent_edges(node, direction)));
239            }
240        };
241
242        for node in
243            Some(entry_node).into_iter().chain(self.enumerated_nodes().map(|(node, _)| node))
244        {
245            push_node(&mut stack, node);
246            while let Some((node, mut iter)) = stack.pop() {
247                if let Some((_, child)) = iter.next() {
248                    let target = child.source_or_target(direction);
249                    stack.push((node, iter));
252                    push_node(&mut stack, target);
254                } else {
255                    result.push(node);
256                }
257            }
258        }
259
260        assert_eq!(result.len(), self.len_nodes());
261        result
262    }
263}
264
265pub struct AdjacentEdges<'g, N, E> {
268    graph: &'g LinkedGraph<N, E>,
269    direction: Direction,
270    next: EdgeIndex,
271}
272
273impl<'g, N: Debug, E: Debug> AdjacentEdges<'g, N, E> {
274    fn targets(self) -> impl Iterator<Item = NodeIndex> {
275        self.map(|(_, edge)| edge.target)
276    }
277
278    fn sources(self) -> impl Iterator<Item = NodeIndex> {
279        self.map(|(_, edge)| edge.source)
280    }
281}
282
283impl<'g, N: Debug, E: Debug> Iterator for AdjacentEdges<'g, N, E> {
284    type Item = (EdgeIndex, &'g Edge<E>);
285
286    fn next(&mut self) -> Option<(EdgeIndex, &'g Edge<E>)> {
287        let edge_index = self.next;
288        if edge_index == INVALID_EDGE_INDEX {
289            return None;
290        }
291
292        let edge = self.graph.edge(edge_index);
293        self.next = edge.next_edge[self.direction.repr];
294        Some((edge_index, edge))
295    }
296
297    fn size_hint(&self) -> (usize, Option<usize>) {
298        (0, Some(self.graph.len_edges()))
300    }
301}
302
303pub struct DepthFirstTraversal<'g, N, E> {
304    graph: &'g LinkedGraph<N, E>,
305    stack: Vec<NodeIndex>,
306    visited: DenseBitSet<usize>,
307    direction: Direction,
308}
309
310impl<'g, N: Debug, E: Debug> DepthFirstTraversal<'g, N, E> {
311    pub fn with_start_node(
312        graph: &'g LinkedGraph<N, E>,
313        start_node: NodeIndex,
314        direction: Direction,
315    ) -> Self {
316        let mut visited = DenseBitSet::new_empty(graph.len_nodes());
317        visited.insert(start_node.node_id());
318        DepthFirstTraversal { graph, stack: vec![start_node], visited, direction }
319    }
320
321    fn visit(&mut self, node: NodeIndex) {
322        if self.visited.insert(node.node_id()) {
323            self.stack.push(node);
324        }
325    }
326}
327
328impl<'g, N: Debug, E: Debug> Iterator for DepthFirstTraversal<'g, N, E> {
329    type Item = NodeIndex;
330
331    fn next(&mut self) -> Option<NodeIndex> {
332        let next = self.stack.pop();
333        if let Some(idx) = next {
334            for (_, edge) in self.graph.adjacent_edges(idx, self.direction) {
335                let target = edge.source_or_target(self.direction);
336                self.visit(target);
337            }
338        }
339        next
340    }
341
342    fn size_hint(&self) -> (usize, Option<usize>) {
343        let remaining = self.graph.len_nodes() - self.visited.count();
345        (remaining, Some(remaining))
346    }
347}
348
349impl<'g, N: Debug, E: Debug> ExactSizeIterator for DepthFirstTraversal<'g, N, E> {}
350
351impl<E> Edge<E> {
352    pub fn source(&self) -> NodeIndex {
353        self.source
354    }
355
356    pub fn target(&self) -> NodeIndex {
357        self.target
358    }
359
360    pub fn source_or_target(&self, direction: Direction) -> NodeIndex {
361        if direction == OUTGOING { self.target } else { self.source }
362    }
363}