Skip to content

Derive Macro for MapEntites (#17611) #18311

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 8 commits into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions crates/bevy_ecs/macros/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
extern crate proc_macro;

mod component;
mod map_entities;
mod query_data;
mod query_filter;
mod states;
Expand Down Expand Up @@ -613,6 +614,11 @@ pub fn derive_substates(input: TokenStream) -> TokenStream {
states::derive_substates(input)
}

#[proc_macro_derive(MapEntities, attributes(skip_mapping))]
pub fn derive_map_entities(input: TokenStream) -> TokenStream {
map_entities::derive_map_entities(input)
}

#[proc_macro_derive(FromWorld, attributes(from_world))]
pub fn derive_from_world(input: TokenStream) -> TokenStream {
let bevy_ecs_path = bevy_ecs_path();
Expand Down
148 changes: 148 additions & 0 deletions crates/bevy_ecs/macros/src/map_entities.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,148 @@
use proc_macro::TokenStream;
use quote::quote;
use syn::{parse_macro_input, Data, DeriveInput, Fields};

pub(super) fn derive_map_entities(input: TokenStream) -> TokenStream {
let input = parse_macro_input!(input as DeriveInput);
let name = input.ident.clone();

match input.data {
Data::Struct(data_struct) => map_struct(name, &data_struct.fields),
Data::Enum(data_enum) => map_enum(name, &data_enum.variants),
Data::Union(_) => panic!("MapEntities is only valid on structs or enums"),
}
.into()
}

fn map_struct(name: syn::Ident, fields: &Fields) -> proc_macro2::TokenStream {
let mut map_entities = vec![];

for (i, field) in fields.iter().enumerate() {
if skip_mapping(&field.attrs) {
continue;
}

let ty = &field.ty;
let member = match &field.ident {
// a named field (struct)
Some(field_name) => syn::Member::Named(field_name.clone()),
// Unnamed field (tuple-like struct)
None => syn::Member::Unnamed(syn::Index::from(i)),
};

let map_field = quote! {
<#ty as bevy_ecs::entity::MapEntities>::map_entities(&mut self.#member, entity_mapper);
};

map_entities.push(map_field);
}

quote! {
impl MapEntities for #name {
fn map_entities<M: bevy_ecs::entity::EntityMapper>(&mut self, entity_mapper: &mut M) {
#(#map_entities)*
}
}
}
}

// fn map_enum(name:syn::Ident, variantes: &[])

fn map_enum(
name: syn::Ident,
variants: &syn::punctuated::Punctuated<syn::Variant, syn::token::Comma>,
) -> proc_macro2::TokenStream {
let variants = variants.iter().map(|variant| {
let variant_name = &variant.ident;
let pattern = get_enum_pattern(&variant.fields);
let map_entities = map_enum_variant(&variant.fields);

quote! {
#name::#variant_name #pattern => {
#map_entities
}
}
});

quote! {
impl MapEntities for #name {
fn map_entities<M: bevy_ecs::entity::EntityMapper>(&mut self, entity_mapper: &mut M) {
match self {
#(#variants)*
// allows empty enums to derive aswell
_ => unreachable!()
}
}
}
}
}

/// Generates the match pattern for unnamed fields (e.g., `(field_0, field_1)`)
///
/// Or named fields (e.g., `(item, quantity)`)
fn get_enum_pattern(fields: &Fields) -> proc_macro2::TokenStream {
match fields {
Fields::Named(named_fields) => {
let field_vars = named_fields.named.iter().map(|field| {
let var_name = field.ident.as_ref().unwrap();
quote! { #var_name }
});
quote! { {#(#field_vars),*} }
}
Fields::Unnamed(unnamed_fields) => {
let field_vars = (0..unnamed_fields.unnamed.len()).map(|i| {
let var_name =
syn::Ident::new(&format!("field_{}", i), proc_macro2::Span::call_site());
quote! { #var_name }
});
quote! { ( #(#field_vars),* ) }
}
Fields::Unit => quote! {},
}
}

// Function to generate code to map fields for an enum variant
fn map_enum_variant(fields: &Fields) -> proc_macro2::TokenStream {
let mut map_entities = vec![];

for (i, field) in fields.iter().enumerate() {
if skip_mapping(&field.attrs) {
continue;
}

let ty = &field.ty;
let member = match &field.ident {
// named field
Some(field_name) => field_name.clone(),
// Unnamed field
None => syn::Ident::new(&format!("field_{}", i), proc_macro2::Span::call_site()),
};

let map_field = quote! {
<#ty as bevy_ecs::entity::MapEntities>::map_entities(#member, entity_mapper);
};

map_entities.push(map_field);
}

quote! {
#(#map_entities)*
}
}

/// check if any attribute contains `skip_mapping`
fn skip_mapping(attrs: &[syn::Attribute]) -> bool {
for attr in attrs {
if !attr.path().is_ident("skip_mapping") {
continue;
}
match &attr.meta {
// skip that fields
syn::Meta::Path(_) => return true,
// the attribute has been used incorrectly
_ => panic!("just use the bare `#[skip_mapping]`"),
}
}

false
}
56 changes: 53 additions & 3 deletions crates/bevy_ecs/src/entity/map_entities.rs
Original file line number Diff line number Diff line change
@@ -1,10 +1,10 @@
use super::{hash_map::EntityHashMap, VisitEntitiesMut};
use crate::{
entity::Entity,
identifier::masks::{IdentifierMask, HIGH_MASK},
world::World,
};

use super::{hash_map::EntityHashMap, VisitEntitiesMut};
pub use bevy_ecs_macros::MapEntities;

/// Operation to map all contained [`Entity`] fields in a type to new values.
///
Expand Down Expand Up @@ -44,6 +44,35 @@ use super::{hash_map::EntityHashMap, VisitEntitiesMut};
/// }
/// }
/// ```
///
/// a derive macro is available to reduce boilerplate
/// ```
/// use bevy_ecs::prelude::*;
/// use bevy_ecs::entity::MapEntities;
///
/// #[derive(MapEntities)]
/// pub struct A;
/// pub struct B;
///
/// #[derive(MapEntities)]
/// struct MyStruct {
/// a: A,
/// #[skip_mapping]
/// b: B,
/// }
///
/// # #[derive(MapEntities)]
/// # struct TestTuple(A);

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

An extra # here

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

to eliminate that double newline?

/// #
/// /* the above derive macro is equivalent to this
/// impl MapEntities for MyStruct {
/// fn map_entities<M: EntityMapper>(&mut self, entity_mapper: &mut M) {
/// self.a.map_entities(entity_mapper);
/// // self.b is skipped because of the #[skip_mapping] attribute
/// }
/// }
/// */
/// ```
pub trait MapEntities {
/// Updates all [`Entity`] references stored inside using `entity_mapper`.
///
Expand Down Expand Up @@ -247,7 +276,7 @@ impl<'m> SceneEntityMapper<'m> {
#[cfg(test)]
mod tests {
use crate::{
entity::{hash_map::EntityHashMap, Entity, EntityMapper, SceneEntityMapper},
entity::{hash_map::EntityHashMap, Entity, EntityMapper, MapEntities, SceneEntityMapper},
world::World,
};

Expand Down Expand Up @@ -312,4 +341,25 @@ mod tests {
// The SceneEntityMapper should leave `Entities` in a flushed state.
assert!(!world.entities.needs_flush());
}

#[test]
#[expect(dead_code, reason = "just test if it compiles")]
fn entity_mapper_macro_on_enum() {
#[derive(MapEntities)]
enum EmptyEnum {}

#[derive(MapEntities)]
struct A;
struct B;

#[derive(MapEntities)]
enum TestEnum {
Unnamed(A, #[skip_mapping] B),
Named {
a: A,
#[skip_mapping]
_b: B,
},
}
}
}