wasmtime_environ/graphs/
dfs.rs

1use super::*;
2
3/// An iterative depth-first traversal.
4pub struct Dfs<Node> {
5    stack: Vec<DfsEvent<Node>>,
6}
7
8impl<Node> Default for Dfs<Node> {
9    fn default() -> Self {
10        Self {
11            stack: Default::default(),
12        }
13    }
14}
15
16impl<Node> Dfs<Node> {
17    /// Create a new DFS traversal, starting at the given roots.
18    pub fn new(roots: impl IntoIterator<Item = Node>) -> Self {
19        let mut dfs = Self::default();
20        dfs.add_roots(roots);
21        dfs
22    }
23
24    /// Add a single new root to this traversal, to be visited immediately.
25    pub fn add_root(&mut self, root: Node) {
26        self.stack.push(DfsEvent::Pre(root));
27    }
28
29    /// Add multiple new roots to this traversal, to be visited immediately.
30    pub fn add_roots(&mut self, roots: impl IntoIterator<Item = Node>) {
31        self.stack
32            .extend(roots.into_iter().map(|v| DfsEvent::Pre(v)));
33    }
34}
35
36/// An event during a DFS traversal.
37#[derive(Clone, Copy, Debug, PartialEq, Eq)]
38pub enum DfsEvent<Node> {
39    /// The first time seeing this node.
40    Pre(Node),
41
42    /// After having just visited the given edge.
43    AfterEdge(Node, Node),
44
45    /// Finished visiting this node and all of its successors.
46    Post(Node),
47}
48
49impl<Node> Dfs<Node>
50where
51    Node: Copy,
52{
53    /// Pump the traversal, yielding the next `DfsEvent`.
54    ///
55    /// Returns `None` when the traversal is complete.
56    pub fn next<G>(&mut self, graph: G, seen: impl Fn(Node) -> bool) -> Option<DfsEvent<Node>>
57    where
58        G: Graph<Node>,
59    {
60        loop {
61            let event = self.stack.pop()?;
62
63            if let DfsEvent::Pre(node) = event {
64                if seen(node) {
65                    continue;
66                }
67
68                let successors = graph.successors(node);
69
70                let (min, max) = successors.size_hint();
71                let estimated_successors_len = max.unwrap_or_else(|| 2 * min);
72                self.stack.reserve(
73                    // We push an after-edge and pre event for each successor.
74                    2 * estimated_successors_len
75                        // And we push one post event for this node.
76                        + 1,
77                );
78
79                self.stack.push(DfsEvent::Post(node));
80                for succ in successors {
81                    self.stack.push(DfsEvent::AfterEdge(node, succ));
82                    if !seen(succ) {
83                        self.stack.push(DfsEvent::Pre(succ));
84                    }
85                }
86            }
87
88            return Some(event);
89        }
90    }
91}
wasmtime_environ/graphs/dfs.rs

wasmtime_environ/graphs/
dfs.rs
