Move toposorting onto DiGraph directly, and cleanup error handling

crates/bevy_ecs/src/schedule/auto_insert_apply_deferred.rs

@@ -9,7 +9,7 @@ use crate::{
 };
 
 use super::{
-    is_apply_deferred, ApplyDeferred, DiGraph, Direction, NodeId, ReportCycles, ScheduleBuildError,
+    is_apply_deferred, ApplyDeferred, DiGraph, Direction, NodeId, ScheduleBuildError,
     ScheduleBuildPass, ScheduleGraph,
 };
 
@@ -75,7 +75,9 @@ impl ScheduleBuildPass for AutoInsertApplyDeferredPass {
         dependency_flattened: &mut DiGraph<SystemKey>,
     ) -> Result<(), ScheduleBuildError> {
         let mut sync_point_graph = dependency_flattened.clone();
-        let topo = graph.topsort_graph(dependency_flattened, ReportCycles::Dependency)?;
+        let topo = dependency_flattened
+            .toposort()
+            .map_err(ScheduleBuildError::FlatDependencySort)?;
 
         fn set_has_conditions(graph: &ScheduleGraph, set: SystemSetKey) -> bool {
             graph.system_sets.has_conditions(set)

crates/bevy_ecs/src/schedule/error.rs

@@ -5,26 +5,26 @@ use thiserror::Error;
 
 use crate::{
     component::{ComponentId, Components},
-    schedule::{graph::GraphNodeId, NodeId, ScheduleGraph, SystemKey, SystemSetKey},
+    schedule::{
+        graph::{DiGraphToposortError, GraphNodeId},
+        NodeId, ScheduleGraph, SystemKey, SystemSetKey,
+    },
     world::World,
 };
 
 /// Category of errors encountered during [`Schedule::initialize`](crate::schedule::Schedule::initialize).
 #[non_exhaustive]
 #[derive(Error, Debug)]
 pub enum ScheduleBuildError {
-    /// A system set contains itself.
-    #[error("System set `{0:?}` contains itself.")]
-    HierarchyLoop(NodeId),
-    /// The hierarchy of system sets contains a cycle.
-    #[error("The hierarchy of system sets contains a cycle: {0:?}")]
-    HierarchyCycle(Vec<Vec<NodeId>>),
-    /// A system (set) has been told to run before itself.
-    #[error("`{0:?}` has been told to run before itself.")]
-    DependencyLoop(NodeId),
-    /// The dependency graph contains a cycle.
-    #[error("The dependency graph contains a cycle: {0:?}")]
-    DependencyCycle(Vec<Vec<NodeId>>),
+    /// Tried to topologically sort the hierarchy of system sets.
+    #[error("Failed to topologically sort the hierarchy of system sets: {0}")]
+    HierarchySort(DiGraphToposortError<NodeId>),
+    /// Tried to topologically sort the dependency graph.
+    #[error("Failed to topologically sort the dependency graph: {0}")]
+    DependencySort(DiGraphToposortError<NodeId>),
+    /// Tried to topologically sort the flattened dependency graph.
+    #[error("Failed to topologically sort the flattened dependency graph: {0}")]
+    FlatDependencySort(DiGraphToposortError<SystemKey>),
     /// Tried to order a system (set) relative to a system set it belongs to.
     #[error("`{0:?}` and `{1:?}` have both `in_set` and `before`-`after` relationships (these might be transitive). This combination is unsolvable as a system cannot run before or after a set it belongs to.")]
     CrossDependency(NodeId, NodeId),
@@ -83,16 +83,22 @@ impl ScheduleBuildError {
     /// [`Schedule`]: crate::schedule::Schedule
     pub fn to_string(&self, graph: &ScheduleGraph, world: &World) -> String {
         match self {
-            ScheduleBuildError::HierarchyLoop(node_id) => {
+            ScheduleBuildError::HierarchySort(DiGraphToposortError::Loop(node_id)) => {
                 Self::hierarchy_loop_to_string(node_id, graph)
             }
-            ScheduleBuildError::HierarchyCycle(cycles) => {
+            ScheduleBuildError::HierarchySort(DiGraphToposortError::Cycle(cycles)) => {
                 Self::hierarchy_cycle_to_string(cycles, graph)
             }
-            ScheduleBuildError::DependencyLoop(node_id) => {
+            ScheduleBuildError::DependencySort(DiGraphToposortError::Loop(node_id)) => {
                 Self::dependency_loop_to_string(node_id, graph)
             }
-            ScheduleBuildError::DependencyCycle(cycles) => {
+            ScheduleBuildError::DependencySort(DiGraphToposortError::Cycle(cycles)) => {
+                Self::dependency_cycle_to_string(cycles, graph)
+            }
+            ScheduleBuildError::FlatDependencySort(DiGraphToposortError::Loop(node_id)) => {
+                Self::dependency_loop_to_string(&NodeId::System(*node_id), graph)
+            }
+            ScheduleBuildError::FlatDependencySort(DiGraphToposortError::Cycle(cycles)) => {
                 Self::dependency_cycle_to_string(cycles, graph)
             }
             ScheduleBuildError::CrossDependency(a, b) => {
@@ -164,10 +170,15 @@ impl ScheduleBuildError {
         )
     }
 
-    fn dependency_cycle_to_string(cycles: &[Vec<NodeId>], graph: &ScheduleGraph) -> String {
+    fn dependency_cycle_to_string<N: GraphNodeId + Into<NodeId>>(
+        cycles: &[Vec<N>],
+        graph: &ScheduleGraph,
+    ) -> String {
         let mut message = format!("schedule has {} before/after cycle(s):\n", cycles.len());
         for (i, cycle) in cycles.iter().enumerate() {
-            let mut names = cycle.iter().map(|id| (id.kind(), graph.get_node_name(id)));
+            let mut names = cycle
+                .iter()
+                .map(|&id| (id.kind(), graph.get_node_name(&id.into())));
             let (first_kind, first_name) = names.next().unwrap();
             writeln!(
                 message,

crates/bevy_ecs/src/schedule/graph/graph_map.rs

@@ -4,13 +4,17 @@
 //!
 //! [`petgraph`]: https://docs.rs/petgraph/0.6.5/petgraph/
 
-use alloc::vec::Vec;
+use alloc::{vec, vec::Vec};
 use core::{
     fmt::{self, Debug},
     hash::{BuildHasher, Hash},
 };
+use thiserror::Error;
 
-use bevy_platform::{collections::HashSet, hash::FixedHasher};
+use bevy_platform::{
+    collections::{HashMap, HashSet},
+    hash::FixedHasher,
+};
 use indexmap::IndexMap;
 use smallvec::SmallVec;
 
@@ -336,12 +340,177 @@ where
 }
 
 impl<N: GraphNodeId, S: BuildHasher> DiGraph<N, S> {
+    /// Tries to topologically sort this directed graph.
+    ///
+    /// If the graph is acyclic, returns [`Ok`] with the list of [`GraphNodeId`]s
+    /// in a valid topological order. If the graph contains cycles, returns [`Err`]
+    /// with the list of strongly-connected components that contain cycles
+    /// (also in a valid topological order).
+    ///
+    /// # Errors
+    ///
+    /// - If the graph contains a self-loop, returns [`DiGraphToposortError::Loop`].
+    /// - If the graph contains cycles, returns [`DiGraphToposortError::Cycle`].
+    pub fn toposort(&self) -> Result<Vec<N>, DiGraphToposortError<N>> {
+        // Check explicitly for self-edges.
+        // `iter_sccs` won't report them as cycles because they still form components of one node.
+        if let Some((node, _)) = self.all_edges().find(|(left, right)| left == right) {
+            return Err(DiGraphToposortError::Loop(node));
+        }
+
+        // Tarjan's SCC algorithm returns elements in *reverse* topological order.
+        let mut top_sorted_nodes = Vec::with_capacity(self.node_count());
+        let mut sccs_with_cycles = Vec::new();
+
+        for scc in self.iter_sccs() {
+            // A strongly-connected component is a group of nodes who can all reach each other
+            // through one or more paths. If an SCC contains more than one node, there must be
+            // at least one cycle within them.
+            top_sorted_nodes.extend_from_slice(&scc);
+            if scc.len() > 1 {
+                sccs_with_cycles.push(scc);
+            }
+        }
+
+        if sccs_with_cycles.is_empty() {
+            // reverse to get topological order
+            top_sorted_nodes.reverse();
+            Ok(top_sorted_nodes)
+        } else {
+            let mut cycles = Vec::new();
+            for scc in &sccs_with_cycles {
+                cycles.append(&mut self.simple_cycles_in_component(scc));
+            }
+
+            Err(DiGraphToposortError::Cycle(cycles))
+        }
+    }
+
+    /// Returns the simple cycles in a strongly-connected component of a directed graph.
+    ///
+    /// The algorithm implemented comes from
+    /// ["Finding all the elementary circuits of a directed graph"][1] by D. B. Johnson.
+    ///
+    /// [1]: https://doi.org/10.1137/0204007
+    pub fn simple_cycles_in_component(&self, scc: &[N]) -> Vec<Vec<N>> {
+        let mut cycles = vec![];
+        let mut sccs = vec![SmallVec::from_slice(scc)];
+
+        while let Some(mut scc) = sccs.pop() {
+            // only look at nodes and edges in this strongly-connected component
+            let mut subgraph = DiGraph::<N>::default();
+            for &node in &scc {
+                subgraph.add_node(node);
+            }
+
+            for &node in &scc {
+                for successor in self.neighbors(node) {
+                    if subgraph.contains_node(successor) {
+                        subgraph.add_edge(node, successor);
+                    }
+                }
+            }
+
+            // path of nodes that may form a cycle
+            let mut path = Vec::with_capacity(subgraph.node_count());
+            // we mark nodes as "blocked" to avoid finding permutations of the same cycles
+            let mut blocked: HashSet<_> =
+                HashSet::with_capacity_and_hasher(subgraph.node_count(), Default::default());
+            // connects nodes along path segments that can't be part of a cycle (given current root)
+            // those nodes can be unblocked at the same time
+            let mut unblock_together: HashMap<N, HashSet<N>> =
+                HashMap::with_capacity_and_hasher(subgraph.node_count(), Default::default());
+            // stack for unblocking nodes
+            let mut unblock_stack = Vec::with_capacity(subgraph.node_count());
+            // nodes can be involved in multiple cycles
+            let mut maybe_in_more_cycles: HashSet<N> =
+                HashSet::with_capacity_and_hasher(subgraph.node_count(), Default::default());
+            // stack for DFS
+            let mut stack = Vec::with_capacity(subgraph.node_count());
+
+            // we're going to look for all cycles that begin and end at this node
+            let root = scc.pop().unwrap();
+            // start a path at the root
+            path.clear();
+            path.push(root);
+            // mark this node as blocked
+            blocked.insert(root);
+
+            // DFS
+            stack.clear();
+            stack.push((root, subgraph.neighbors(root)));
+            while !stack.is_empty() {
+                let &mut (ref node, ref mut successors) = stack.last_mut().unwrap();
+                if let Some(next) = successors.next() {
+                    if next == root {
+                        // found a cycle
+                        maybe_in_more_cycles.extend(path.iter());
+                        cycles.push(path.clone());
+                    } else if !blocked.contains(&next) {
+                        // first time seeing `next` on this path
+                        maybe_in_more_cycles.remove(&next);
+                        path.push(next);
+                        blocked.insert(next);
+                        stack.push((next, subgraph.neighbors(next)));
+                        continue;
+                    } else {
+                        // not first time seeing `next` on this path
+                    }
+                }
+
+                if successors.peekable().peek().is_none() {
+                    if maybe_in_more_cycles.contains(node) {
+                        unblock_stack.push(*node);
+                        // unblock this node's ancestors
+                        while let Some(n) = unblock_stack.pop() {
+                            if blocked.remove(&n) {
+                                let unblock_predecessors = unblock_together.entry(n).or_default();
+                                unblock_stack.extend(unblock_predecessors.iter());
+                                unblock_predecessors.clear();
+                            }
+                        }
+                    } else {
+                        // if its descendants can be unblocked later, this node will be too
+                        for successor in subgraph.neighbors(*node) {
+                            unblock_together.entry(successor).or_default().insert(*node);
+                        }
+                    }
+
+                    // remove node from path and DFS stack
+                    path.pop();
+                    stack.pop();
+                }
+            }
+
+            drop(stack);
+
+            // remove node from subgraph
+            subgraph.remove_node(root);
+
+            // divide remainder into smaller SCCs
+            sccs.extend(subgraph.iter_sccs().filter(|scc| scc.len() > 1));
+        }
+
+        cycles
+    }
+
     /// Iterate over all *Strongly Connected Components* in this graph.
     pub(crate) fn iter_sccs(&self) -> impl Iterator<Item = SmallVec<[N; 4]>> + '_ {
         super::tarjan_scc::new_tarjan_scc(self)
     }
 }
 
+/// Error returned when topologically sorting a directed graph fails.
+#[derive(Error, Debug)]
+pub enum DiGraphToposortError<N: GraphNodeId> {
+    /// A self-loop was detected.
+    #[error("self-loop detected at node `{0:?}`")]
+    Loop(N),
+    /// Cycles were detected.
+    #[error("cycles detected: {0:?}")]
+    Cycle(Vec<Vec<N>>),
+}
+
 /// Edge direction.
 #[derive(Clone, Copy, Debug, PartialEq, PartialOrd, Ord, Eq, Hash)]
 #[repr(u8)]