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