BufferKeyMap macro (#56)

Signed-off-by: Michael X. Grey <[email protected]>

Author: mxgrey

macros/src/buffer.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, attributes(key))]
+pub fn derive_buffer_key_map(input: TokenStream) -> TokenStream {
+    let input = parse_macro_input!(input as ItemStruct);
+    match impl_buffer_key_map(&input) {
+        Ok(tokens) => tokens.into(),
+        Err(msg) => quote! {
+            compile_error!(#msg);
+        }
+        .into(),
+    }
+}

macros/src/lib.rs

@@ -17,9 +17,9 @@
 
 use bevy_ecs::{
     change_detection::Mut,
-    prelude::{Commands, Entity, Query},
+    prelude::{Commands, Entity, Query, World},
     query::QueryEntityError,
-    system::SystemParam,
+    system::{SystemParam, SystemState},
 };
 
 use std::{ops::RangeBounds, sync::Arc};
@@ -38,7 +38,7 @@ pub use buffer_access_lifecycle::BufferKeyLifecycle;
 pub(crate) use buffer_access_lifecycle::*;
 
 mod buffer_key_builder;
-pub(crate) use buffer_key_builder::*;
+pub use buffer_key_builder::*;
 
 mod buffer_map;
 pub use buffer_map::*;
@@ -402,6 +402,67 @@ where
     }
 }
 
+/// This trait allows [`World`] to give you access to any buffer using a [`BufferKey`]
+pub trait BufferWorldAccess {
+    /// Call this to get read-only access to a buffer from a [`World`].
+    ///
+    /// Alternatively you can use [`BufferAccess`] as a regular bevy system parameter,
+    /// which does not need direct world access.
+    fn buffer_view<T>(&self, key: &BufferKey<T>) -> Result<BufferView<'_, T>, BufferError>
+    where
+        T: 'static + Send + Sync;
+
+    /// Call this to get mutable access to a buffer.
+    ///
+    /// Pass in a callback that will receive [`BufferMut`], allowing it to view
+    /// and modify the contents of the buffer.
+    fn buffer_mut<T, U>(
+        &mut self,
+        key: &BufferKey<T>,
+        f: impl FnOnce(BufferMut<T>) -> U,
+    ) -> Result<U, BufferError>
+    where
+        T: 'static + Send + Sync;
+}
+
+impl BufferWorldAccess for World {
+    fn buffer_view<T>(&self, key: &BufferKey<T>) -> Result<BufferView<'_, T>, BufferError>
+    where
+        T: 'static + Send + Sync,
+    {
+        let buffer_ref = self
+            .get_entity(key.tag.buffer)
+            .ok_or(BufferError::BufferMissing)?;
+        let storage = buffer_ref
+            .get::<BufferStorage<T>>()
+            .ok_or(BufferError::BufferMissing)?;
+        let gate = buffer_ref
+            .get::<GateState>()
+            .ok_or(BufferError::BufferMissing)?;
+        Ok(BufferView {
+            storage,
+            gate,
+            session: key.tag.session,
+        })
+    }
+
+    fn buffer_mut<T, U>(
+        &mut self,
+        key: &BufferKey<T>,
+        f: impl FnOnce(BufferMut<T>) -> U,
+    ) -> Result<U, BufferError>
+    where
+        T: 'static + Send + Sync,
+    {
+        let mut state = SystemState::<BufferAccessMut<T>>::new(self);
+        let mut buffer_access_mut = state.get_mut(self);
+        let buffer_mut = buffer_access_mut
+            .get_mut(key)
+            .map_err(|_| BufferError::BufferMissing)?;
+        Ok(f(buffer_mut))
+    }
+}
+
 /// Access to view a buffer that exists inside a workflow.
 pub struct BufferView<'a, T>
 where

src/buffer.rs

@@ -517,7 +517,7 @@ pub trait AnyBufferWorldAccess {
     /// For technical reasons this requires direct [`World`] access, but you can
     /// do other read-only queries on the world while holding onto the
     /// [`AnyBufferView`].
-    fn any_buffer_view<'a>(&self, key: &AnyBufferKey) -> Result<AnyBufferView<'_>, BufferError>;
+    fn any_buffer_view(&self, key: &AnyBufferKey) -> Result<AnyBufferView<'_>, BufferError>;
 
     /// Call this to get mutable access to any buffer.
     ///
@@ -531,7 +531,7 @@ pub trait AnyBufferWorldAccess {
 }
 
 impl AnyBufferWorldAccess for World {
-    fn any_buffer_view<'a>(&self, key: &AnyBufferKey) -> Result<AnyBufferView<'_>, BufferError> {
+    fn any_buffer_view(&self, key: &AnyBufferKey) -> Result<AnyBufferView<'_>, BufferError> {
         key.interface.create_any_buffer_view(key, self)
     }