Schedule build pass (#11094)

# Objective

This is a follow up to #9822, which automatically adds sync points
during the Schedule build process.

However, the implementation in #9822 feels very "special case" to me. As
the number of things we want to do with the `Schedule` grows, we need a
modularized way to manage those behaviors. For example, in one of my
current experiments I want to automatically add systems to apply GPU
pipeline barriers between systems accessing GPU resources.

For dynamic modifications of the schedule, we mostly need these
capabilities:
- Storing custom data on schedule edges
- Storing custom data on schedule nodes
- Modify the schedule graph whenever it builds

These should be enough to allows us to add "hooks" to the schedule build
process for various reasons.

cc @hymm 

## Solution
This PR abstracts the process of schedule modification and created a new
trait, `ScheduleBuildPass`. Most of the logics in #9822 were moved to an
implementation of `ScheduleBuildPass`, `AutoInsertApplyDeferredPass`.

Whether a dependency edge should "ignore deferred" is now indicated by
the presence of a marker struct, `IgnoreDeferred`.

This PR has no externally visible effects. However, in a future PR I
propose to change the `before_ignore_deferred` and
`after_ignore_deferred` API into a more general form,
`before_with_options` and `after_with_options`.

```rs
schedule.add_systems(
    system.before_with_options(another_system, IgnoreDeferred)
);

schedule.add_systems(
    system.before_with_options(another_system, (
        IgnoreDeferred,
        AnyOtherOption {
            key: value
        }
    ))
);

schedule.add_systems(
    system.before_with_options(another_system, ())
);
```

Author: PixelDust22

crates/bevy_ecs/src/schedule/auto_insert_apply_deferred.rs

@@ -0,0 +1,161 @@
+use alloc::{boxed::Box, collections::BTreeSet, vec::Vec};
+
+use bevy_platform_support::collections::HashMap;
+
+use crate::system::IntoSystem;
+use crate::world::World;
+
+use super::{
+    is_apply_deferred, ApplyDeferred, DiGraph, Direction, NodeId, ReportCycles, ScheduleBuildError,
+    ScheduleBuildPass, ScheduleGraph, SystemNode,
+};
+
+/// A [`ScheduleBuildPass`] that inserts [`ApplyDeferred`] systems into the schedule graph
+/// when there are [`Deferred`](crate::prelude::Deferred)
+/// in one system and there are ordering dependencies on that system. [`Commands`](crate::system::Commands) is one
+/// such deferred buffer.
+///
+/// This pass is typically automatically added to the schedule. You can disable this by setting
+/// [`ScheduleBuildSettings::auto_insert_apply_deferred`](crate::schedule::ScheduleBuildSettings::auto_insert_apply_deferred)
+/// to `false`. You may want to disable this if you only want to sync deferred params at the end of the schedule,
+/// or want to manually insert all your sync points.
+#[derive(Debug, Default)]
+pub struct AutoInsertApplyDeferredPass {
+    /// Dependency edges that will **not** automatically insert an instance of `ApplyDeferred` on the edge.
+    no_sync_edges: BTreeSet<(NodeId, NodeId)>,
+    auto_sync_node_ids: HashMap<u32, NodeId>,
+}
+
+/// If added to a dependency edge, the edge will not be considered for auto sync point insertions.
+pub struct IgnoreDeferred;
+
+impl AutoInsertApplyDeferredPass {
+    /// Returns the `NodeId` of the cached auto sync point. Will create
+    /// a new one if needed.
+    fn get_sync_point(&mut self, graph: &mut ScheduleGraph, distance: u32) -> NodeId {
+        self.auto_sync_node_ids
+            .get(&distance)
+            .copied()
+            .or_else(|| {
+                let node_id = self.add_auto_sync(graph);
+                self.auto_sync_node_ids.insert(distance, node_id);
+                Some(node_id)
+            })
+            .unwrap()
+    }
+    /// add an [`ApplyDeferred`] system with no config
+    fn add_auto_sync(&mut self, graph: &mut ScheduleGraph) -> NodeId {
+        let id = NodeId::System(graph.systems.len());
+
+        graph
+            .systems
+            .push(SystemNode::new(Box::new(IntoSystem::into_system(
+                ApplyDeferred,
+            ))));
+        graph.system_conditions.push(Vec::new());
+
+        // ignore ambiguities with auto sync points
+        // They aren't under user control, so no one should know or care.
+        graph.ambiguous_with_all.insert(id);
+
+        id
+    }
+}
+
+impl ScheduleBuildPass for AutoInsertApplyDeferredPass {
+    type EdgeOptions = IgnoreDeferred;
+
+    fn add_dependency(&mut self, from: NodeId, to: NodeId, options: Option<&Self::EdgeOptions>) {
+        if options.is_some() {
+            self.no_sync_edges.insert((from, to));
+        }
+    }
+
+    fn build(
+        &mut self,
+        _world: &mut World,
+        graph: &mut ScheduleGraph,
+        dependency_flattened: &mut DiGraph,
+    ) -> Result<(), ScheduleBuildError> {
+        let mut sync_point_graph = dependency_flattened.clone();
+        let topo = graph.topsort_graph(dependency_flattened, ReportCycles::Dependency)?;
+
+        // calculate the number of sync points each sync point is from the beginning of the graph
+        // use the same sync point if the distance is the same
+        let mut distances: HashMap<usize, Option<u32>> =
+            HashMap::with_capacity_and_hasher(topo.len(), Default::default());
+        for node in &topo {
+            let add_sync_after = graph.systems[node.index()].get().unwrap().has_deferred();
+
+            for target in dependency_flattened.neighbors_directed(*node, Direction::Outgoing) {
+                let add_sync_on_edge = add_sync_after
+                    && !is_apply_deferred(graph.systems[target.index()].get().unwrap())
+                    && !self.no_sync_edges.contains(&(*node, target));
+
+                let weight = if add_sync_on_edge { 1 } else { 0 };
+
+                let distance = distances
+                    .get(&target.index())
+                    .unwrap_or(&None)
+                    .or(Some(0))
+                    .map(|distance| {
+                        distance.max(
+                            distances.get(&node.index()).unwrap_or(&None).unwrap_or(0) + weight,
+                        )
+                    });
+
+                distances.insert(target.index(), distance);
+
+                if add_sync_on_edge {
+                    let sync_point =
+                        self.get_sync_point(graph, distances[&target.index()].unwrap());
+                    sync_point_graph.add_edge(*node, sync_point);
+                    sync_point_graph.add_edge(sync_point, target);
+
+                    // edge is now redundant
+                    sync_point_graph.remove_edge(*node, target);
+                }
+            }
+        }
+
+        *dependency_flattened = sync_point_graph;
+        Ok(())
+    }
+
+    fn collapse_set(
+        &mut self,
+        set: NodeId,
+        systems: &[NodeId],
+        dependency_flattened: &DiGraph,
+    ) -> impl Iterator<Item = (NodeId, NodeId)> {
+        if systems.is_empty() {
+            // collapse dependencies for empty sets
+            for a in dependency_flattened.neighbors_directed(set, Direction::Incoming) {
+                for b in dependency_flattened.neighbors_directed(set, Direction::Outgoing) {
+                    if self.no_sync_edges.contains(&(a, set))
+                        && self.no_sync_edges.contains(&(set, b))
+                    {
+                        self.no_sync_edges.insert((a, b));
+                    }
+                }
+            }
+        } else {
+            for a in dependency_flattened.neighbors_directed(set, Direction::Incoming) {
+                for &sys in systems {
+                    if self.no_sync_edges.contains(&(a, set)) {
+                        self.no_sync_edges.insert((a, sys));
+                    }
+                }
+            }
+
+            for b in dependency_flattened.neighbors_directed(set, Direction::Outgoing) {
+                for &sys in systems {
+                    if self.no_sync_edges.contains(&(set, b)) {
+                        self.no_sync_edges.insert((sys, b));
+                    }
+                }
+            }
+        }
+        core::iter::empty()
+    }
+}

crates/bevy_ecs/src/schedule/config.rs

@@ -4,6 +4,7 @@ use variadics_please::all_tuples;
 use crate::{
     result::Result,
     schedule::{
+        auto_insert_apply_deferred::IgnoreDeferred,
         condition::{BoxedCondition, Condition},
         graph::{Ambiguity, Dependency, DependencyKind, GraphInfo},
         set::{InternedSystemSet, IntoSystemSet, SystemSet},
@@ -137,7 +138,7 @@ impl<T> NodeConfigs<T> {
                 config
                     .graph_info
                     .dependencies
-                    .push(Dependency::new(DependencyKind::BeforeNoSync, set));
+                    .push(Dependency::new(DependencyKind::Before, set).add_config(IgnoreDeferred));
             }
             Self::Configs { configs, .. } => {
                 for config in configs {
@@ -153,7 +154,7 @@ impl<T> NodeConfigs<T> {
                 config
                     .graph_info
                     .dependencies
-                    .push(Dependency::new(DependencyKind::AfterNoSync, set));
+                    .push(Dependency::new(DependencyKind::After, set).add_config(IgnoreDeferred));
             }
             Self::Configs { configs, .. } => {
                 for config in configs {
@@ -224,17 +225,17 @@ impl<T> NodeConfigs<T> {
         match &mut self {
             Self::NodeConfig(_) => { /* no op */ }
             Self::Configs { chained, .. } => {
-                *chained = Chain::Yes;
+                chained.set_chained();
             }
-        }
+        };
         self
     }
 
     fn chain_ignore_deferred_inner(mut self) -> Self {
         match &mut self {
             Self::NodeConfig(_) => { /* no op */ }
             Self::Configs { chained, .. } => {
-                *chained = Chain::YesIgnoreDeferred;
+                chained.set_chained_with_config(IgnoreDeferred);
             }
         }
         self
@@ -582,7 +583,7 @@ macro_rules! impl_system_collection {
                 SystemConfigs::Configs {
                     configs: vec![$($sys.into_configs(),)*],
                     collective_conditions: Vec::new(),
-                    chained: Chain::No,
+                    chained: Default::default(),
                 }
             }
         }
@@ -820,7 +821,7 @@ macro_rules! impl_system_set_collection {
                 SystemSetConfigs::Configs {
                     configs: vec![$($set.into_configs(),)*],
                     collective_conditions: Vec::new(),
-                    chained: Chain::No,
+                    chained: Default::default(),
                 }
             }
         }

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

@@ -65,7 +65,7 @@ where
     S: BuildHasher,
 {
     /// Create a new `Graph` with estimated capacity.
-    pub(crate) fn with_capacity(nodes: usize, edges: usize) -> Self
+    pub fn with_capacity(nodes: usize, edges: usize) -> Self
     where
         S: Default,
     {
@@ -89,14 +89,14 @@ where
     }
 
     /// Add node `n` to the graph.
-    pub(crate) fn add_node(&mut self, n: NodeId) {
+    pub fn add_node(&mut self, n: NodeId) {
         self.nodes.entry(n).or_default();
     }
 
     /// Remove a node `n` from the graph.
     ///
     /// Computes in **O(N)** time, due to the removal of edges with other nodes.
-    pub(crate) fn remove_node(&mut self, n: NodeId) {
+    pub fn remove_node(&mut self, n: NodeId) {
         let Some(links) = self.nodes.swap_remove(&n) else {
             return;
         };
@@ -166,7 +166,7 @@ where
     /// Remove edge from `a` to `b` from the graph.
     ///
     /// Return `false` if the edge didn't exist.
-    pub(crate) fn remove_edge(&mut self, a: NodeId, b: NodeId) -> bool {
+    pub fn remove_edge(&mut self, a: NodeId, b: NodeId) -> bool {
         let exist1 = self.remove_single_edge(a, b, Outgoing);
         let exist2 = if a != b {
             self.remove_single_edge(b, a, Incoming)

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

@@ -1,8 +1,13 @@
-use alloc::{vec, vec::Vec};
-use core::fmt::Debug;
+use alloc::{boxed::Box, vec, vec::Vec};
+use core::{
+    any::{Any, TypeId},
+    fmt::Debug,
+};
 use smallvec::SmallVec;
 
 use bevy_platform_support::collections::{HashMap, HashSet};
+use bevy_utils::TypeIdMap;
+
 use fixedbitset::FixedBitSet;
 
 use crate::schedule::set::*;
@@ -21,22 +26,26 @@ pub(crate) enum DependencyKind {
     Before,
     /// A node that should be succeeded.
     After,
-    /// A node that should be preceded and will **not** automatically insert an instance of `ApplyDeferred` on the edge.
-    BeforeNoSync,
-    /// A node that should be succeeded and will **not** automatically insert an instance of `ApplyDeferred` on the edge.
-    AfterNoSync,
 }
 
 /// An edge to be added to the dependency graph.
-#[derive(Clone)]
 pub(crate) struct Dependency {
     pub(crate) kind: DependencyKind,
     pub(crate) set: InternedSystemSet,
+    pub(crate) options: TypeIdMap<Box<dyn Any>>,
 }
 
 impl Dependency {
     pub fn new(kind: DependencyKind, set: InternedSystemSet) -> Self {
-        Self { kind, set }
+        Self {
+            kind,
+            set,
+            options: Default::default(),
+        }
+    }
+    pub fn add_config<T: 'static>(mut self, option: T) -> Self {
+        self.options.insert(TypeId::of::<T>(), Box::new(option));
+        self
     }
 }
 
@@ -52,7 +61,7 @@ pub(crate) enum Ambiguity {
 }
 
 /// Metadata about how the node fits in the schedule graph
-#[derive(Clone, Default)]
+#[derive(Default)]
 pub(crate) struct GraphInfo {
     /// the sets that the node belongs to (hierarchy)
     pub(crate) hierarchy: Vec<InternedSystemSet>,

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

@@ -13,7 +13,7 @@ pub enum NodeId {
 
 impl NodeId {
     /// Returns the internal integer value.
-    pub(crate) const fn index(&self) -> usize {
+    pub const fn index(&self) -> usize {
         match self {
             NodeId::System(index) | NodeId::Set(index) => *index,
         }

crates/bevy_ecs/src/schedule/mod.rs

@@ -1,18 +1,28 @@
 //! Contains APIs for ordering systems and executing them on a [`World`](crate::world::World)
 
+mod auto_insert_apply_deferred;
 mod condition;
 mod config;
 mod executor;
-mod graph;
+mod pass;
 mod schedule;
 mod set;
 mod stepping;
 
 use self::graph::*;
 pub use self::{condition::*, config::*, executor::*, schedule::*, set::*};
+pub use pass::ScheduleBuildPass;
 
 pub use self::graph::NodeId;
 
+/// An implementation of a graph data structure.
+pub mod graph;
+
+/// Included optional schedule build passes.
+pub mod passes {
+    pub use crate::schedule::auto_insert_apply_deferred::*;
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -1082,7 +1092,7 @@ mod tests {
 
             schedule.graph_mut().initialize(&mut world);
             let _ = schedule.graph_mut().build_schedule(
-                world.components(),
+                &mut world,
                 TestSchedule.intern(),
                 &BTreeSet::new(),
             );
@@ -1131,7 +1141,7 @@ mod tests {
             let mut world = World::new();
             schedule.graph_mut().initialize(&mut world);
             let _ = schedule.graph_mut().build_schedule(
-                world.components(),
+                &mut world,
                 TestSchedule.intern(),
                 &BTreeSet::new(),
             );