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