Remove recursion from sccc walking

This allows constructing the sccc for large that visit many nodes before
finding a single cycle of sccc, for example lists. When used to find
dependencies in borrow checking the list case is what occurs in very
long functions.

Author: HeroicKatora

compiler/rustc_data_structures/src/graph/scc/mod.rs

@@ -231,20 +231,30 @@ where
 
         let scc_indices = (0..num_nodes)
             .map(G::Node::new)
-            .map(|node| match this.walk_node(0, node) {
+            .map(|node| match this.start_walk_from(node) {
                 WalkReturn::Complete { scc_index } => scc_index,
-                WalkReturn::Cycle { min_depth } => {
-                    panic!("`walk_node(0, {:?})` returned cycle with depth {:?}", node, min_depth)
-                }
+                WalkReturn::Cycle { min_depth } => panic!(
+                    "`start_walk_node({:?})` returned cycle with depth {:?}",
+                    node, min_depth
+                ),
             })
             .collect();
 
         Sccs { scc_indices, scc_data: this.scc_data }
     }
 
-    /// Visits a node during the DFS. We first examine its current
-    /// state -- if it is not yet visited (`NotVisited`), we can push
-    /// it onto the stack and start walking its successors.
+    fn start_walk_from(&mut self, node: G::Node) -> WalkReturn<S> {
+        if let Some(result) = self.inspect_node(node) {
+            result
+        } else {
+            self.walk_unvisited_node(node)
+        }
+    }
+
+    /// Inspect a node during the DFS. We first examine its current
+    /// state -- if it is not yet visited (`NotVisited`), return `None` so
+    /// that the caller might push it onto the stack and start walking its
+    /// successors.
     ///
     /// If it is already on the DFS stack it will be in the state
     /// `BeingVisited`. In that case, we have found a cycle and we
@@ -253,20 +263,19 @@ where
     /// Otherwise, we are looking at a node that has already been
     /// completely visited. We therefore return `WalkReturn::Complete`
     /// with its associated SCC index.
-    fn walk_node(&mut self, depth: usize, node: G::Node) -> WalkReturn<S> {
-        debug!("walk_node(depth = {:?}, node = {:?})", depth, node);
-        match self.find_state(node) {
+    fn inspect_node(&mut self, node: G::Node) -> Option<WalkReturn<S>> {
+        Some(match self.find_state(node) {
             NodeState::InCycle { scc_index } => WalkReturn::Complete { scc_index },
 
             NodeState::BeingVisited { depth: min_depth } => WalkReturn::Cycle { min_depth },
 
-            NodeState::NotVisited => self.walk_unvisited_node(depth, node),
+            NodeState::NotVisited => return None,
 
             NodeState::InCycleWith { parent } => panic!(
                 "`find_state` returned `InCycleWith({:?})`, which ought to be impossible",
                 parent
             ),
-        }
+        })
     }
 
     /// Fetches the state of the node `r`. If `r` is recorded as being
@@ -392,74 +401,174 @@ where
     }
 
     /// Walks a node that has never been visited before.
-    fn walk_unvisited_node(&mut self, depth: usize, node: G::Node) -> WalkReturn<S> {
-        debug!("walk_unvisited_node(depth = {:?}, node = {:?})", depth, node);
-
-        debug_assert!(matches!(self.node_states[node], NodeState::NotVisited));
-
-        // Push `node` onto the stack.
-        self.node_states[node] = NodeState::BeingVisited { depth };
-        self.node_stack.push(node);
-
-        // Walk each successor of the node, looking to see if any of
-        // them can reach a node that is presently on the stack. If
-        // so, that means they can also reach us.
-        let mut min_depth = depth;
-        let mut min_cycle_root = node;
-        let successors_len = self.successors_stack.len();
-        for successor_node in self.graph.successors(node) {
-            debug!("walk_unvisited_node: node = {:?} successor_ode = {:?}", node, successor_node);
-            match self.walk_node(depth + 1, successor_node) {
-                WalkReturn::Cycle { min_depth: successor_min_depth } => {
-                    // Track the minimum depth we can reach.
-                    assert!(successor_min_depth <= depth);
-                    if successor_min_depth < min_depth {
+    ///
+    /// Call this method when `inspect_node` has returned `None`. Having the
+    /// caller decide avoids mutual recursion between the two methods and allows
+    /// us to maintain an allocated stack for nodes on the path between calls.
+    fn walk_unvisited_node(&mut self, initial: G::Node) -> WalkReturn<S> {
+        struct VisitingNodeFrame<G: DirectedGraph, Successors> {
+            node: G::Node,
+            iter: Option<Successors>,
+            depth: usize,
+            min_depth: usize,
+            successors_len: usize,
+            min_cycle_root: G::Node,
+            successor_node: G::Node,
+        }
+
+        // Move the stack to a local variable. We want to utilize the existing allocation and
+        // mutably borrow it without borrowing self at the same time.
+        let mut successors_stack = core::mem::take(&mut self.successors_stack);
+        debug_assert_eq!(successors_stack.len(), 0);
+
+        let mut stack: Vec<VisitingNodeFrame<G, _>> = vec![VisitingNodeFrame {
+            node: initial,
+            depth: 0,
+            min_depth: 0,
+            iter: None,
+            successors_len: 0,
+            min_cycle_root: initial,
+            successor_node: initial,
+        }];
+
+        let mut return_value = None;
+
+        'recurse: while let Some(frame) = stack.last_mut() {
+            let VisitingNodeFrame {
+                node,
+                depth,
+                iter,
+                successors_len,
+                min_depth,
+                min_cycle_root,
+                successor_node,
+            } = frame;
+
+            let node = *node;
+            let depth = *depth;
+
+            let successors = match iter {
+                Some(iter) => iter,
+                None => {
+                    // This None marks that we still have the initialize this node's frame.
+                    debug!("walk_unvisited_node(depth = {:?}, node = {:?})", depth, node);
+
+                    debug_assert!(matches!(self.node_states[node], NodeState::NotVisited));
+
+                    // Push `node` onto the stack.
+                    self.node_states[node] = NodeState::BeingVisited { depth };
+                    self.node_stack.push(node);
+
+                    // Walk each successor of the node, looking to see if any of
+                    // them can reach a node that is presently on the stack. If
+                    // so, that means they can also reach us.
+                    *successors_len = successors_stack.len();
+                    // Set and return a reference, this is currently empty.
+                    iter.get_or_insert(self.graph.successors(node))
+                }
+            };
+
+            // Now that iter is initialized, this is a constant for this frame.
+            let successors_len = *successors_len;
+
+            // Construct iterators for the nodes and walk results. There are two cases:
+            // * The walk of a successor node returned.
+            // * The remaining successor nodes.
+            let returned_walk =
+                return_value.take().into_iter().map(|walk| (*successor_node, Some(walk)));
+
+            let successor_walk = successors.by_ref().map(|successor_node| {
+                debug!(
+                    "walk_unvisited_node: node = {:?} successor_ode = {:?}",
+                    node, successor_node
+                );
+                (successor_node, self.inspect_node(successor_node))
+            });
+
+            for (successor_node, walk) in returned_walk.chain(successor_walk) {
+                match walk {
+                    Some(WalkReturn::Cycle { min_depth: successor_min_depth }) => {
+                        // Track the minimum depth we can reach.
+                        assert!(successor_min_depth <= depth);
+                        if successor_min_depth < *min_depth {
+                            debug!(
+                                "walk_unvisited_node: node = {:?} successor_min_depth = {:?}",
+                                node, successor_min_depth
+                            );
+                            *min_depth = successor_min_depth;
+                            *min_cycle_root = successor_node;
+                        }
+                    }
+
+                    Some(WalkReturn::Complete { scc_index: successor_scc_index }) => {
+                        // Push the completed SCC indices onto
+                        // the `successors_stack` for later.
                         debug!(
-                            "walk_unvisited_node: node = {:?} successor_min_depth = {:?}",
-                            node, successor_min_depth
+                            "walk_unvisited_node: node = {:?} successor_scc_index = {:?}",
+                            node, successor_scc_index
                         );
-                        min_depth = successor_min_depth;
-                        min_cycle_root = successor_node;
+                        successors_stack.push(successor_scc_index);
                     }
-                }
 
-                WalkReturn::Complete { scc_index: successor_scc_index } => {
-                    // Push the completed SCC indices onto
-                    // the `successors_stack` for later.
-                    debug!(
-                        "walk_unvisited_node: node = {:?} successor_scc_index = {:?}",
-                        node, successor_scc_index
-                    );
-                    self.successors_stack.push(successor_scc_index);
+                    None => {
+                        let depth = depth + 1;
+                        debug!("walk_node(depth = {:?}, node = {:?})", depth, successor_node);
+                        // Remember which node the return value will come from.
+                        frame.successor_node = successor_node;
+                        // Start a new stack frame the step into it.
+                        stack.push(VisitingNodeFrame {
+                            node: successor_node,
+                            depth,
+                            iter: None,
+                            successors_len: 0,
+                            min_depth: depth,
+                            min_cycle_root: successor_node,
+                            successor_node: successor_node,
+                        });
+                        continue 'recurse;
+                    }
                 }
             }
-        }
 
-        // Completed walk, remove `node` from the stack.
-        let r = self.node_stack.pop();
-        debug_assert_eq!(r, Some(node));
-
-        // If `min_depth == depth`, then we are the root of the
-        // cycle: we can't reach anyone further down the stack.
-        if min_depth == depth {
-            // Note that successor stack may have duplicates, so we
-            // want to remove those:
-            let deduplicated_successors = {
-                let duplicate_set = &mut self.duplicate_set;
-                duplicate_set.clear();
-                self.successors_stack
-                    .drain(successors_len..)
-                    .filter(move |&i| duplicate_set.insert(i))
-            };
-            let scc_index = self.scc_data.create_scc(deduplicated_successors);
-            self.node_states[node] = NodeState::InCycle { scc_index };
-            WalkReturn::Complete { scc_index }
-        } else {
-            // We are not the head of the cycle. Return back to our
-            // caller. They will take ownership of the
-            // `self.successors` data that we pushed.
-            self.node_states[node] = NodeState::InCycleWith { parent: min_cycle_root };
-            WalkReturn::Cycle { min_depth }
+            // Completed walk, remove `node` from the stack.
+            let r = self.node_stack.pop();
+            debug_assert_eq!(r, Some(node));
+
+            // Remove the frame, it's done.
+            let frame = stack.pop().unwrap();
+
+            // If `min_depth == depth`, then we are the root of the
+            // cycle: we can't reach anyone further down the stack.
+
+            // Pass the 'return value' down the stack.
+            // We return one frame at a time so there can't be another return value.
+            debug_assert!(return_value.is_none());
+            return_value = Some(if frame.min_depth == depth {
+                // Note that successor stack may have duplicates, so we
+                // want to remove those:
+                let deduplicated_successors = {
+                    let duplicate_set = &mut self.duplicate_set;
+                    duplicate_set.clear();
+                    successors_stack
+                        .drain(successors_len..)
+                        .filter(move |&i| duplicate_set.insert(i))
+                };
+                let scc_index = self.scc_data.create_scc(deduplicated_successors);
+                self.node_states[node] = NodeState::InCycle { scc_index };
+                WalkReturn::Complete { scc_index }
+            } else {
+                // We are not the head of the cycle. Return back to our
+                // caller. They will take ownership of the
+                // `self.successors` data that we pushed.
+                self.node_states[node] = NodeState::InCycleWith { parent: frame.min_cycle_root };
+                WalkReturn::Cycle { min_depth: frame.min_depth }
+            });
         }
+
+        // Keep the allocation we used for successors_stack.
+        self.successors_stack = successors_stack;
+        debug_assert_eq!(self.successors_stack.len(), 0);
+
+        return_value.unwrap()
     }
 }