Use explicitly added ApplyDeferred stages when computing automatically inserted sync points.

crates/bevy_ecs/src/schedule/schedule.rs

@@ -1152,39 +1152,69 @@ impl ScheduleGraph {
         let topo = self.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>> =
+        let mut distances: HashMap<usize, u32> =
             HashMap::with_capacity_and_hasher(topo.len(), Default::default());
+        // Keep track of any explicit sync nodes for a specific distance.
+        let mut distance_to_explicit_sync_node: HashMap<u32, NodeId> = HashMap::default();
         for node in &topo {
-            let add_sync_after = self.systems[node.index()].get().unwrap().has_deferred();
+            let node_system = self.systems[node.index()].get().unwrap();
+
+            let node_needs_sync = if is_apply_deferred(node_system) {
+                distance_to_explicit_sync_node.insert(
+                    distances.get(&node.index()).copied().unwrap_or_default(),
+                    *node,
+                );
+
+                // This node just did a sync, so the only reason to do another sync is if one was
+                // explicitly scheduled afterwards.
+                false
+            } else {
+                node_system.has_deferred()
+            };
 
             for target in dependency_flattened.neighbors_directed(*node, Outgoing) {
-                let add_sync_on_edge = add_sync_after
-                    && !is_apply_deferred(self.systems[target.index()].get().unwrap())
-                    && !self.no_sync_edges.contains(&(*node, target));
+                let edge_needs_sync = node_needs_sync
+                    && !self.no_sync_edges.contains(&(*node, target))
+                    || is_apply_deferred(self.systems[target.index()].get().unwrap());
 
-                let weight = if add_sync_on_edge { 1 } else { 0 };
+                let weight = if edge_needs_sync { 1 } else { 0 };
 
+                // Use whichever distance is larger, either the current distance, or the distance to
+                // the parent plus the weight.
                 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,
-                        )
-                    });
+                    .copied()
+                    .unwrap_or_default()
+                    .max(distances.get(&node.index()).copied().unwrap_or_default() + weight);
 
                 distances.insert(target.index(), distance);
+            }
+        }
 
-                if add_sync_on_edge {
-                    let sync_point = self.get_sync_point(distances[&target.index()].unwrap());
-                    sync_point_graph.add_edge(*node, sync_point);
-                    sync_point_graph.add_edge(sync_point, target);
+        // Find any edges which have a different number of sync points between them and make sure
+        for node in &topo {
+            let node_distance = distances.get(&node.index()).copied().unwrap_or_default();
+            for target in dependency_flattened.neighbors_directed(*node, Outgoing) {
+                let target_distance = distances.get(&target.index()).copied().unwrap_or_default();
+                if node_distance == target_distance {
+                    // These nodes are the same distance, so they don't need an edge between them.
+                    continue;
+                }
 
-                    // edge is now redundant
-                    sync_point_graph.remove_edge(*node, target);
+                if is_apply_deferred(self.systems[target.index()].get().unwrap()) {
+                    // We don't need to insert a sync point since ApplyDeferred is a sync point
+                    // already!
+                    continue;
                 }
+                let sync_point = distance_to_explicit_sync_node
+                    .get(&target_distance)
+                    .copied()
+                    .unwrap_or_else(|| self.get_sync_point(target_distance));
+
+                sync_point_graph.add_edge(*node, sync_point);
+                sync_point_graph.add_edge(sync_point, target);
+
+                sync_point_graph.remove_edge(*node, target);
             }
         }
 
@@ -2050,7 +2080,7 @@ mod tests {
 
     use crate::{
         self as bevy_ecs,
-        prelude::{Res, Resource},
+        prelude::{ApplyDeferred, Res, Resource},
         schedule::{
             tests::ResMut, IntoSystemConfigs, IntoSystemSetConfigs, Schedule,
             ScheduleBuildSettings, SystemSet,
@@ -2102,6 +2132,24 @@ mod tests {
         assert_eq!(schedule.executable.systems.len(), 3);
     }
 
+    #[test]
+    fn explicit_sync_point_used_as_auto_sync_point() {
+        let mut schedule = Schedule::default();
+        let mut world = World::default();
+        schedule.add_systems(
+            (
+                |mut commands: Commands| commands.insert_resource(Resource1),
+                |_: Res<Resource1>| {},
+            )
+                .chain(),
+        );
+        schedule.add_systems((|| {}, ApplyDeferred, || {}).chain());
+        schedule.run(&mut world);
+
+        // No sync point was inserted, since we can reuse the explicit sync point.
+        assert_eq!(schedule.executable.systems.len(), 5);
+    }
+
     #[test]
     fn merges_sync_points_into_one() {
         let mut schedule = Schedule::default();