compiler/rustc_data_structures/src/graph/iterate/mod.rs

   1 use super::{DirectedGraph, WithNumNodes, WithStartNode, WithSuccessors};
   2 use rustc_index::bit_set::BitSet;
   3 use rustc_index::vec::IndexVec;
   4 use std::ops::ControlFlow;
   5
   6 #[cfg(test)]
   7 mod tests;
   8
   9 pub fn post_order_from<G: DirectedGraph + WithSuccessors + WithNumNodes>(
  10     graph: &G,
  11     start_node: G::Node,
  12 ) -> Vec<G::Node> {
  13     post_order_from_to(graph, start_node, None)
  14 }
  15
  16 pub fn post_order_from_to<G: DirectedGraph + WithSuccessors + WithNumNodes>(
  17     graph: &G,
  18     start_node: G::Node,
  19     end_node: Option<G::Node>,
  20 ) -> Vec<G::Node> {
  21     let mut visited: IndexVec<G::Node, bool> = IndexVec::from_elem_n(false, graph.num_nodes());
  22     let mut result: Vec<G::Node> = Vec::with_capacity(graph.num_nodes());
  23     if let Some(end_node) = end_node {
  24         visited[end_node] = true;
  25     }
  26     post_order_walk(graph, start_node, &mut result, &mut visited);
  27     result
  28 }
  29
  30 fn post_order_walk<G: DirectedGraph + WithSuccessors + WithNumNodes>(
  31     graph: &G,
  32     node: G::Node,
  33     result: &mut Vec<G::Node>,
  34     visited: &mut IndexVec<G::Node, bool>,
  35 ) {
  36     if visited[node] {
  37         return;
  38     }
  39     visited[node] = true;
  40
  41     for successor in graph.successors(node) {
  42         post_order_walk(graph, successor, result, visited);
  43     }
  44
  45     result.push(node);
  46 }
  47
  48 pub fn reverse_post_order<G: DirectedGraph + WithSuccessors + WithNumNodes>(
  49     graph: &G,
  50     start_node: G::Node,
  51 ) -> Vec<G::Node> {
  52     let mut vec = post_order_from(graph, start_node);
  53     vec.reverse();
  54     vec
  55 }
  56
  57 /// A "depth-first search" iterator for a directed graph.
  58 pub struct DepthFirstSearch<'graph, G>
  59 where
  60     G: ?Sized + DirectedGraph + WithNumNodes + WithSuccessors,
  61 {
  62     graph: &'graph G,
  63     stack: Vec<G::Node>,
  64     visited: BitSet<G::Node>,
  65 }
  66
  67 impl<G> DepthFirstSearch<'graph, G>
  68 where
  69     G: ?Sized + DirectedGraph + WithNumNodes + WithSuccessors,
  70 {
  71     pub fn new(graph: &'graph G, start_node: G::Node) -> Self {
  72         Self { graph, stack: vec![start_node], visited: BitSet::new_empty(graph.num_nodes()) }
  73     }
  74 }
  75
  76 impl<G> Iterator for DepthFirstSearch<'_, G>
  77 where
  78     G: ?Sized + DirectedGraph + WithNumNodes + WithSuccessors,
  79 {
  80     type Item = G::Node;
  81
  82     fn next(&mut self) -> Option<G::Node> {
  83         let DepthFirstSearch { stack, visited, graph } = self;
  84         let n = stack.pop()?;
  85         stack.extend(graph.successors(n).filter(|&m| visited.insert(m)));
  86         Some(n)
  87     }
  88 }
  89
  90 /// The status of a node in the depth-first search.
  91 ///
  92 /// See the documentation of `TriColorDepthFirstSearch` to see how a node's status is updated
  93 /// during DFS.
  94 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
  95 pub enum NodeStatus {
  96     /// This node has been examined by the depth-first search but is not yet `Settled`.
  97     ///
  98     /// Also referred to as "gray" or "discovered" nodes in [CLR].
  99     ///
 100     /// [CLR]: https://en.wikipedia.org/wiki/Introduction_to_Algorithms
 101     Visited,
 102
 103     /// This node and all nodes reachable from it have been examined by the depth-first search.
 104     ///
 105     /// Also referred to as "black" or "finished" nodes in [CLR].
 106     ///
 107     /// [CLR]: https://en.wikipedia.org/wiki/Introduction_to_Algorithms
 108     Settled,
 109 }
 110
 111 struct Event<N> {
 112     node: N,
 113     becomes: NodeStatus,
 114 }
 115
 116 /// A depth-first search that also tracks when all successors of a node have been examined.
 117 ///
 118 /// This is based on the DFS described in [Introduction to Algorithms (1st ed.)][CLR], hereby
 119 /// referred to as **CLR**. However, we use the terminology in [`NodeStatus`] above instead of
 120 /// "discovered"/"finished" or "white"/"grey"/"black". Each node begins the search with no status,
 121 /// becomes `Visited` when it is first examined by the DFS and is `Settled` when all nodes
 122 /// reachable from it have been examined. This allows us to differentiate between "tree", "back"
 123 /// and "forward" edges (see [`TriColorVisitor::node_examined`]).
 124 ///
 125 /// Unlike the pseudocode in [CLR], this implementation is iterative and does not use timestamps.
 126 /// We accomplish this by storing `Event`s on the stack that result in a (possible) state change
 127 /// for each node. A `Visited` event signifies that we should examine this node if it has not yet
 128 /// been `Visited` or `Settled`. When a node is examined for the first time, we mark it as
 129 /// `Visited` and push a `Settled` event for it on stack followed by `Visited` events for all of
 130 /// its predecessors, scheduling them for examination. Multiple `Visited` events for a single node
 131 /// may exist on the stack simultaneously if a node has multiple predecessors, but only one
 132 /// `Settled` event will ever be created for each node. After all `Visited` events for a node's
 133 /// successors have been popped off the stack (as well as any new events triggered by visiting
 134 /// those successors), we will pop off that node's `Settled` event.
 135 ///
 136 /// [CLR]: https://en.wikipedia.org/wiki/Introduction_to_Algorithms
 137 /// [`NodeStatus`]: ./enum.NodeStatus.html
 138 /// [`TriColorVisitor::node_examined`]: ./trait.TriColorVisitor.html#method.node_examined
 139 pub struct TriColorDepthFirstSearch<'graph, G>
 140 where
 141     G: ?Sized + DirectedGraph + WithNumNodes + WithSuccessors,
 142 {
 143     graph: &'graph G,
 144     stack: Vec<Event<G::Node>>,
 145     visited: BitSet<G::Node>,
 146     settled: BitSet<G::Node>,
 147 }
 148
 149 impl<G> TriColorDepthFirstSearch<'graph, G>
 150 where
 151     G: ?Sized + DirectedGraph + WithNumNodes + WithSuccessors,
 152 {
 153     pub fn new(graph: &'graph G) -> Self {
 154         TriColorDepthFirstSearch {
 155             graph,
 156             stack: vec![],
 157             visited: BitSet::new_empty(graph.num_nodes()),
 158             settled: BitSet::new_empty(graph.num_nodes()),
 159         }
 160     }
 161
 162     /// Performs a depth-first search, starting from the given `root`.
 163     ///
 164     /// This won't visit nodes that are not reachable from `root`.
 165     pub fn run_from<V>(mut self, root: G::Node, visitor: &mut V) -> Option<V::BreakVal>
 166     where
 167         V: TriColorVisitor<G>,
 168     {
 169         use NodeStatus::{Settled, Visited};
 170
 171         self.stack.push(Event { node: root, becomes: Visited });
 172
 173         loop {
 174             match self.stack.pop()? {
 175                 Event { node, becomes: Settled } => {
 176                     let not_previously_settled = self.settled.insert(node);
 177                     assert!(not_previously_settled, "A node should be settled exactly once");
 178                     if let ControlFlow::Break(val) = visitor.node_settled(node) {
 179                         return Some(val);
 180                     }
 181                 }
 182
 183                 Event { node, becomes: Visited } => {
 184                     let not_previously_visited = self.visited.insert(node);
 185                     let prior_status = if not_previously_visited {
 186                         None
 187                     } else if self.settled.contains(node) {
 188                         Some(Settled)
 189                     } else {
 190                         Some(Visited)
 191                     };
 192
 193                     if let ControlFlow::Break(val) = visitor.node_examined(node, prior_status) {
 194                         return Some(val);
 195                     }
 196
 197                     // If this node has already been examined, we are done.
 198                     if prior_status.is_some() {
 199                         continue;
 200                     }
 201
 202                     // Otherwise, push a `Settled` event for this node onto the stack, then
 203                     // schedule its successors for examination.
 204                     self.stack.push(Event { node, becomes: Settled });
 205                     for succ in self.graph.successors(node) {
 206                         if !visitor.ignore_edge(node, succ) {
 207                             self.stack.push(Event { node: succ, becomes: Visited });
 208                         }
 209                     }
 210                 }
 211             }
 212         }
 213     }
 214 }
 215
 216 impl<G> TriColorDepthFirstSearch<'graph, G>
 217 where
 218     G: ?Sized + DirectedGraph + WithNumNodes + WithSuccessors + WithStartNode,
 219 {
 220     /// Performs a depth-first search, starting from `G::start_node()`.
 221     ///
 222     /// This won't visit nodes that are not reachable from the start node.
 223     pub fn run_from_start<V>(self, visitor: &mut V) -> Option<V::BreakVal>
 224     where
 225         V: TriColorVisitor<G>,
 226     {
 227         let root = self.graph.start_node();
 228         self.run_from(root, visitor)
 229     }
 230 }
 231
 232 /// What to do when a node is examined or becomes `Settled` during DFS.
 233 pub trait TriColorVisitor<G>
 234 where
 235     G: ?Sized + DirectedGraph,
 236 {
 237     /// The value returned by this search.
 238     type BreakVal;
 239
 240     /// Called when a node is examined by the depth-first search.
 241     ///
 242     /// By checking the value of `prior_status`, this visitor can determine whether the edge
 243     /// leading to this node was a tree edge (`None`), forward edge (`Some(Settled)`) or back edge
 244     /// (`Some(Visited)`). For a full explanation of each edge type, see the "Depth-first Search"
 245     /// chapter in [CLR] or [wikipedia].
 246     ///
 247     /// If you want to know *both* nodes linked by each edge, you'll need to modify
 248     /// `TriColorDepthFirstSearch` to store a `source` node for each `Visited` event.
 249     ///
 250     /// [wikipedia]: https://en.wikipedia.org/wiki/Depth-first_search#Output_of_a_depth-first_search
 251     /// [CLR]: https://en.wikipedia.org/wiki/Introduction_to_Algorithms
 252     fn node_examined(
 253         &mut self,
 254         _node: G::Node,
 255         _prior_status: Option<NodeStatus>,
 256     ) -> ControlFlow<Self::BreakVal> {
 257         ControlFlow::CONTINUE
 258     }
 259
 260     /// Called after all nodes reachable from this one have been examined.
 261     fn node_settled(&mut self, _node: G::Node) -> ControlFlow<Self::BreakVal> {
 262         ControlFlow::CONTINUE
 263     }
 264
 265     /// Behave as if no edges exist from `source` to `target`.
 266     fn ignore_edge(&mut self, _source: G::Node, _target: G::Node) -> bool {
 267         false
 268     }
 269 }
 270
 271 /// This `TriColorVisitor` looks for back edges in a graph, which indicate that a cycle exists.
 272 pub struct CycleDetector;
 273
 274 impl<G> TriColorVisitor<G> for CycleDetector
 275 where
 276     G: ?Sized + DirectedGraph,
 277 {
 278     type BreakVal = ();
 279
 280     fn node_examined(
 281         &mut self,
 282         _node: G::Node,
 283         prior_status: Option<NodeStatus>,
 284     ) -> ControlFlow<Self::BreakVal> {
 285         match prior_status {
 286             Some(NodeStatus::Visited) => ControlFlow::BREAK,
 287             _ => ControlFlow::CONTINUE,
 288         }
 289     }
 290 }