Round out Mut<T> api using lenses + auto-unwrapping

crates/bevy_ecs/macros/src/component.rs

@@ -2,7 +2,9 @@ use bevy_macro_utils::{get_lit_str, Symbol};
 use proc_macro::TokenStream;
 use proc_macro2::{Span, TokenStream as TokenStream2};
 use quote::{quote, ToTokens};
-use syn::{parse_macro_input, parse_quote, DeriveInput, Error, Ident, Path, Result};
+use syn::{
+    parse_macro_input, parse_quote, DataStruct, DeriveInput, Error, Fields, Ident, Path, Result,
+};
 
 pub fn derive_component(input: TokenStream) -> TokenStream {
     let mut ast = parse_macro_input!(input as DeriveInput);
@@ -23,18 +25,61 @@ pub fn derive_component(input: TokenStream) -> TokenStream {
     let struct_name = &ast.ident;
     let (impl_generics, type_generics, where_clause) = &ast.generics.split_for_impl();
 
+    let (lens_type, ext) = match attrs.lens {
+        Lens::None => (quote! {#bevy_ecs_path::lens::NoopLens<Self>}, None),
+        Lens::Declared(decl) => (decl, None),
+        Lens::Derived => {
+            let inner_type = match &ast.data {
+                syn::Data::Struct(DataStruct {
+                    fields: Fields::Unnamed(f),
+                    ..
+                }) if f.unnamed.len() == 1 => &f.unnamed[0].ty,
+                _ => return quote! {
+                    compile_error!("Automatic lensing is only available for unit structs with one element")
+                }.into(),
+            };
+            (
+                quote! {Self},
+                Some(quote! {
+                    impl #impl_generics #bevy_ecs_path::lens::Lens for #struct_name #type_generics #where_clause {
+                        type In = Self;
+                        type Out = #inner_type;
+
+                        fn get(input: &Self::In) -> &Self::Out {
+                            &input.0
+                        }
+                        fn get_mut(input: &mut Self::In) -> &mut Self::Out {
+                            &mut input.0
+                        }
+                    }
+                }),
+            )
+        }
+    };
+
     TokenStream::from(quote! {
+        #ext
+
         impl #impl_generics #bevy_ecs_path::component::Component for #struct_name #type_generics #where_clause {
             type Storage = #storage;
+            type DefaultLens = #lens_type;
         }
     })
 }
 
 pub const COMPONENT: Symbol = Symbol("component");
 pub const STORAGE: Symbol = Symbol("storage");
+pub const LENS: Symbol = Symbol("lens");
 
 struct Attrs {
     storage: StorageTy,
+    lens: Lens,
+}
+
+enum Lens {
+    None,
+    Declared(TokenStream2),
+    Derived,
 }
 
 #[derive(Clone, Copy)]
@@ -52,11 +97,12 @@ fn parse_component_attr(ast: &DeriveInput) -> Result<Attrs> {
 
     let mut attrs = Attrs {
         storage: StorageTy::Table,
+        lens: Lens::None,
     };
 
     for meta in meta_items {
         use syn::{
-            Meta::NameValue,
+            Meta::{NameValue, Path},
             NestedMeta::{Lit, Meta},
         };
         match meta {
@@ -75,6 +121,16 @@ fn parse_component_attr(ast: &DeriveInput) -> Result<Attrs> {
                     }
                 };
             }
+            Meta(NameValue(m)) if m.path == LENS => {
+                attrs.lens = Lens::Declared(get_lit_str(LENS, &m.lit)?.value().as_str().parse()?);
+            }
+            Meta(Path(p))
+                if p.leading_colon.is_none()
+                    && p.segments.len() == 1
+                    && p.segments[0].ident == LENS =>
+            {
+                attrs.lens = Lens::Derived;
+            }
             Meta(meta_item) => {
                 return Err(Error::new_spanned(
                     meta_item.path(),

crates/bevy_ecs/src/change_detection.rs

@@ -1,6 +1,6 @@
 //! Types that detect when their internal data mutate.
 
-use crate::{component::ComponentTicks, system::Resource};
+use crate::{component::ComponentTicks, lens::Lens, system::Resource};
 #[cfg(feature = "bevy_reflect")]
 use bevy_reflect::Reflect;
 use std::ops::{Deref, DerefMut};
@@ -116,6 +116,32 @@ macro_rules! impl_into_inner {
     };
 }
 
+macro_rules! impl_lens {
+    ($name:ident < $( $generics:tt ),+ >, $target:ty, $($traits:ident)?) => {
+        impl<$($generics),* $(: $traits)?> $name<$($generics),*> {
+            #[inline]
+            pub fn lens<L: Lens<In = $target>>(self) -> Mut<'a, L::Out> {
+                Mut {
+                    value: L::get_mut(self.value),
+                    ticks: self.ticks,
+                }
+            }
+
+            #[inline]
+            pub fn lens_borrow<'b, L: Lens<In = $target>>(&'b mut self) -> Mut<'b, L::Out> {
+                Mut {
+                    value: L::get_mut(self.value),
+                    ticks: Ticks {
+                        component_ticks: self.ticks.component_ticks,
+                        last_change_tick: self.ticks.last_change_tick,
+                        change_tick: self.ticks.change_tick,
+                    },
+                }
+            }
+        }
+    };
+}
+
 macro_rules! impl_debug {
     ($name:ident < $( $generics:tt ),+ >, $($traits:ident)?) => {
         impl<$($generics),* $(: $traits)?> std::fmt::Debug for $name<$($generics),*>
@@ -154,6 +180,7 @@ pub struct ResMut<'a, T: Resource> {
 }
 
 change_detection_impl!(ResMut<'a, T>, T, Resource);
+impl_lens!(ResMut<'a, T>, T, Resource);
 impl_into_inner!(ResMut<'a, T>, T, Resource);
 impl_debug!(ResMut<'a, T>, Resource);
 
@@ -175,6 +202,7 @@ pub struct NonSendMut<'a, T: 'static> {
 }
 
 change_detection_impl!(NonSendMut<'a, T>, T,);
+impl_lens!(NonSendMut<'a, T>, T,);
 impl_into_inner!(NonSendMut<'a, T>, T,);
 impl_debug!(NonSendMut<'a, T>,);
 
@@ -185,6 +213,7 @@ pub struct Mut<'a, T> {
 }
 
 change_detection_impl!(Mut<'a, T>, T,);
+impl_lens!(Mut<'a, T>, T,);
 impl_into_inner!(Mut<'a, T>, T,);
 impl_debug!(Mut<'a, T>,);
 
@@ -199,3 +228,70 @@ pub struct ReflectMut<'a> {
 change_detection_impl!(ReflectMut<'a>, dyn Reflect,);
 #[cfg(feature = "bevy_reflect")]
 impl_into_inner!(ReflectMut<'a>, dyn Reflect,);
+
+mod ops_passthrough {
+    use std::ops::*;
+
+    use super::Mut;
+
+    macro_rules! binary_ops {
+        ($trait:ident, $method:ident) => {
+            impl<Rhs, T: $trait<Rhs> + Copy> $trait<Rhs> for Mut<'_, T> {
+                type Output = T::Output;
+
+                fn $method(self, rhs: Rhs) -> Self::Output {
+                    T::$method(*self, rhs)
+                }
+            }
+        };
+    }
+    macro_rules! unary_ops {
+        ($trait:ident, $method:ident) => {
+            impl<T: $trait + Copy> $trait for Mut<'_, T> {
+                type Output = T::Output;
+
+                fn $method(self) -> Self::Output {
+                    T::$method(*self)
+                }
+            }
+        };
+    }
+    macro_rules! assign_ops {
+        ($trait:ident, $method:ident) => {
+            impl<Rhs, T: $trait<Rhs>> $trait<Rhs> for Mut<'_, T> {
+                fn $method(&mut self, rhs: Rhs) {
+                    T::$method(&mut *self, rhs)
+                }
+            }
+        };
+    }
+
+    binary_ops!(Add, add);
+    binary_ops!(Sub, sub);
+    binary_ops!(Mul, mul);
+    binary_ops!(Div, div);
+
+    binary_ops!(Rem, rem);
+
+    binary_ops!(BitAnd, bitand);
+    binary_ops!(BitOr, bitor);
+    binary_ops!(BitXor, bitxor);
+    binary_ops!(Shr, shr);
+    binary_ops!(Shl, shl);
+
+    unary_ops!(Neg, neg);
+    unary_ops!(Not, not);
+
+    assign_ops!(AddAssign, add_assign);
+    assign_ops!(SubAssign, sub_assign);
+    assign_ops!(MulAssign, mul_assign);
+    assign_ops!(DivAssign, div_assign);
+
+    assign_ops!(RemAssign, rem_assign);
+
+    assign_ops!(BitAndAssign, bitand_assign);
+    assign_ops!(BitOrAssign, bitor_assign);
+    assign_ops!(BitXorAssign, bitxor_assign);
+    assign_ops!(ShrAssign, shr_assign);
+    assign_ops!(ShlAssign, shl_assign);
+}

crates/bevy_ecs/src/component.rs

@@ -1,6 +1,7 @@
 //! Types for declaring and storing [`Component`]s.
 
 use crate::{
+    lens::{Lens, NoopLens},
     storage::{SparseSetIndex, Storages},
     system::Resource,
 };
@@ -32,6 +33,8 @@ use std::{
 /// Components can be grouped together into a [`Bundle`](crate::bundle::Bundle).
 pub trait Component: Send + Sync + 'static {
     type Storage: ComponentStorage;
+
+    type DefaultLens: Lens<In = Self>;
 }
 
 pub struct TableStorage;
@@ -61,6 +64,7 @@ where
     Self: Send + Sync + 'static,
 {
     type Storage = TableStorage;
+    type DefaultLens = NoopLens<Self>;
 }
 
 /// The storage used for a specific component type.

crates/bevy_ecs/src/lens.rs

@@ -0,0 +1,76 @@
+use std::marker::PhantomData;
+
+pub trait Lens {
+    type In: 'static;
+    type Out: 'static;
+
+    fn get(input: &Self::In) -> &Self::Out;
+    fn get_mut(input: &mut Self::In) -> &mut Self::Out;
+}
+
+pub struct ComposedLens<A, B>(PhantomData<(A, B)>);
+
+impl<A, B> Lens for ComposedLens<A, B>
+where
+    A: Lens,
+    B: Lens<In = A::Out>,
+{
+    type In = A::In;
+    type Out = B::Out;
+
+    fn get(input: &Self::In) -> &Self::Out {
+        B::get(A::get(input))
+    }
+
+    fn get_mut(input: &mut Self::In) -> &mut Self::Out {
+        B::get_mut(A::get_mut(input))
+    }
+}
+
+pub struct NoopLens<T>(PhantomData<T>);
+
+impl<T: 'static> Lens for NoopLens<T> {
+    type In = T;
+    type Out = T;
+
+    fn get(input: &Self::In) -> &Self::Out {
+        input
+    }
+
+    fn get_mut(input: &mut Self::In) -> &mut Self::Out {
+        input
+    }
+}
+
+#[macro_export]
+macro_rules! declare_lens {
+    ($name:ident, $in:ty, $out:ty, $($path:tt).*) => {
+        struct $name;
+
+        impl $crate::lens::Lens for $name {
+            type In = $in;
+            type Out = $out;
+
+            fn get(input: &Self::In) -> &Self::Out {
+                &input$(.$path)*
+            }
+
+            fn get_mut(input: &mut Self::In) -> &mut Self::Out {
+                &mut input$(.$path)*
+            }
+        }
+    };
+}
+
+#[macro_export]
+macro_rules! composed_lens {
+    ($last:ty) => {
+        $last
+    };
+    ($first:ty, $second:ty) => {
+        $crate::lens::ComposedLens<$first, $second>
+    };
+    ($first:ty, $second:ty, $($rest:ty),*) => {
+        $crate::lens::ComposedLens<composed_lens![$first, $second], composed_lens![$($rest),*]>
+    };
+}

crates/bevy_ecs/src/lib.rs

@@ -6,6 +6,7 @@ pub mod change_detection;
 pub mod component;
 pub mod entity;
 pub mod event;
+pub mod lens;
 pub mod query;
 #[cfg(feature = "bevy_reflect")]
 pub mod reflect;
@@ -26,7 +27,9 @@ pub mod prelude {
         component::Component,
         entity::Entity,
         event::{EventReader, EventWriter},
-        query::{Added, AnyOf, ChangeTrackers, Changed, Or, QueryState, With, Without},
+        query::{
+            Added, AnyOf, ChangeTrackers, Changed, Or, QueryState, Raw, RawMut, With, Without,
+        },
         schedule::{
             AmbiguitySetLabel, ExclusiveSystemDescriptorCoercion, ParallelSystemDescriptorCoercion,
             RunCriteria, RunCriteriaDescriptorCoercion, RunCriteriaLabel, RunCriteriaPiping,