rustc_data_structures/graph/implementation/
mod.rs

1//! A graph module for use in dataflow, region resolution, and elsewhere.
2//!
3//! # Interface details
4//!
5//! You customize the graph by specifying a "node data" type `N` and an
6//! "edge data" type `E`. You can then later gain access (mutable or
7//! immutable) to these "user-data" bits. Currently, you can only add
8//! nodes or edges to the graph. You cannot remove or modify them once
9//! added. This could be changed if we have a need.
10//!
11//! # Implementation details
12//!
13//! The main tricky thing about this code is the way that edges are
14//! stored. The edges are stored in a central array, but they are also
15//! threaded onto two linked lists for each node, one for incoming edges
16//! and one for outgoing edges. Note that every edge is a member of some
17//! incoming list and some outgoing list. Basically you can load the
18//! first index of the linked list from the node data structures (the
19//! field `first_edge`) and then, for each edge, load the next index from
20//! the field `next_edge`). Each of those fields is an array that should
21//! be indexed by the direction (see the type `Direction`).
22
23use std::fmt::Debug;
24
25use rustc_index::bit_set::DenseBitSet;
26use tracing::debug;
27
28#[cfg(test)]
29mod tests;
30
31pub struct Graph<N, E> {
32    nodes: Vec<Node<N>>,
33    edges: Vec<Edge<E>>,
34}
35
36pub struct Node<N> {
37    first_edge: [EdgeIndex; 2], // see module comment
38    pub data: N,
39}
40
41#[derive(Debug)]
42pub struct Edge<E> {
43    next_edge: [EdgeIndex; 2], // see module comment
44    source: NodeIndex,
45    target: NodeIndex,
46    pub data: E,
47}
48
49#[derive(Copy, Clone, PartialEq, Debug)]
50pub struct NodeIndex(pub usize);
51
52#[derive(Copy, Clone, PartialEq, Debug)]
53pub struct EdgeIndex(pub usize);
54
55pub const INVALID_EDGE_INDEX: EdgeIndex = EdgeIndex(usize::MAX);
56
57// Use a private field here to guarantee no more instances are created:
58#[derive(Copy, Clone, Debug, PartialEq)]
59pub struct Direction {
60    repr: usize,
61}
62
63pub const OUTGOING: Direction = Direction { repr: 0 };
64
65pub const INCOMING: Direction = Direction { repr: 1 };
66
67impl NodeIndex {
68    /// Returns unique ID (unique with respect to the graph holding associated node).
69    pub fn node_id(self) -> usize {
70        self.0
71    }
72}
73
74impl<N: Debug, E: Debug> Graph<N, E> {
75    pub fn new() -> Graph<N, E> {
76        Graph { nodes: Vec::new(), edges: Vec::new() }
77    }
78
79    pub fn with_capacity(nodes: usize, edges: usize) -> Graph<N, E> {
80        Graph { nodes: Vec::with_capacity(nodes), edges: Vec::with_capacity(edges) }
81    }
82
83    // # Simple accessors
84
85    #[inline]
86    pub fn all_nodes(&self) -> &[Node<N>] {
87        &self.nodes
88    }
89
90    #[inline]
91    pub fn len_nodes(&self) -> usize {
92        self.nodes.len()
93    }
94
95    #[inline]
96    pub fn all_edges(&self) -> &[Edge<E>] {
97        &self.edges
98    }
99
100    #[inline]
101    pub fn len_edges(&self) -> usize {
102        self.edges.len()
103    }
104
105    // # Node construction
106
107    pub fn next_node_index(&self) -> NodeIndex {
108        NodeIndex(self.nodes.len())
109    }
110
111    pub fn add_node(&mut self, data: N) -> NodeIndex {
112        let idx = self.next_node_index();
113        self.nodes.push(Node { first_edge: [INVALID_EDGE_INDEX, INVALID_EDGE_INDEX], data });
114        idx
115    }
116
117    pub fn mut_node_data(&mut self, idx: NodeIndex) -> &mut N {
118        &mut self.nodes[idx.0].data
119    }
120
121    pub fn node_data(&self, idx: NodeIndex) -> &N {
122        &self.nodes[idx.0].data
123    }
124
125    pub fn node(&self, idx: NodeIndex) -> &Node<N> {
126        &self.nodes[idx.0]
127    }
128
129    // # Edge construction and queries
130
131    pub fn next_edge_index(&self) -> EdgeIndex {
132        EdgeIndex(self.edges.len())
133    }
134
135    pub fn add_edge(&mut self, source: NodeIndex, target: NodeIndex, data: E) -> EdgeIndex {
136        debug!("graph: add_edge({:?}, {:?}, {:?})", source, target, data);
137
138        let idx = self.next_edge_index();
139
140        // read current first of the list of edges from each node
141        let source_first = self.nodes[source.0].first_edge[OUTGOING.repr];
142        let target_first = self.nodes[target.0].first_edge[INCOMING.repr];
143
144        // create the new edge, with the previous firsts from each node
145        // as the next pointers
146        self.edges.push(Edge { next_edge: [source_first, target_first], source, target, data });
147
148        // adjust the firsts for each node target be the next object.
149        self.nodes[source.0].first_edge[OUTGOING.repr] = idx;
150        self.nodes[target.0].first_edge[INCOMING.repr] = idx;
151
152        idx
153    }
154
155    pub fn edge(&self, idx: EdgeIndex) -> &Edge<E> {
156        &self.edges[idx.0]
157    }
158
159    // # Iterating over nodes, edges
160
161    pub fn enumerated_nodes(&self) -> impl Iterator<Item = (NodeIndex, &Node<N>)> {
162        self.nodes.iter().enumerate().map(|(idx, n)| (NodeIndex(idx), n))
163    }
164
165    pub fn enumerated_edges(&self) -> impl Iterator<Item = (EdgeIndex, &Edge<E>)> {
166        self.edges.iter().enumerate().map(|(idx, e)| (EdgeIndex(idx), e))
167    }
168
169    pub fn each_node<'a>(&'a self, mut f: impl FnMut(NodeIndex, &'a Node<N>) -> bool) -> bool {
170        //! Iterates over all edges defined in the graph.
171        self.enumerated_nodes().all(|(node_idx, node)| f(node_idx, node))
172    }
173
174    pub fn each_edge<'a>(&'a self, mut f: impl FnMut(EdgeIndex, &'a Edge<E>) -> bool) -> bool {
175        //! Iterates over all edges defined in the graph
176        self.enumerated_edges().all(|(edge_idx, edge)| f(edge_idx, edge))
177    }
178
179    pub fn outgoing_edges(&self, source: NodeIndex) -> AdjacentEdges<'_, N, E> {
180        self.adjacent_edges(source, OUTGOING)
181    }
182
183    pub fn incoming_edges(&self, source: NodeIndex) -> AdjacentEdges<'_, N, E> {
184        self.adjacent_edges(source, INCOMING)
185    }
186
187    pub fn adjacent_edges(
188        &self,
189        source: NodeIndex,
190        direction: Direction,
191    ) -> AdjacentEdges<'_, N, E> {
192        let first_edge = self.node(source).first_edge[direction.repr];
193        AdjacentEdges { graph: self, direction, next: first_edge }
194    }
195
196    pub fn successor_nodes(&self, source: NodeIndex) -> impl Iterator<Item = NodeIndex> + '_ {
197        self.outgoing_edges(source).targets()
198    }
199
200    pub fn predecessor_nodes(&self, target: NodeIndex) -> impl Iterator<Item = NodeIndex> + '_ {
201        self.incoming_edges(target).sources()
202    }
203
204    pub fn depth_traverse(
205        &self,
206        start: NodeIndex,
207        direction: Direction,
208    ) -> DepthFirstTraversal<'_, N, E> {
209        DepthFirstTraversal::with_start_node(self, start, direction)
210    }
211
212    pub fn nodes_in_postorder(
213        &self,
214        direction: Direction,
215        entry_node: NodeIndex,
216    ) -> Vec<NodeIndex> {
217        let mut visited = DenseBitSet::new_empty(self.len_nodes());
218        let mut stack = vec![];
219        let mut result = Vec::with_capacity(self.len_nodes());
220        let mut push_node = |stack: &mut Vec<_>, node: NodeIndex| {
221            if visited.insert(node.0) {
222                stack.push((node, self.adjacent_edges(node, direction)));
223            }
224        };
225
226        for node in
227            Some(entry_node).into_iter().chain(self.enumerated_nodes().map(|(node, _)| node))
228        {
229            push_node(&mut stack, node);
230            while let Some((node, mut iter)) = stack.pop() {
231                if let Some((_, child)) = iter.next() {
232                    let target = child.source_or_target(direction);
233                    // the current node needs more processing, so
234                    // add it back to the stack
235                    stack.push((node, iter));
236                    // and then push the new node
237                    push_node(&mut stack, target);
238                } else {
239                    result.push(node);
240                }
241            }
242        }
243
244        assert_eq!(result.len(), self.len_nodes());
245        result
246    }
247}
248
249// # Iterators
250
251pub struct AdjacentEdges<'g, N, E> {
252    graph: &'g Graph<N, E>,
253    direction: Direction,
254    next: EdgeIndex,
255}
256
257impl<'g, N: Debug, E: Debug> AdjacentEdges<'g, N, E> {
258    fn targets(self) -> impl Iterator<Item = NodeIndex> + 'g {
259        self.map(|(_, edge)| edge.target)
260    }
261
262    fn sources(self) -> impl Iterator<Item = NodeIndex> + 'g {
263        self.map(|(_, edge)| edge.source)
264    }
265}
266
267impl<'g, N: Debug, E: Debug> Iterator for AdjacentEdges<'g, N, E> {
268    type Item = (EdgeIndex, &'g Edge<E>);
269
270    fn next(&mut self) -> Option<(EdgeIndex, &'g Edge<E>)> {
271        let edge_index = self.next;
272        if edge_index == INVALID_EDGE_INDEX {
273            return None;
274        }
275
276        let edge = self.graph.edge(edge_index);
277        self.next = edge.next_edge[self.direction.repr];
278        Some((edge_index, edge))
279    }
280
281    fn size_hint(&self) -> (usize, Option<usize>) {
282        // At most, all the edges in the graph.
283        (0, Some(self.graph.len_edges()))
284    }
285}
286
287pub struct DepthFirstTraversal<'g, N, E> {
288    graph: &'g Graph<N, E>,
289    stack: Vec<NodeIndex>,
290    visited: DenseBitSet<usize>,
291    direction: Direction,
292}
293
294impl<'g, N: Debug, E: Debug> DepthFirstTraversal<'g, N, E> {
295    pub fn with_start_node(
296        graph: &'g Graph<N, E>,
297        start_node: NodeIndex,
298        direction: Direction,
299    ) -> Self {
300        let mut visited = DenseBitSet::new_empty(graph.len_nodes());
301        visited.insert(start_node.node_id());
302        DepthFirstTraversal { graph, stack: vec![start_node], visited, direction }
303    }
304
305    fn visit(&mut self, node: NodeIndex) {
306        if self.visited.insert(node.node_id()) {
307            self.stack.push(node);
308        }
309    }
310}
311
312impl<'g, N: Debug, E: Debug> Iterator for DepthFirstTraversal<'g, N, E> {
313    type Item = NodeIndex;
314
315    fn next(&mut self) -> Option<NodeIndex> {
316        let next = self.stack.pop();
317        if let Some(idx) = next {
318            for (_, edge) in self.graph.adjacent_edges(idx, self.direction) {
319                let target = edge.source_or_target(self.direction);
320                self.visit(target);
321            }
322        }
323        next
324    }
325
326    fn size_hint(&self) -> (usize, Option<usize>) {
327        // We will visit every node in the graph exactly once.
328        let remaining = self.graph.len_nodes() - self.visited.count();
329        (remaining, Some(remaining))
330    }
331}
332
333impl<'g, N: Debug, E: Debug> ExactSizeIterator for DepthFirstTraversal<'g, N, E> {}
334
335impl<E> Edge<E> {
336    pub fn source(&self) -> NodeIndex {
337        self.source
338    }
339
340    pub fn target(&self) -> NodeIndex {
341        self.target
342    }
343
344    pub fn source_or_target(&self, direction: Direction) -> NodeIndex {
345        if direction == OUTGOING { self.target } else { self.source }
346    }
347}