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) -> Self {
  87         Self { graph, stack: vec![], visited: BitSet::new_empty(graph.num_nodes()) }
  88     }
  89
  90     /// Version of `push_start_node` that is convenient for chained
  91     /// use.
  92     pub fn with_start_node(mut self, start_node: G::Node) -> Self {
  93         self.push_start_node(start_node);
  94         self
  95     }
  96
  97     /// Pushes another start node onto the stack. If the node
  98     /// has not already been visited, then you will be able to
  99     /// walk its successors (and so forth) after the current
 100     /// contents of the stack are drained. If multiple start nodes
 101     /// are added into the walk, then their mutual successors
 102     /// will all be walked. You can use this method once the
 103     /// iterator has been completely drained to add additional
 104     /// start nodes.
 105     pub fn push_start_node(&mut self, start_node: G::Node) {
 106         if self.visited.insert(start_node) {
 107             self.stack.push(start_node);
 108         }
 109     }
 110
 111     /// Searches all nodes reachable from the current start nodes.
 112     /// This is equivalent to just invoke `next` repeatedly until
 113     /// you get a `None` result.
 114     pub fn complete_search(&mut self) {
 115         while let Some(_) = self.next() {}
 116     }
 117
 118     /// Returns true if node has been visited thus far.
 119     /// A node is considered "visited" once it is pushed
 120     /// onto the internal stack; it may not yet have been yielded
 121     /// from the iterator. This method is best used after
 122     /// the iterator is completely drained.
 123     pub fn visited(&self, node: G::Node) -> bool {
 124         self.visited.contains(node)
 125     }
 126 }
 127
 128 impl<G> std::fmt::Debug for DepthFirstSearch<'_, G>
 129 where
 130     G: ?Sized + DirectedGraph + WithNumNodes + WithSuccessors,
 131 {
 132     fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
 133         let mut f = fmt.debug_set();
 134         for n in self.visited.iter() {
 135             f.entry(&n);
 136         }
 137         f.finish()
 138     }
 139 }
 140
 141 impl<G> Iterator for DepthFirstSearch<'_, G>
 142 where
 143     G: ?Sized + DirectedGraph + WithNumNodes + WithSuccessors,
 144 {
 145     type Item = G::Node;
 146
 147     fn next(&mut self) -> Option<G::Node> {
 148         let DepthFirstSearch { stack, visited, graph } = self;
 149         let n = stack.pop()?;
 150         stack.extend(graph.successors(n).filter(|&m| visited.insert(m)));
 151         Some(n)
 152     }
 153 }
 154
 155 /// The status of a node in the depth-first search.
 156 ///
 157 /// See the documentation of `TriColorDepthFirstSearch` to see how a node's status is updated
 158 /// during DFS.
 159 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 160 pub enum NodeStatus {
 161     /// This node has been examined by the depth-first search but is not yet `Settled`.
 162     ///
 163     /// Also referred to as "gray" or "discovered" nodes in [CLR].
 164     ///
 165     /// [CLR]: https://en.wikipedia.org/wiki/Introduction_to_Algorithms
 166     Visited,
 167
 168     /// This node and all nodes reachable from it have been examined by the depth-first search.
 169     ///
 170     /// Also referred to as "black" or "finished" nodes in [CLR].
 171     ///
 172     /// [CLR]: https://en.wikipedia.org/wiki/Introduction_to_Algorithms
 173     Settled,
 174 }
 175
 176 struct Event<N> {
 177     node: N,
 178     becomes: NodeStatus,
 179 }
 180
 181 /// A depth-first search that also tracks when all successors of a node have been examined.
 182 ///
 183 /// This is based on the DFS described in [Introduction to Algorithms (1st ed.)][CLR], hereby
 184 /// referred to as **CLR**. However, we use the terminology in [`NodeStatus`] above instead of
 185 /// "discovered"/"finished" or "white"/"grey"/"black". Each node begins the search with no status,
 186 /// becomes `Visited` when it is first examined by the DFS and is `Settled` when all nodes
 187 /// reachable from it have been examined. This allows us to differentiate between "tree", "back"
 188 /// and "forward" edges (see [`TriColorVisitor::node_examined`]).
 189 ///
 190 /// Unlike the pseudocode in [CLR], this implementation is iterative and does not use timestamps.
 191 /// We accomplish this by storing `Event`s on the stack that result in a (possible) state change
 192 /// for each node. A `Visited` event signifies that we should examine this node if it has not yet
 193 /// been `Visited` or `Settled`. When a node is examined for the first time, we mark it as
 194 /// `Visited` and push a `Settled` event for it on stack followed by `Visited` events for all of
 195 /// its predecessors, scheduling them for examination. Multiple `Visited` events for a single node
 196 /// may exist on the stack simultaneously if a node has multiple predecessors, but only one
 197 /// `Settled` event will ever be created for each node. After all `Visited` events for a node's
 198 /// successors have been popped off the stack (as well as any new events triggered by visiting
 199 /// those successors), we will pop off that node's `Settled` event.
 200 ///
 201 /// [CLR]: https://en.wikipedia.org/wiki/Introduction_to_Algorithms
 202 pub struct TriColorDepthFirstSearch<'graph, G>
 203 where
 204     G: ?Sized + DirectedGraph + WithNumNodes + WithSuccessors,
 205 {
 206     graph: &'graph G,
 207     stack: Vec<Event<G::Node>>,
 208     visited: BitSet<G::Node>,
 209     settled: BitSet<G::Node>,
 210 }
 211
 212 impl<G> TriColorDepthFirstSearch<'graph, G>
 213 where
 214     G: ?Sized + DirectedGraph + WithNumNodes + WithSuccessors,
 215 {
 216     pub fn new(graph: &'graph G) -> Self {
 217         TriColorDepthFirstSearch {
 218             graph,
 219             stack: vec![],
 220             visited: BitSet::new_empty(graph.num_nodes()),
 221             settled: BitSet::new_empty(graph.num_nodes()),
 222         }
 223     }
 224
 225     /// Performs a depth-first search, starting from the given `root`.
 226     ///
 227     /// This won't visit nodes that are not reachable from `root`.
 228     pub fn run_from<V>(mut self, root: G::Node, visitor: &mut V) -> Option<V::BreakVal>
 229     where
 230         V: TriColorVisitor<G>,
 231     {
 232         use NodeStatus::{Settled, Visited};
 233
 234         self.stack.push(Event { node: root, becomes: Visited });
 235
 236         loop {
 237             match self.stack.pop()? {
 238                 Event { node, becomes: Settled } => {
 239                     let not_previously_settled = self.settled.insert(node);
 240                     assert!(not_previously_settled, "A node should be settled exactly once");
 241                     if let ControlFlow::Break(val) = visitor.node_settled(node) {
 242                         return Some(val);
 243                     }
 244                 }
 245
 246                 Event { node, becomes: Visited } => {
 247                     let not_previously_visited = self.visited.insert(node);
 248                     let prior_status = if not_previously_visited {
 249                         None
 250                     } else if self.settled.contains(node) {
 251                         Some(Settled)
 252                     } else {
 253                         Some(Visited)
 254                     };
 255
 256                     if let ControlFlow::Break(val) = visitor.node_examined(node, prior_status) {
 257                         return Some(val);
 258                     }
 259
 260                     // If this node has already been examined, we are done.
 261                     if prior_status.is_some() {
 262                         continue;
 263                     }
 264
 265                     // Otherwise, push a `Settled` event for this node onto the stack, then
 266                     // schedule its successors for examination.
 267                     self.stack.push(Event { node, becomes: Settled });
 268                     for succ in self.graph.successors(node) {
 269                         if !visitor.ignore_edge(node, succ) {
 270                             self.stack.push(Event { node: succ, becomes: Visited });
 271                         }
 272                     }
 273                 }
 274             }
 275         }
 276     }
 277 }
 278
 279 impl<G> TriColorDepthFirstSearch<'graph, G>
 280 where
 281     G: ?Sized + DirectedGraph + WithNumNodes + WithSuccessors + WithStartNode,
 282 {
 283     /// Performs a depth-first search, starting from `G::start_node()`.
 284     ///
 285     /// This won't visit nodes that are not reachable from the start node.
 286     pub fn run_from_start<V>(self, visitor: &mut V) -> Option<V::BreakVal>
 287     where
 288         V: TriColorVisitor<G>,
 289     {
 290         let root = self.graph.start_node();
 291         self.run_from(root, visitor)
 292     }
 293 }
 294
 295 /// What to do when a node is examined or becomes `Settled` during DFS.
 296 pub trait TriColorVisitor<G>
 297 where
 298     G: ?Sized + DirectedGraph,
 299 {
 300     /// The value returned by this search.
 301     type BreakVal;
 302
 303     /// Called when a node is examined by the depth-first search.
 304     ///
 305     /// By checking the value of `prior_status`, this visitor can determine whether the edge
 306     /// leading to this node was a tree edge (`None`), forward edge (`Some(Settled)`) or back edge
 307     /// (`Some(Visited)`). For a full explanation of each edge type, see the "Depth-first Search"
 308     /// chapter in [CLR] or [wikipedia].
 309     ///
 310     /// If you want to know *both* nodes linked by each edge, you'll need to modify
 311     /// `TriColorDepthFirstSearch` to store a `source` node for each `Visited` event.
 312     ///
 313     /// [wikipedia]: https://en.wikipedia.org/wiki/Depth-first_search#Output_of_a_depth-first_search
 314     /// [CLR]: https://en.wikipedia.org/wiki/Introduction_to_Algorithms
 315     fn node_examined(
 316         &mut self,
 317         _node: G::Node,
 318         _prior_status: Option<NodeStatus>,
 319     ) -> ControlFlow<Self::BreakVal> {
 320         ControlFlow::CONTINUE
 321     }
 322
 323     /// Called after all nodes reachable from this one have been examined.
 324     fn node_settled(&mut self, _node: G::Node) -> ControlFlow<Self::BreakVal> {
 325         ControlFlow::CONTINUE
 326     }
 327
 328     /// Behave as if no edges exist from `source` to `target`.
 329     fn ignore_edge(&mut self, _source: G::Node, _target: G::Node) -> bool {
 330         false
 331     }
 332 }
 333
 334 /// This `TriColorVisitor` looks for back edges in a graph, which indicate that a cycle exists.
 335 pub struct CycleDetector;
 336
 337 impl<G> TriColorVisitor<G> for CycleDetector
 338 where
 339     G: ?Sized + DirectedGraph,
 340 {
 341     type BreakVal = ();
 342
 343     fn node_examined(
 344         &mut self,
 345         _node: G::Node,
 346         prior_status: Option<NodeStatus>,
 347     ) -> ControlFlow<Self::BreakVal> {
 348         match prior_status {
 349             Some(NodeStatus::Visited) => ControlFlow::BREAK,
 350             _ => ControlFlow::CONTINUE,
 351         }
 352     }
 353 }