diff --git a/crates/bevy_ecs/macros/src/component.rs b/crates/bevy_ecs/macros/src/component.rs index ef7fad99f4..040c1b26b6 100644 --- a/crates/bevy_ecs/macros/src/component.rs +++ b/crates/bevy_ecs/macros/src/component.rs @@ -128,9 +128,11 @@ pub fn derive_component(input: TokenStream) -> TokenStream { let map_entities = map_entities( &ast.data, + &bevy_ecs_path, Ident::new("this", Span::call_site()), relationship.is_some(), relationship_target.is_some(), + attrs.map_entities ).map(|map_entities_impl| quote! { fn map_entities(this: &mut Self, mapper: &mut M) { use #bevy_ecs_path::entity::MapEntities; @@ -339,10 +341,19 @@ const ENTITIES: &str = "entities"; pub(crate) fn map_entities( data: &Data, + bevy_ecs_path: &Path, self_ident: Ident, is_relationship: bool, is_relationship_target: bool, + map_entities_attr: Option, ) -> Option { + if let Some(map_entities_override) = map_entities_attr { + let map_entities_tokens = map_entities_override.to_token_stream(bevy_ecs_path); + return Some(quote!( + #map_entities_tokens(#self_ident, mapper) + )); + } + match data { Data::Struct(DataStruct { fields, .. }) => { let mut map = Vec::with_capacity(fields.len()); @@ -430,6 +441,7 @@ pub const ON_INSERT: &str = "on_insert"; pub const ON_REPLACE: &str = "on_replace"; pub const ON_REMOVE: &str = "on_remove"; pub const ON_DESPAWN: &str = "on_despawn"; +pub const MAP_ENTITIES: &str = "map_entities"; pub const IMMUTABLE: &str = "immutable"; pub const CLONE_BEHAVIOR: &str = "clone_behavior"; @@ -484,6 +496,56 @@ impl Parse for HookAttributeKind { } } +#[derive(Debug)] +pub(super) enum MapEntitiesAttributeKind { + /// expressions like function or struct names + /// + /// structs will throw compile errors on the code generation so this is safe + Path(ExprPath), + /// When no value is specified + Default, +} + +impl MapEntitiesAttributeKind { + fn from_expr(value: Expr) -> Result { + match value { + Expr::Path(path) => Ok(Self::Path(path)), + // throw meaningful error on all other expressions + _ => Err(syn::Error::new( + value.span(), + [ + "Not supported in this position, please use one of the following:", + "- path to function", + "- nothing to default to MapEntities implementation", + ] + .join("\n"), + )), + } + } + + fn to_token_stream(&self, bevy_ecs_path: &Path) -> TokenStream2 { + match self { + MapEntitiesAttributeKind::Path(path) => path.to_token_stream(), + MapEntitiesAttributeKind::Default => { + quote!( + ::map_entities + ) + } + } + } +} + +impl Parse for MapEntitiesAttributeKind { + fn parse(input: syn::parse::ParseStream) -> Result { + if input.peek(Token![=]) { + input.parse::()?; + input.parse::().and_then(Self::from_expr) + } else { + Ok(Self::Default) + } + } +} + struct Attrs { storage: StorageTy, requires: Option>, @@ -496,6 +558,7 @@ struct Attrs { relationship_target: Option, immutable: bool, clone_behavior: Option, + map_entities: Option, } #[derive(Clone, Copy)] @@ -535,6 +598,7 @@ fn parse_component_attr(ast: &DeriveInput) -> Result { relationship_target: None, immutable: false, clone_behavior: None, + map_entities: None, }; let mut require_paths = HashSet::new(); @@ -573,6 +637,9 @@ fn parse_component_attr(ast: &DeriveInput) -> Result { } else if nested.path.is_ident(CLONE_BEHAVIOR) { attrs.clone_behavior = Some(nested.value()?.parse()?); Ok(()) + } else if nested.path.is_ident(MAP_ENTITIES) { + attrs.map_entities = Some(nested.input.parse::()?); + Ok(()) } else { Err(nested.error("Unsupported attribute")) } diff --git a/crates/bevy_ecs/macros/src/lib.rs b/crates/bevy_ecs/macros/src/lib.rs index 7750f97259..9bc3e5913e 100644 --- a/crates/bevy_ecs/macros/src/lib.rs +++ b/crates/bevy_ecs/macros/src/lib.rs @@ -220,9 +220,11 @@ pub fn derive_map_entities(input: TokenStream) -> TokenStream { let map_entities_impl = map_entities( &ast.data, + &ecs_path, Ident::new("self", Span::call_site()), false, false, + None, ); let struct_name = &ast.ident; diff --git a/crates/bevy_ecs/src/component.rs b/crates/bevy_ecs/src/component.rs index 831402240d..615c5903f8 100644 --- a/crates/bevy_ecs/src/component.rs +++ b/crates/bevy_ecs/src/component.rs @@ -578,6 +578,65 @@ pub trait Component: Send + Sync + 'static { /// items: Vec> /// } /// ``` + /// + /// You might need more specialized logic. A likely cause of this is your component contains collections of entities that + /// don't implement [`MapEntities`](crate::entity::MapEntities). In that case, you can annotate your component with + /// `#[component(map_entities)]`. Using this attribute, you must implement `MapEntities` for the + /// component itself, and this method will simply call that implementation. + /// + /// ``` + /// # use bevy_ecs::{component::Component, entity::{Entity, MapEntities, EntityMapper}}; + /// # use std::collections::HashMap; + /// #[derive(Component)] + /// #[component(map_entities)] + /// struct Inventory { + /// items: HashMap + /// } + /// + /// impl MapEntities for Inventory { + /// fn map_entities(&mut self, entity_mapper: &mut M) { + /// self.items = self.items + /// .drain() + /// .map(|(id, count)|(entity_mapper.get_mapped(id), count)) + /// .collect(); + /// } + /// } + /// # let a = Entity::from_bits(0x1_0000_0001); + /// # let b = Entity::from_bits(0x1_0000_0002); + /// # let mut inv = Inventory { items: Default::default() }; + /// # inv.items.insert(a, 10); + /// # ::map_entities(&mut inv, &mut (a,b)); + /// # assert_eq!(inv.items.get(&b), Some(&10)); + /// ```` + /// + /// Alternatively, you can specify the path to a function with `#[component(map_entities = function_path)]`, similar to component hooks. + /// In this case, the inputs of the function should mirror the inputs to this method, with the second parameter being generic. + /// + /// ``` + /// # use bevy_ecs::{component::Component, entity::{Entity, MapEntities, EntityMapper}}; + /// # use std::collections::HashMap; + /// #[derive(Component)] + /// #[component(map_entities = map_the_map)] + /// // Also works: map_the_map:: or map_the_map::<_> + /// struct Inventory { + /// items: HashMap + /// } + /// + /// fn map_the_map(inv: &mut Inventory, entity_mapper: &mut M) { + /// inv.items = inv.items + /// .drain() + /// .map(|(id, count)|(entity_mapper.get_mapped(id), count)) + /// .collect(); + /// } + /// # let a = Entity::from_bits(0x1_0000_0001); + /// # let b = Entity::from_bits(0x1_0000_0002); + /// # let mut inv = Inventory { items: Default::default() }; + /// # inv.items.insert(a, 10); + /// # ::map_entities(&mut inv, &mut (a,b)); + /// # assert_eq!(inv.items.get(&b), Some(&10)); + /// ```` + /// + /// You can use the turbofish (`::`) to specify parameters when a function is generic, using either M or _ for the type of the mapper parameter. #[inline] fn map_entities(_this: &mut Self, _mapper: &mut E) {} }