derive BufferKeyMap

Signed-off-by: Teo Koon Peng <[email protected]>

Author: koonpeng

macros/Cargo.toml

@@ -16,3 +16,4 @@ proc-macro = true
 [dependencies]
 syn = "2.0"
 quote = "1.0"
+proc-macro2 = "1.0.93"

macros/src/buffer.rs

@@ -6,15 +6,86 @@ use crate::Result;
 
 pub(crate) fn impl_joined_value(ast: DeriveInput) -> Result<TokenStream> {
     let struct_ident = ast.ident;
-    let (field_ident, field_type) = match ast.data {
-        syn::Data::Struct(data) => get_fields_map(data.fields)?,
+    let (field_ident, field_type): (Vec<Ident>, Vec<Type>) = match ast.data {
+        syn::Data::Struct(data) => get_fields_map(data.fields)?.into_iter().unzip(),
         _ => return Err("expected struct".to_string()),
     };
     let map_key: Vec<String> = field_ident.iter().map(|v| v.to_string()).collect();
     let struct_buffer_ident = format_ident!("__{}Buffers", struct_ident);
 
+    let impl_buffer_map_layout =
+        buffer_map_layout(&struct_ident, &field_ident, &field_type, &map_key);
+
     let gen = quote! {
-        impl BufferMapLayout for #struct_ident {
+        #impl_buffer_map_layout
+
+        impl ::bevy_impulse::JoinedValue for #struct_ident {
+            type Buffers = #struct_buffer_ident;
+
+            fn pull(
+                buffers: &::bevy_impulse::BufferMap,
+                session: ::bevy_ecs::prelude::Entity,
+                world: &mut ::bevy_ecs::prelude::World,
+            ) -> Result<Self, ::bevy_impulse::OperationError> {
+                use ::bevy_impulse::{ManageBuffer, OrBroken};
+
+                #(
+                    let #field_ident = world
+                        .get_entity_mut(buffers.get(#map_key).or_broken()?.id())
+                        .or_broken()?
+                        .pull_from_buffer::<#field_type>(session)?;
+                )*
+
+                Ok(Self {
+                    #(
+                        #field_ident
+                    ),*
+                })
+            }
+        }
+
+        struct #struct_buffer_ident {
+            #(
+                #field_ident: ::bevy_impulse::Buffer<#field_type>
+            ),*
+        }
+
+        impl From<#struct_buffer_ident> for ::bevy_impulse::BufferMap {
+            fn from(value: #struct_buffer_ident) -> Self {
+                let mut buffers = ::bevy_impulse::BufferMap::default();
+                #(
+                    buffers.insert(std::borrow::Cow::Borrowed(#map_key), value.#field_ident);
+                )*
+                buffers
+            }
+        }
+    };
+
+    Ok(gen.into())
+}
+
+fn get_fields_map(fields: syn::Fields) -> Result<Vec<(Ident, Type)>> {
+    match fields {
+        syn::Fields::Named(data) => {
+            let mut idents_types = Vec::with_capacity(data.named.len());
+            for field in data.named {
+                let ident = field.ident.ok_or("expected named fields".to_string())?;
+                idents_types.push((ident, field.ty));
+            }
+            Ok(idents_types)
+        }
+        _ => return Err("expected named fields".to_string()),
+    }
+}
+
+fn buffer_map_layout(
+    struct_ident: &Ident,
+    field_ident: &Vec<Ident>,
+    field_type: &Vec<Type>,
+    map_key: &Vec<String>,
+) -> proc_macro2::TokenStream {
+    quote! {
+        impl ::bevy_impulse::BufferMapLayout for #struct_ident {
             fn is_compatible(buffers: &BufferMap) -> Result<(), ::bevy_impulse::IncompatibleLayout> {
                 let mut compatibility = ::bevy_impulse::IncompatibleLayout::default();
                 #(
@@ -61,64 +132,57 @@ pub(crate) fn impl_joined_value(ast: DeriveInput) -> Result<TokenStream> {
                 Ok(())
             }
         }
+    }
+}
 
-        impl ::bevy_impulse::JoinedValue for #struct_ident {
-            type Buffers = #struct_buffer_ident;
+pub(crate) fn impl_buffer_key_map(ast: DeriveInput) -> Result<TokenStream> {
+    let struct_ident = ast.ident;
+    let (field_ident, field_type): (Vec<Ident>, Vec<Type>) = match ast.data {
+        syn::Data::Struct(data) => get_fields_map(data.fields)?.into_iter().unzip(),
+        _ => return Err("expected struct".to_string()),
+    };
+    let map_key: Vec<String> = field_ident.iter().map(|v| v.to_string()).collect();
 
-            fn pull(
-                buffers: &::bevy_impulse::BufferMap,
-                session: ::bevy_ecs::prelude::Entity,
-                world: &mut ::bevy_ecs::prelude::World,
-            ) -> Result<Self, ::bevy_impulse::OperationError> {
-                use ::bevy_impulse::{ManageBuffer, OrBroken};
+    let impl_buffer_map_layout =
+        buffer_map_layout(&struct_ident, &field_ident, &field_type, &map_key);
+
+    // FIXME(koonpeng): `create_key` does not allow failure and we can't guarantee that the buffer
+    // from buffer is valid.
+    let gen = quote! {
+        impl ::bevy_impulse::BufferKeyMap for #struct_ident {
+            fn add_accessor(buffers: &::bevy_impulse::BufferMap, accessor: ::bevy_ecs::prelude::Entity, world: &mut ::bevy_ecs::prelude::World) -> ::bevy_impulse::OperationResult {
+                use ::bevy_impulse::{Accessed, OrBroken};
 
                 #(
-                    let #field_ident = world
-                        .get_entity_mut(buffers.get(#map_key).or_broken()?.id())
-                        .or_broken()?
-                        .pull_from_buffer::<#field_type>(session)?;
+                    buffers.get(#map_key).or_broken()?.add_accessor(accessor, world)?;
                 )*
 
-                Ok(Self {
-                    #(
-                        #field_ident
-                    ),*
-                })
+                Ok(())
             }
-        }
 
-        struct #struct_buffer_ident {
-            #(
-                #field_ident: ::bevy_impulse::Buffer<#field_type>
-            ),*
-        }
+            fn create_key(buffers: &::bevy_impulse::BufferMap, builder: &::bevy_impulse::BufferKeyBuilder) -> Self {
+                use ::bevy_impulse::{Accessed, OrBroken};
 
-        impl From<#struct_buffer_ident> for ::bevy_impulse::BufferMap {
-            fn from(value: #struct_buffer_ident) -> Self {
-                let mut buffers = ::bevy_impulse::BufferMap::default();
+                Self {#(
+                    #field_ident: buffers.get(#map_key).unwrap().create_key(builder).try_into().unwrap(),
+                )*}
+            }
+
+            fn deep_clone_key(&self) -> Self {
+                Self {#(
+                    #field_ident: self.#field_ident.deep_clone(),
+                )*}
+            }
+
+            fn is_key_in_use(&self) -> bool {
                 #(
-                    buffers.insert(std::borrow::Cow::Borrowed(#map_key), value.#field_ident);
-                )*
-                buffers
+                    self.#field_ident.is_in_use()
+                )||*
             }
         }
+
+        #impl_buffer_map_layout
     };
 
     Ok(gen.into())
 }
-
-fn get_fields_map(fields: syn::Fields) -> Result<(Vec<Ident>, Vec<Type>)> {
-    match fields {
-        syn::Fields::Named(data) => {
-            let mut idents = Vec::with_capacity(data.named.len());
-            let mut types = Vec::with_capacity(data.named.len());
-            for field in data.named {
-                let ident = field.ident.ok_or("expected named fields".to_string())?;
-                idents.push(ident);
-                types.push(field.ty);
-            }
-            Ok((idents, types))
-        }
-        _ => return Err("expected named fields".to_string()),
-    }
-}

macros/src/lib.rs

@@ -16,7 +16,7 @@
 */
 
 mod buffer;
-use buffer::impl_joined_value;
+use buffer::{impl_buffer_key_map, impl_joined_value};
 
 use proc_macro::TokenStream;
 use quote::quote;
@@ -76,3 +76,15 @@ pub fn derive_joined_value(input: TokenStream) -> TokenStream {
         .into(),
     }
 }
+
+#[proc_macro_derive(BufferKeyMap)]
+pub fn derive_buffer_key_map(input: TokenStream) -> TokenStream {
+    let input = parse_macro_input!(input as DeriveInput);
+    match impl_buffer_key_map(input) {
+        Ok(tokens) => tokens,
+        Err(msg) => quote! {
+            compile_error!(#msg);
+        }
+        .into(),
+    }
+}

src/buffer/any_buffer.rs

@@ -176,11 +176,11 @@ impl AnyBufferKey {
         self.tag.session
     }
 
-    fn is_in_use(&self) -> bool {
+    pub fn is_in_use(&self) -> bool {
         self.tag.is_in_use()
     }
 
-    fn deep_clone(&self) -> Self {
+    pub fn deep_clone(&self) -> Self {
         Self {
             tag: self.tag.deep_clone(),
             interface: self.interface,
@@ -207,6 +207,14 @@ impl<T: 'static + Send + Sync + Any> From<BufferKey<T>> for AnyBufferKey {
     }
 }
 
+impl<T: 'static + Send + Sync + Any> TryFrom<AnyBufferKey> for BufferKey<T> {
+    type Error = OperationError;
+
+    fn try_from(value: AnyBufferKey) -> Result<Self, Self::Error> {
+        value.downcast_for_message().or_broken()
+    }
+}
+
 /// Similar to [`BufferView`][crate::BufferView], but this can be unlocked with
 /// an [`AnyBufferKey`], so it can work for any buffer whose message types
 /// support serialization and deserialization.

src/buffer/buffer_map.rs

@@ -29,7 +29,7 @@ use crate::{
     OperationRoster, Output, UnusedTarget,
 };
 
-pub use bevy_impulse_derive::JoinedValue;
+pub use bevy_impulse_derive::{BufferKeyMap, JoinedValue};
 
 #[derive(Clone, Default)]
 pub struct BufferMap {
@@ -409,4 +409,12 @@ mod tests {
         assert_eq!(value.string, "hello");
         assert!(context.no_unhandled_errors());
     }
+
+    // TODO(koonpeng): Add test that uses `TestBufferKeyMap`.
+    #[derive(Clone, BufferKeyMap)]
+    struct TestBufferKeyMap {
+        integer: BufferKey<i64>,
+        float: BufferKey<f64>,
+        string: AnyBufferKey,
+    }
 }