Add default system sets based on World accesses

crates/bevy_ecs/macros/src/component.rs

@@ -1,7 +1,7 @@
-use bevy_macro_utils::{get_lit_str, Symbol};
+use bevy_macro_utils::{get_lit_str, NamedArg, Symbol};
 use proc_macro::TokenStream;
 use proc_macro2::{Span, TokenStream as TokenStream2};
-use quote::{quote, ToTokens};
+use quote::quote;
 use syn::{parse_macro_input, parse_quote, DeriveInput, Error, Ident, Path, Result};
 
 pub fn derive_resource(input: TokenStream) -> TokenStream {
@@ -33,6 +33,9 @@ pub fn derive_component(input: TokenStream) -> TokenStream {
 
     let storage = storage_path(&bevy_ecs_path, attrs.storage);
 
+    let read_sets = attrs.read_sets;
+    let write_sets = attrs.write_sets;
+
     ast.generics
         .make_where_clause()
         .predicates
@@ -44,6 +47,14 @@ pub fn derive_component(input: TokenStream) -> TokenStream {
     TokenStream::from(quote! {
         impl #impl_generics #bevy_ecs_path::component::Component for #struct_name #type_generics #where_clause {
             type Storage = #storage;
+
+            fn add_read_sets(_sets: &mut Vec<#bevy_ecs_path::schedule::BoxedSystemSet>) {
+                #(_sets.push(#read_sets.dyn_clone());)*
+            }
+
+            fn add_write_sets(_sets: &mut Vec<#bevy_ecs_path::schedule::BoxedSystemSet>) {
+                #(_sets.push(Box::new(#write_sets));)*
+            }
         }
     })
 }
@@ -53,6 +64,8 @@ pub const STORAGE: Symbol = Symbol("storage");
 
 struct Attrs {
     storage: StorageTy,
+    read_sets: Vec<syn::Expr>,
+    write_sets: Vec<syn::Expr>,
 }
 
 #[derive(Clone, Copy)]
@@ -66,47 +79,36 @@ const TABLE: &str = "Table";
 const SPARSE_SET: &str = "SparseSet";
 
 fn parse_component_attr(ast: &DeriveInput) -> Result<Attrs> {
-    let meta_items = bevy_macro_utils::parse_attrs(ast, COMPONENT)?;
-
     let mut attrs = Attrs {
         storage: StorageTy::Table,
+        read_sets: vec![],
+        write_sets: vec![],
     };
 
-    for meta in meta_items {
-        use syn::{
-            Meta::NameValue,
-            NestedMeta::{Lit, Meta},
-        };
-        match meta {
-            Meta(NameValue(m)) if m.path == STORAGE => {
-                attrs.storage = match get_lit_str(STORAGE, &m.lit)?.value().as_str() {
+    for NamedArg { path, expr } in bevy_macro_utils::parse_attrs(ast, COMPONENT)? {
+        match path.get_ident().unwrap().to_string().as_str() {
+            "storage" => {
+                attrs.storage = match get_lit_str(STORAGE, &expr)?.value().as_str() {
                     TABLE => StorageTy::Table,
                     SPARSE_SET => StorageTy::SparseSet,
                     s => {
                         return Err(Error::new_spanned(
-                            m.lit,
+                            expr,
                             format!(
                                 "Invalid storage type `{s}`, expected '{TABLE}' or '{SPARSE_SET}'.",
                             ),
                         ))
                     }
                 };
             }
-            Meta(meta_item) => {
+            "read_set" => attrs.read_sets.push(expr),
+            "write_set" => attrs.write_sets.push(expr),
+            _ => {
                 return Err(Error::new_spanned(
-                    meta_item.path(),
-                    format!(
-                        "unknown component attribute `{}`",
-                        meta_item.path().into_token_stream()
-                    ),
+                    path,
+                    "Invalid component attribute format: expected `storages`, `read_set`, or `write_set`",
                 ));
             }
-            Lit(lit) => {
-                return Err(Error::new_spanned(
-                    lit,
-                    "unexpected literal in component attribute",
-                ))
-            }
         }
     }
 

crates/bevy_ecs/macros/src/fetch.rs

@@ -323,6 +323,10 @@ pub fn derive_world_query_impl(ast: DeriveInput) -> TokenStream {
                 fn matches_component_set(state: &Self::State, _set_contains_id: &impl Fn(#path::component::ComponentId) -> bool) -> bool {
                     true #(&& <#field_types>::matches_component_set(&state.#field_idents, _set_contains_id))*
                 }
+
+                fn add_default_sets(_sets: &mut Vec<#path::schedule::BoxedSystemSet>) {
+                    #( <#field_types>::add_default_sets(_sets); )*
+                }
             }
         }
     };

crates/bevy_ecs/macros/src/lib.rs

@@ -236,6 +236,11 @@ pub fn impl_param_set(_input: TokenStream) -> TokenStream {
                         change_tick,
                     }
                 }
+
+                #[inline]
+                fn add_default_sets(_sets: &mut Vec<BoxedSystemSet>) {
+                    #( #param::add_default_sets(_sets); )*
+                }
             }
 
             impl<'w, 's, #(#param: SystemParam,)*> ParamSet<'w, 's, (#(#param,)*)>
@@ -441,6 +446,11 @@ pub fn derive_system_param(input: TokenStream) -> TokenStream {
                         #(#ignored_fields: std::default::Default::default(),)*
                     }
                 }
+
+                #[inline]
+                fn add_default_sets(_sets: &mut Vec<#path::schedule::BoxedSystemSet>) {
+                    <(#(#tuple_types,)*) as #path::system::SystemParam>::add_default_sets(_sets);
+                }
             }
 
             // Safety: Each field is `ReadOnlySystemParam`, so this can only read from the `World`

crates/bevy_ecs/src/component.rs

@@ -2,6 +2,7 @@
 
 use crate::{
     change_detection::MAX_CHANGE_AGE,
+    schedule::BoxedSystemSet,
     storage::{SparseSetIndex, Storages},
     system::{Local, Resource},
     world::{FromWorld, World},
@@ -147,6 +148,12 @@ use std::{
 /// [`Exclusive`]: https://doc.rust-lang.org/nightly/std/sync/struct.Exclusive.html
 pub trait Component: Send + Sync + 'static {
     type Storage: ComponentStorage;
+
+    #[allow(unused_variables)]
+    fn add_read_sets(sets: &mut Vec<BoxedSystemSet>) {}
+
+    #[allow(unused_variables)]
+    fn add_write_sets(sets: &mut Vec<BoxedSystemSet>) {}
 }
 
 pub struct TableStorage;

crates/bevy_ecs/src/lib.rs

@@ -65,6 +65,8 @@ type TypeIdMap<V> = rustc_hash::FxHashMap<TypeId, V>;
 mod tests {
     use crate as bevy_ecs;
     use crate::prelude::Or;
+    use crate::schedule::{IntoSystemConfig, Schedule, ScheduleBuildSettings};
+    use crate::system::Query;
     use crate::{
         bundle::Bundle,
         change_detection::Ref,
@@ -74,6 +76,7 @@ mod tests {
         system::Resource,
         world::{Mut, World},
     };
+    use bevy_ecs_macros::SystemSet;
     use bevy_tasks::{ComputeTaskPool, TaskPool};
     use std::{
         any::TypeId,
@@ -1738,4 +1741,31 @@ mod tests {
             "new entity was spawned and received C component"
         );
     }
+
+    #[test]
+    fn component_with_access_set() {
+        #[derive(SystemSet, Debug, Clone, Copy, PartialEq, Eq, Hash)]
+        struct WriteSet;
+
+        #[derive(Component)]
+        #[component(write_set = WriteSet)]
+        struct WithWriteSet;
+
+        fn does_write(_q: Query<&mut WithWriteSet>) {}
+        fn after_write(_q: Query<&WithWriteSet>) {}
+
+        let mut world = World::new();
+        let mut schedule = Schedule::new();
+
+        schedule.add_system(does_write);
+        schedule.add_system(after_write.after(WriteSet));
+
+        // Panic if there are ambiguities.
+        schedule.set_build_settings(ScheduleBuildSettings {
+            ambiguity_detection: bevy_ecs::schedule::LogLevel::Error,
+            ..Default::default()
+        });
+
+        schedule.run(&mut world);
+    }
 }

crates/bevy_ecs/src/query/fetch.rs

@@ -4,6 +4,7 @@ use crate::{
     component::{Component, ComponentId, ComponentStorage, ComponentTicks, StorageType, Tick},
     entity::Entity,
     query::{Access, DebugCheckedUnwrap, FilteredAccess},
+    schedule::BoxedSystemSet,
     storage::{ComponentSparseSet, Table, TableRow},
     world::{Mut, Ref, World},
 };
@@ -428,6 +429,9 @@ pub unsafe trait WorldQuery {
         state: &Self::State,
         set_contains_id: &impl Fn(ComponentId) -> bool,
     ) -> bool;
+
+    #[allow(unused_variables)]
+    fn add_default_sets(sets: &mut Vec<BoxedSystemSet>) {}
 }
 
 /// A world query that is read only.
@@ -648,6 +652,10 @@ unsafe impl<T: Component> WorldQuery for &T {
     ) -> bool {
         set_contains_id(state)
     }
+
+    fn add_default_sets(sets: &mut Vec<BoxedSystemSet>) {
+        T::add_read_sets(sets);
+    }
 }
 
 /// SAFETY: access is read only
@@ -809,6 +817,10 @@ unsafe impl<'__w, T: Component> WorldQuery for Ref<'__w, T> {
     ) -> bool {
         set_contains_id(state)
     }
+
+    fn add_default_sets(sets: &mut Vec<BoxedSystemSet>) {
+        T::add_read_sets(sets);
+    }
 }
 
 /// SAFETY: access is read only
@@ -974,6 +986,10 @@ unsafe impl<'__w, T: Component> WorldQuery for &'__w mut T {
     ) -> bool {
         set_contains_id(state)
     }
+
+    fn add_default_sets(sets: &mut Vec<BoxedSystemSet>) {
+        T::add_write_sets(sets);
+    }
 }
 
 #[doc(hidden)]
@@ -1078,6 +1094,10 @@ unsafe impl<T: WorldQuery> WorldQuery for Option<T> {
     ) -> bool {
         true
     }
+
+    fn add_default_sets(sets: &mut Vec<BoxedSystemSet>) {
+        T::add_default_sets(sets);
+    }
 }
 
 /// SAFETY: [`OptionFetch`] is read only because `T` is read only
@@ -1559,6 +1579,10 @@ macro_rules! impl_anytuple_fetch {
                 let ($($name,)*) = _state;
                 false $(|| $name::matches_component_set($name, _set_contains_id))*
             }
+
+            fn add_default_sets(_sets: &mut Vec<BoxedSystemSet>) {
+                $( $name::add_default_sets(_sets); )*
+            }
         }
 
         /// SAFETY: each item in the tuple is read only

crates/bevy_ecs/src/schedule/config.rs

@@ -42,10 +42,9 @@ pub struct SystemConfig {
 
 impl SystemConfig {
     fn new(system: BoxedSystem) -> Self {
-        // include system in its default sets
-        let sets = system.default_system_sets().into_iter().collect();
         let mut graph_info = GraphInfo::system();
-        graph_info.sets = sets;
+        // Include the system in its default sets.
+        graph_info.sets = system.default_system_sets();
         Self {
             system,
             graph_info,

crates/bevy_ecs/src/system/function_system.rs

@@ -4,6 +4,7 @@ use crate::{
     component::ComponentId,
     prelude::FromWorld,
     query::{Access, FilteredAccessSet},
+    schedule::BoxedSystemSet,
     system::{check_system_change_tick, ReadOnlySystemParam, System, SystemParam, SystemParamItem},
     world::{World, WorldId},
 };
@@ -515,9 +516,11 @@ where
         );
     }
 
-    fn default_system_sets(&self) -> Vec<Box<dyn crate::schedule::SystemSet>> {
+    fn default_system_sets(&self) -> Vec<BoxedSystemSet> {
         let set = crate::schedule::SystemTypeSet::<F>::new();
-        vec![Box::new(set)]
+        let mut sets = vec![Box::new(set) as BoxedSystemSet];
+        F::Param::add_default_sets(&mut sets);
+        sets
     }
 }
 

crates/bevy_ecs/src/system/system_param.rs

@@ -8,6 +8,7 @@ use crate::{
     query::{
         Access, FilteredAccess, FilteredAccessSet, QueryState, ReadOnlyWorldQuery, WorldQuery,
     },
+    schedule::BoxedSystemSet,
     system::{Query, SystemMeta},
     world::{FromWorld, World},
 };
@@ -171,6 +172,9 @@ pub unsafe trait SystemParam: Sized {
         world: &'world World,
         change_tick: u32,
     ) -> Self::Item<'world, 'state>;
+
+    #[allow(unused_variables)]
+    fn add_default_sets(sets: &mut Vec<BoxedSystemSet>) {}
 }
 
 /// A [`SystemParam`] that only reads a given [`World`].
@@ -237,6 +241,11 @@ unsafe impl<Q: WorldQuery + 'static, F: ReadOnlyWorldQuery + 'static> SystemPara
             false,
         )
     }
+
+    fn add_default_sets(sets: &mut Vec<BoxedSystemSet>) {
+        Q::add_default_sets(sets);
+        F::add_default_sets(sets);
+    }
 }
 
 fn assert_component_access_compatibility(
@@ -1416,6 +1425,11 @@ macro_rules! impl_system_param_tuple {
                 let ($($param,)*) = state;
                 ($($param::get_param($param, _system_meta, _world, _change_tick),)*)
             }
+
+            #[inline]
+            fn add_default_sets(_sets: &mut Vec<BoxedSystemSet>) {
+                $( $param::add_default_sets(_sets); )*
+            }
         }
     };
 }
@@ -1539,6 +1553,11 @@ unsafe impl<P: SystemParam + 'static> SystemParam for StaticSystemParam<'_, '_,
         // SAFETY: Defer to the safety of P::SystemParam
         StaticSystemParam(P::get_param(state, system_meta, world, change_tick))
     }
+
+    #[inline]
+    fn add_default_sets(sets: &mut Vec<BoxedSystemSet>) {
+        P::add_default_sets(sets);
+    }
 }
 
 #[cfg(test)]

crates/bevy_macro_utils/Cargo.toml

@@ -9,6 +9,7 @@ license = "MIT OR Apache-2.0"
 keywords = ["bevy"]
 
 [dependencies]
+proc-macro2 = "1"
 toml_edit = "0.19"
 syn = "1.0"
 quote = "1.0"