Add System Set Filters

examples/print_schedule_graph_advanced_filters.rs

@@ -0,0 +1,38 @@
+use bevy::log::LogPlugin;
+use bevy::prelude::*;
+use bevy_app::DynEq;
+use bevy_mod_debugdump::schedule_graph::Settings;
+
+fn special_system_a() {}
+fn special_system_b() {}
+fn regular_system_c() {}
+
+#[derive(Debug, Clone, Hash, PartialEq, Eq, SystemSet)]
+enum ExampleSystemSet {
+    Special,
+    Regular,
+}
+
+fn main() {
+    let mut app = App::new();
+    app.add_plugins(DefaultPlugins.build().disable::<LogPlugin>())
+        .configure_sets(
+            Update,
+            (ExampleSystemSet::Special, ExampleSystemSet::Regular),
+        )
+        .add_systems(
+            Update,
+            (special_system_a, special_system_b, regular_system_c).chain(),
+        );
+
+    let settings = Settings::default()
+        .filter_in_crate("print_schedule_graph_advanced_filters")
+        .with_system_filter(|system| system.name().contains("special"))
+        .with_system_set_filter(|system_set| {
+            system_set
+                .as_dyn_eq()
+                .dyn_eq(ExampleSystemSet::Special.as_dyn_eq())
+        });
+    let dot = bevy_mod_debugdump::schedule_graph_dot(&mut app, Update, &settings);
+    println!("{dot}");
+}

src/schedule_graph/mod.rs

@@ -477,19 +477,28 @@ impl ScheduleGraphContext<'_> {
 }
 
 fn included_systems_sets(graph: &ScheduleGraph, settings: &Settings) -> HashSet<NodeId> {
+    let system_set_filter = settings
+        .include_system_set
+        .as_deref()
+        .unwrap_or(&include_all_system_sets);
+
     let Some(include_system) = &settings.include_system else {
-        return graph
-            .systems()
-            .map(|(id, ..)| id)
-            .chain(graph.system_sets().map(|(id, ..)| id))
-            .collect();
+        let systems = graph.systems().map(|(id, ..)| id);
+        let system_sets = graph
+            .system_sets()
+            .filter_map(|(id, set, _)| system_set_filter(set).then_some(id));
+        return systems.chain(system_sets).collect();
     };
 
     let hierarchy = graph.hierarchy().graph();
 
     let root_sets = hierarchy.nodes().filter(|&node| {
-        node.is_set()
-            && graph.set_at(node).system_type().is_none()
+        if !node.is_set() {
+            return false;
+        }
+        let system_set = graph.set_at(node);
+        system_set.system_type().is_none()
+            && system_set_filter(system_set)
             && hierarchy
                 .neighbors_directed(node, Direction::Incoming)
                 .next()
@@ -502,24 +511,39 @@ fn included_systems_sets(graph: &ScheduleGraph, settings: &Settings) -> HashSet<
         .map(|(id, ..)| id)
         .collect();
 
+    fn include_all_system_sets(_: &dyn SystemSet) -> bool {
+        true
+    }
+
     fn include_ancestors(
         id: NodeId,
+        graph: &ScheduleGraph,
         hierarchy: &DiGraphMap<NodeId, ()>,
+        filter: &impl Fn(&dyn SystemSet) -> bool,
         included_systems_sets: &mut HashSet<NodeId>,
     ) {
         let parents = hierarchy.neighbors_directed(id, Direction::Incoming);
 
         for parent in parents {
-            included_systems_sets.insert(parent);
-            include_ancestors(parent, hierarchy, included_systems_sets);
+            let system_set = graph.set_at(parent);
+            if filter(system_set) {
+                included_systems_sets.insert(parent);
+                include_ancestors(parent, graph, hierarchy, filter, included_systems_sets);
+            }
         }
     }
 
     let mut included_systems_sets = systems_of_interest.clone();
     included_systems_sets.extend(root_sets);
 
     for &id in &systems_of_interest {
-        include_ancestors(id, hierarchy, &mut included_systems_sets);
+        include_ancestors(
+            id,
+            graph,
+            hierarchy,
+            &system_set_filter,
+            &mut included_systems_sets,
+        );
     }
 
     if settings.ambiguity_enable {
@@ -540,12 +564,24 @@ fn included_systems_sets(graph: &ScheduleGraph, settings: &Settings) -> HashSet<
     for (from, to, ()) in graph.dependency().graph().all_edges() {
         if systems_of_interest.contains(&from) {
             included_systems_sets.insert(to);
-            include_ancestors(to, hierarchy, &mut included_systems_sets);
+            include_ancestors(
+                to,
+                graph,
+                hierarchy,
+                &system_set_filter,
+                &mut included_systems_sets,
+            );
         }
 
         if systems_of_interest.contains(&to) {
             included_systems_sets.insert(from);
-            include_ancestors(to, hierarchy, &mut included_systems_sets);
+            include_ancestors(
+                to,
+                graph,
+                hierarchy,
+                &system_set_filter,
+                &mut included_systems_sets,
+            );
         }
     }
 

src/schedule_graph/settings.rs

@@ -183,6 +183,7 @@ pub struct Settings {
 
     /// When set to `Some`, will only include systems matching the predicate, and their ancestor sets
     pub include_system: Option<SystemMapperFn<bool>>,
+    pub include_system_set: Option<SystemSetMapperFn<bool>>,
     pub collapse_single_system_sets: bool,
 
     pub ambiguity_enable: bool,
@@ -196,29 +197,37 @@ pub struct Settings {
 
 impl Settings {
     /// Set the `include_system` predicate to match only systems for which their names matches `filter`
-    pub fn filter_name(mut self, filter: impl Fn(&str) -> bool + 'static) -> Self {
-        self.include_system = Some(Box::new(move |system| {
-            let name = system.name();
-            filter(&name)
-        }));
-        self
+    pub fn filter_name(self, filter: impl Fn(&str) -> bool + 'static) -> Self {
+        self.with_system_filter(move |system| filter(&system.name()))
     }
     /// Set the `include_system` predicate to only match systems from the specified crate
-    pub fn filter_in_crate(mut self, crate_: &str) -> Self {
+    pub fn filter_in_crate(self, crate_: &str) -> Self {
         let crate_ = crate_.to_owned();
-        self.include_system = Some(Box::new(move |system| {
-            let name = system.name();
-            name.starts_with(&crate_)
-        }));
-        self
+        self.with_system_filter(move |system| system.name().starts_with(&crate_))
     }
     /// Set the `include_system` predicate to only match systems from the specified crates
-    pub fn filter_in_crates(mut self, crates: &[&str]) -> Self {
+    pub fn filter_in_crates(self, crates: &[&str]) -> Self {
         let crates: Vec<_> = crates.iter().map(|&s| s.to_owned()).collect();
-        self.include_system = Some(Box::new(move |system| {
-            let name = system.name();
-            crates.iter().any(|crate_| name.starts_with(crate_))
-        }));
+        self.with_system_filter(move |system| {
+            crates
+                .iter()
+                .any(|crate_| system.name().starts_with(crate_))
+        })
+    }
+
+    pub fn with_system_filter(
+        mut self,
+        filter: impl Fn(&dyn System<In = (), Out = ()>) -> bool + 'static,
+    ) -> Self {
+        self.include_system = Some(Box::new(filter));
+        self
+    }
+
+    pub fn with_system_set_filter(
+        mut self,
+        filter: impl Fn(&dyn SystemSet) -> bool + 'static,
+    ) -> Self {
+        self.include_system_set = Some(Box::new(filter));
         self
     }
 
@@ -301,6 +310,7 @@ impl Default for Settings {
             system_style: Box::new(system_to_style),
 
             include_system: None,
+            include_system_set: None,
             collapse_single_system_sets: false,
 
             ambiguity_enable: true,
@@ -324,4 +334,4 @@ pub fn full_system_name(system: &dyn System<In = (), Out = ()>) -> String {
 
 pub fn default_system_set_name(system_set: &dyn SystemSet) -> String {
     format!("{:?}", system_set)
-}
+}