tsort proposal

src/uu/tsort/src/tsort.rs

@@ -3,7 +3,7 @@
 // For the full copyright and license information, please view the LICENSE
 // file that was distributed with this source code.
 use clap::{crate_version, Arg, Command};
-use std::collections::{BTreeMap, BTreeSet};
+use std::collections::{BTreeMap, HashSet, VecDeque};
 use std::fs::File;
 use std::io::{stdin, BufReader, Read};
 use std::path::Path;
@@ -100,10 +100,28 @@ pub fn uu_app() -> Command {
 
 // We use String as a representation of node here
 // but using integer may improve performance.
+// Representation of a node in the graph
+struct Node<'input> {
+    successors: Vec<&'input str>,
+    predecessors_count: usize,
+}
+
+impl<'input> Node<'input> {
+    fn new() -> Self {
+        Node {
+            successors: Vec::new(),
+            predecessors_count: 0,
+        }
+    }
+
+    fn add_successor(&mut self, successor: &'input str) {
+        self.successors.push(successor);
+    }
+}
+
 #[derive(Default)]
 struct Graph<'input> {
-    in_edges: BTreeMap<&'input str, BTreeSet<&'input str>>,
-    out_edges: BTreeMap<&'input str, Vec<&'input str>>,
+    nodes: BTreeMap<&'input str, Node<'input>>,
     result: Vec<&'input str>,
 }
 
@@ -112,65 +130,98 @@ impl<'input> Graph<'input> {
         Self::default()
     }
 
-    fn has_node(&self, n: &str) -> bool {
-        self.in_edges.contains_key(n)
+    fn add_node(&mut self, name: &'input str) {
+        self.nodes.entry(name).or_insert_with(Node::new);
     }
 
-    fn has_edge(&self, from: &str, to: &str) -> bool {
-        self.in_edges[to].contains(from)
+    fn add_edge(&mut self, from: &'input str, to: &'input str) {
+        self.add_node(from);
+        if from != to {
+            self.add_node(to);
+            {
+                let from_node = self.nodes.get_mut(from).unwrap();
+                from_node.add_successor(to);
+            }
+            {
+                let to_node = self.nodes.get_mut(to).unwrap();
+                to_node.predecessors_count += 1;
+            }
+        }
     }
 
-    fn init_node(&mut self, n: &'input str) {
-        self.in_edges.insert(n, BTreeSet::new());
-        self.out_edges.insert(n, vec![]);
-    }
+    fn run_tsort(&mut self) {
+        let mut zeros: VecDeque<&'input str> = self
+            .nodes
+            .iter()
+            .filter_map(|(&name, node)| {
+                if node.predecessors_count == 0 {
+                    Some(name)
+                } else {
+                    None
+                }
+            })
+            .collect();
 
-    fn add_edge(&mut self, from: &'input str, to: &'input str) {
-        if !self.has_node(to) {
-            self.init_node(to);
-        }
+        while let Some(n) = zeros.pop_front() {
+            let node = self.nodes.get_mut(n).unwrap();
+            self.result.push(n);
 
-        if !self.has_node(from) {
-            self.init_node(from);
+            let successors: Vec<&'input str> = node.successors.clone();
+            for successor in successors {
+                let successor_node = self.nodes.get_mut(successor).unwrap();
+                successor_node.predecessors_count -= 1;
+                if successor_node.predecessors_count == 0 {
+                    zeros.push_back(successor);
+                }
+            }
         }
 
-        if from != to && !self.has_edge(from, to) {
-            self.in_edges.get_mut(to).unwrap().insert(from);
-            self.out_edges.get_mut(from).unwrap().push(to);
+        if !self.is_acyclic() {
+            self.result = self.detect_cycle().unwrap();
         }
     }
-
-    // Kahn's algorithm
-    // O(|V|+|E|)
-    fn run_tsort(&mut self) {
-        let mut start_nodes = vec![];
-        for (n, edges) in &self.in_edges {
-            if edges.is_empty() {
-                start_nodes.push(*n);
+    fn detect_cycle(&self) -> Option<Vec<&'input str>> {
+        let mut visited = HashSet::new();
+        let mut stack = Vec::new();
+        for &node in self.nodes.keys() {
+            if !visited.contains(node) && self.dfs(node, &mut visited, &mut stack) {
+                return Some(stack);
             }
         }
+        None
+    }
 
-        while !start_nodes.is_empty() {
-            let n = start_nodes.remove(0);
-
-            self.result.push(n);
+    fn dfs(
+        &self,
+        node: &'input str,
+        visited: &mut HashSet<&'input str>,
+        stack: &mut Vec<&'input str>,
+    ) -> bool {
+        if stack.contains(&node) {
+            return true;
+        }
+        if visited.contains(&node) {
+            return false;
+        }
 
-            let n_out_edges = self.out_edges.get_mut(&n).unwrap();
-            #[allow(clippy::explicit_iter_loop)]
-            for m in n_out_edges.iter() {
-                let m_in_edges = self.in_edges.get_mut(m).unwrap();
-                m_in_edges.remove(&n);
+        visited.insert(node);
+        stack.push(node);
 
-                // If m doesn't have other in-coming edges add it to start_nodes
-                if m_in_edges.is_empty() {
-                    start_nodes.push(m);
+        if let Some(successors) = self.nodes.get(node).map(|n| &n.successors) {
+            for &successor in successors {
+                if self.dfs(successor, visited, stack) {
+                    return true;
                 }
             }
-            n_out_edges.clear();
         }
+
+        stack.pop();
+        false
     }
 
     fn is_acyclic(&self) -> bool {
-        self.out_edges.values().all(|edge| edge.is_empty())
+        self.nodes
+            .values()
+            .all(|node| node.successors.is_empty() || node.predecessors_count == 0)
     }
 }