tsort: use iterative dfs to avoid crashes on stack overflow

Author: Nekrolm

src/uu/tsort/src/tsort.rs

@@ -4,7 +4,8 @@
 // file that was distributed with this source code.
 //spell-checker:ignore TAOCP indegree
 use clap::{Arg, Command};
-use std::collections::{HashMap, HashSet, VecDeque};
+use std::collections::hash_map::Entry;
+use std::collections::{HashMap, VecDeque};
 use std::ffi::OsString;
 use std::path::Path;
 use thiserror::Error;
@@ -34,13 +35,15 @@ enum TsortError {
     /// The graph contains a cycle.
     #[error("{input}: {message}", input = .0, message = translate!("tsort-error-loop"))]
     Loop(String),
-
-    /// A particular node in a cycle. (This is mainly used for printing.)
-    #[error("{0}")]
-    LoopNode(String),
 }
 
+// Auxiliary struct, just for printing loop nodes via show! macro
+#[derive(Debug, Error)]
+#[error("{0}")]
+struct LoopNode<'a>(&'a str);
+
 impl UError for TsortError {}
+impl UError for LoopNode<'_> {}
 
 #[uucore::main]
 pub fn uumain(args: impl uucore::Args) -> UResult<()> {
@@ -131,6 +134,13 @@ struct Graph<'input> {
     nodes: HashMap<&'input str, Node<'input>>,
 }
 
+#[derive(Clone, Copy, PartialEq, Eq)]
+enum VisitedState {
+    JustAdded,
+    Opened,
+    Closed,
+}
+
 impl<'input> Graph<'input> {
     fn new(name: String) -> Graph<'input> {
         Self {
@@ -224,8 +234,8 @@ impl<'input> Graph<'input> {
     fn find_and_break_cycle(&mut self, frontier: &mut VecDeque<&'input str>) {
         let cycle = self.detect_cycle();
         show!(TsortError::Loop(self.name.clone()));
-        for node in &cycle {
-            show!(TsortError::LoopNode((*node).to_string()));
+        for &node in &cycle {
+            show!(LoopNode(node));
         }
         let u = cycle[0];
         let v = cycle[1];
@@ -240,41 +250,68 @@ impl<'input> Graph<'input> {
         let mut nodes: Vec<_> = self.nodes.keys().collect();
         nodes.sort_unstable();
 
-        let mut visited = HashSet::new();
+        let mut visited = HashMap::new();
         let mut stack = Vec::with_capacity(self.nodes.len());
         for node in nodes {
-            if !visited.contains(node) && self.dfs(node, &mut visited, &mut stack) {
+            if self.dfs(node, &mut visited, &mut stack) {
                 return stack;
             }
         }
-        unreachable!();
+        unreachable!("detect_cycle is expected to be called only on graphs with cycles");
     }
 
     fn dfs(
         &self,
         node: &'input str,
-        visited: &mut HashSet<&'input str>,
+        visited: &mut HashMap<&'input str, VisitedState>,
         stack: &mut Vec<&'input str>,
     ) -> bool {
-        if stack.contains(&node) {
-            return true;
-        }
-        if visited.contains(&node) {
-            return false;
-        }
-
-        visited.insert(node);
         stack.push(node);
 
-        if let Some(successor_names) = self.nodes.get(node).map(|n| &n.successor_names) {
-            for &successor in successor_names {
-                if self.dfs(successor, visited, stack) {
-                    return true;
+        while let Some(node) = stack.pop() {
+            let state = visited.entry(node).or_insert(VisitedState::JustAdded);
+
+            match *state {
+                VisitedState::Closed => continue,
+                VisitedState::Opened => {
+                    // close the node if it was opened
+                    *state = VisitedState::Closed;
+                    continue;
+                }
+                VisitedState::JustAdded => {
+                    // open the node, put it on the stack for tracking & to return
+                    // to this node at the end of processing all its successors
+                    *state = VisitedState::Opened;
+                    stack.push(node);
+                }
+            }
+
+            let Some(successors) = self.nodes.get(node).map(|n| &n.successor_names) else {
+                continue;
+            };
+
+            // to preserve same output as for recursive version:
+            // iterate over successors in reverse order -- because all of them are going to be
+            // inserted on the stack, and popped back in reverse order
+            for &successor in successors.iter().rev() {
+                match visited.entry(successor) {
+                    Entry::Vacant(v) => {
+                        v.insert(VisitedState::JustAdded);
+                        stack.push(successor);
+                    }
+                    Entry::Occupied(o) => {
+                        if *o.get() == VisitedState::Opened {
+                            // we are entering the same node --
+                            // we found loop & stack contains this loop
+                            // in form of Opened nodes. Drop all other nodes
+                            stack.retain(|&n| visited.get(n) == Some(&VisitedState::Opened));
+                            return true;
+                        }
+                    }
                 }
             }
         }
 
-        stack.pop();
         false
     }
 }

tests/by-util/test_tsort.rs

@@ -122,3 +122,17 @@ fn test_two_cycles() {
         .stdout_is("a\nc\nd\nb\n")
         .stderr_is("tsort: -: input contains a loop:\ntsort: b\ntsort: c\ntsort: -: input contains a loop:\ntsort: b\ntsort: d\n");
 }
+
+#[test]
+fn test_long_loop_no_stack_overflow() {
+    let mut input = String::new();
+    const N: usize = 100000;
+    for v in 0..N {
+        let next = (v + 1) % N;
+        input.push_str(&format!("{v} {next} "));
+    }
+    new_ucmd!()
+        .pipe_in(input)
+        .fails_with_code(1)
+        .stderr_contains("tsort: -: input contains a loop");
+}