# rustc_data_structures/graph/iterate/mod.rs

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355

```
use std::ops::ControlFlow;
use rustc_index::bit_set::BitSet;
use rustc_index::{IndexSlice, IndexVec};
use super::{DirectedGraph, StartNode, Successors};
#[cfg(test)]
mod tests;
pub fn post_order_from<G: DirectedGraph + Successors>(
graph: &G,
start_node: G::Node,
) -> Vec<G::Node> {
post_order_from_to(graph, start_node, None)
}
pub fn post_order_from_to<G: DirectedGraph + Successors>(
graph: &G,
start_node: G::Node,
end_node: Option<G::Node>,
) -> Vec<G::Node> {
let mut visited: IndexVec<G::Node, bool> = IndexVec::from_elem_n(false, graph.num_nodes());
let mut result: Vec<G::Node> = Vec::with_capacity(graph.num_nodes());
if let Some(end_node) = end_node {
visited[end_node] = true;
}
post_order_walk(graph, start_node, &mut result, &mut visited);
result
}
fn post_order_walk<G: DirectedGraph + Successors>(
graph: &G,
node: G::Node,
result: &mut Vec<G::Node>,
visited: &mut IndexSlice<G::Node, bool>,
) {
struct PostOrderFrame<Node, Iter> {
node: Node,
iter: Iter,
}
if visited[node] {
return;
}
let mut stack = vec![PostOrderFrame { node, iter: graph.successors(node) }];
'recurse: while let Some(frame) = stack.last_mut() {
let node = frame.node;
visited[node] = true;
for successor in frame.iter.by_ref() {
if !visited[successor] {
stack.push(PostOrderFrame { node: successor, iter: graph.successors(successor) });
continue 'recurse;
}
}
let _ = stack.pop();
result.push(node);
}
}
pub fn reverse_post_order<G: DirectedGraph + Successors>(
graph: &G,
start_node: G::Node,
) -> Vec<G::Node> {
let mut vec = post_order_from(graph, start_node);
vec.reverse();
vec
}
/// A "depth-first search" iterator for a directed graph.
pub struct DepthFirstSearch<G>
where
G: DirectedGraph + Successors,
{
graph: G,
stack: Vec<G::Node>,
visited: BitSet<G::Node>,
}
impl<G> DepthFirstSearch<G>
where
G: DirectedGraph + Successors,
{
pub fn new(graph: G) -> Self {
Self { stack: vec![], visited: BitSet::new_empty(graph.num_nodes()), graph }
}
/// Version of `push_start_node` that is convenient for chained
/// use.
pub fn with_start_node(mut self, start_node: G::Node) -> Self {
self.push_start_node(start_node);
self
}
/// Pushes another start node onto the stack. If the node
/// has not already been visited, then you will be able to
/// walk its successors (and so forth) after the current
/// contents of the stack are drained. If multiple start nodes
/// are added into the walk, then their mutual successors
/// will all be walked. You can use this method once the
/// iterator has been completely drained to add additional
/// start nodes.
pub fn push_start_node(&mut self, start_node: G::Node) {
if self.visited.insert(start_node) {
self.stack.push(start_node);
}
}
/// Searches all nodes reachable from the current start nodes.
/// This is equivalent to just invoke `next` repeatedly until
/// you get a `None` result.
pub fn complete_search(&mut self) {
for _ in self.by_ref() {}
}
/// Returns true if node has been visited thus far.
/// A node is considered "visited" once it is pushed
/// onto the internal stack; it may not yet have been yielded
/// from the iterator. This method is best used after
/// the iterator is completely drained.
pub fn visited(&self, node: G::Node) -> bool {
self.visited.contains(node)
}
}
impl<G> std::fmt::Debug for DepthFirstSearch<G>
where
G: DirectedGraph + Successors,
{
fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let mut f = fmt.debug_set();
for n in self.visited.iter() {
f.entry(&n);
}
f.finish()
}
}
impl<G> Iterator for DepthFirstSearch<G>
where
G: DirectedGraph + Successors,
{
type Item = G::Node;
fn next(&mut self) -> Option<G::Node> {
let DepthFirstSearch { stack, visited, graph } = self;
let n = stack.pop()?;
stack.extend(graph.successors(n).filter(|&m| visited.insert(m)));
Some(n)
}
}
/// The status of a node in the depth-first search.
///
/// See the documentation of `TriColorDepthFirstSearch` to see how a node's status is updated
/// during DFS.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum NodeStatus {
/// This node has been examined by the depth-first search but is not yet `Settled`.
///
/// Also referred to as "gray" or "discovered" nodes in [CLR].
///
/// [CLR]: https://en.wikipedia.org/wiki/Introduction_to_Algorithms
Visited,
/// This node and all nodes reachable from it have been examined by the depth-first search.
///
/// Also referred to as "black" or "finished" nodes in [CLR].
///
/// [CLR]: https://en.wikipedia.org/wiki/Introduction_to_Algorithms
Settled,
}
struct Event<N> {
node: N,
becomes: NodeStatus,
}
/// A depth-first search that also tracks when all successors of a node have been examined.
///
/// This is based on the DFS described in [Introduction to Algorithms (1st ed.)][CLR], hereby
/// referred to as **CLR**. However, we use the terminology in [`NodeStatus`] above instead of
/// "discovered"/"finished" or "white"/"grey"/"black". Each node begins the search with no status,
/// becomes `Visited` when it is first examined by the DFS and is `Settled` when all nodes
/// reachable from it have been examined. This allows us to differentiate between "tree", "back"
/// and "forward" edges (see [`TriColorVisitor::node_examined`]).
///
/// Unlike the pseudocode in [CLR], this implementation is iterative and does not use timestamps.
/// We accomplish this by storing `Event`s on the stack that result in a (possible) state change
/// for each node. A `Visited` event signifies that we should examine this node if it has not yet
/// been `Visited` or `Settled`. When a node is examined for the first time, we mark it as
/// `Visited` and push a `Settled` event for it on stack followed by `Visited` events for all of
/// its predecessors, scheduling them for examination. Multiple `Visited` events for a single node
/// may exist on the stack simultaneously if a node has multiple predecessors, but only one
/// `Settled` event will ever be created for each node. After all `Visited` events for a node's
/// successors have been popped off the stack (as well as any new events triggered by visiting
/// those successors), we will pop off that node's `Settled` event.
///
/// [CLR]: https://en.wikipedia.org/wiki/Introduction_to_Algorithms
pub struct TriColorDepthFirstSearch<'graph, G>
where
G: ?Sized + DirectedGraph + Successors,
{
graph: &'graph G,
stack: Vec<Event<G::Node>>,
visited: BitSet<G::Node>,
settled: BitSet<G::Node>,
}
impl<'graph, G> TriColorDepthFirstSearch<'graph, G>
where
G: ?Sized + DirectedGraph + Successors,
{
pub fn new(graph: &'graph G) -> Self {
TriColorDepthFirstSearch {
graph,
stack: vec![],
visited: BitSet::new_empty(graph.num_nodes()),
settled: BitSet::new_empty(graph.num_nodes()),
}
}
/// Performs a depth-first search, starting from the given `root`.
///
/// This won't visit nodes that are not reachable from `root`.
pub fn run_from<V>(mut self, root: G::Node, visitor: &mut V) -> Option<V::BreakVal>
where
V: TriColorVisitor<G>,
{
use NodeStatus::{Settled, Visited};
self.stack.push(Event { node: root, becomes: Visited });
loop {
match self.stack.pop()? {
Event { node, becomes: Settled } => {
let not_previously_settled = self.settled.insert(node);
assert!(not_previously_settled, "A node should be settled exactly once");
if let ControlFlow::Break(val) = visitor.node_settled(node) {
return Some(val);
}
}
Event { node, becomes: Visited } => {
let not_previously_visited = self.visited.insert(node);
let prior_status = if not_previously_visited {
None
} else if self.settled.contains(node) {
Some(Settled)
} else {
Some(Visited)
};
if let ControlFlow::Break(val) = visitor.node_examined(node, prior_status) {
return Some(val);
}
// If this node has already been examined, we are done.
if prior_status.is_some() {
continue;
}
// Otherwise, push a `Settled` event for this node onto the stack, then
// schedule its successors for examination.
self.stack.push(Event { node, becomes: Settled });
for succ in self.graph.successors(node) {
if !visitor.ignore_edge(node, succ) {
self.stack.push(Event { node: succ, becomes: Visited });
}
}
}
}
}
}
}
impl<G> TriColorDepthFirstSearch<'_, G>
where
G: ?Sized + DirectedGraph + Successors + StartNode,
{
/// Performs a depth-first search, starting from `G::start_node()`.
///
/// This won't visit nodes that are not reachable from the start node.
pub fn run_from_start<V>(self, visitor: &mut V) -> Option<V::BreakVal>
where
V: TriColorVisitor<G>,
{
let root = self.graph.start_node();
self.run_from(root, visitor)
}
}
/// What to do when a node is examined or becomes `Settled` during DFS.
pub trait TriColorVisitor<G>
where
G: ?Sized + DirectedGraph,
{
/// The value returned by this search.
type BreakVal;
/// Called when a node is examined by the depth-first search.
///
/// By checking the value of `prior_status`, this visitor can determine whether the edge
/// leading to this node was a tree edge (`None`), forward edge (`Some(Settled)`) or back edge
/// (`Some(Visited)`). For a full explanation of each edge type, see the "Depth-first Search"
/// chapter in [CLR] or [wikipedia].
///
/// If you want to know *both* nodes linked by each edge, you'll need to modify
/// `TriColorDepthFirstSearch` to store a `source` node for each `Visited` event.
///
/// [wikipedia]: https://en.wikipedia.org/wiki/Depth-first_search#Output_of_a_depth-first_search
/// [CLR]: https://en.wikipedia.org/wiki/Introduction_to_Algorithms
fn node_examined(
&mut self,
_node: G::Node,
_prior_status: Option<NodeStatus>,
) -> ControlFlow<Self::BreakVal> {
ControlFlow::Continue(())
}
/// Called after all nodes reachable from this one have been examined.
fn node_settled(&mut self, _node: G::Node) -> ControlFlow<Self::BreakVal> {
ControlFlow::Continue(())
}
/// Behave as if no edges exist from `source` to `target`.
fn ignore_edge(&mut self, _source: G::Node, _target: G::Node) -> bool {
false
}
}
/// This `TriColorVisitor` looks for back edges in a graph, which indicate that a cycle exists.
pub struct CycleDetector;
impl<G> TriColorVisitor<G> for CycleDetector
where
G: ?Sized + DirectedGraph,
{
type BreakVal = ();
fn node_examined(
&mut self,
_node: G::Node,
prior_status: Option<NodeStatus>,
) -> ControlFlow<Self::BreakVal> {
match prior_status {
Some(NodeStatus::Visited) => ControlFlow::Break(()),
_ => ControlFlow::Continue(()),
}
}
}
```