1 files changed, 45 insertions, 7 deletions

diff --git a/ecs/src/query.rs b/ecs/src/query.rs
index 1d27b7a..1947c18 100644
--- a/ecs/src/query.rs
+++ b/ecs/src/query.rs
@@ -3,7 +3,12 @@ use std::marker::PhantomData;
 
 use seq_macro::seq;
 
-use crate::component::{Component, HandleFromEntityComponentRef, Ref as ComponentRef};
+use crate::component::{
+    Component,
+    Handle as ComponentHandle,
+    HandleFromEntityComponentRef,
+    HandleMut as ComponentHandleMut,
+};
 use crate::entity::Handle as EntityHandle;
 use crate::query::flexible::{Iter as FlexibleQueryIter, Query as FlexibleQuery};
 use crate::system::{Param as SystemParam, System};
@@ -321,15 +326,15 @@ pub trait TermWithField
     ) -> Self::Field<'world>;
 }
 
-impl<ComponentRefT: ComponentRef> TermWithField for ComponentRefT
+impl<ComponentT: Component> TermWithField for &ComponentT
 {
-    type Field<'a> = ComponentRefT::Handle<'a>;
+    type Field<'a> = ComponentHandle<'a, ComponentT>;
 
     fn apply_to_terms_builder<const MAX_TERM_CNT: usize>(
         terms_builder: &mut TermsBuilder<MAX_TERM_CNT>,
     )
     {
-        terms_builder.with::<ComponentRefT::Component>();
+        terms_builder.with::<ComponentT>();
     }
 
     fn get_field<'world>(
@@ -337,18 +342,51 @@ impl<ComponentRefT: ComponentRef> TermWithField for ComponentRefT
         world: &'world World,
     ) -> Self::Field<'world>
     {
-        assert_eq!(ComponentRefT::Component::id().kind(), UidKind::Component);
+        assert_eq!(ComponentT::id().kind(), UidKind::Component);
 
         Self::Field::from_entity_component_ref(
             entity_handle
-                .get_matching_components(ComponentRefT::Component::id())
+                .get_matching_components(ComponentT::id())
                 .next(),
             world,
         )
         .unwrap_or_else(|err| {
             panic!(
                 "Creating handle to component {} failed: {err}",
-                type_name::<ComponentRefT::Component>()
+                type_name::<ComponentT>()
+            );
+        })
+    }
+}
+
+impl<ComponentT: Component> TermWithField for &mut ComponentT
+{
+    type Field<'a> = ComponentHandleMut<'a, ComponentT>;
+
+    fn apply_to_terms_builder<const MAX_TERM_CNT: usize>(
+        terms_builder: &mut TermsBuilder<MAX_TERM_CNT>,
+    )
+    {
+        terms_builder.with::<ComponentT>();
+    }
+
+    fn get_field<'world>(
+        entity_handle: &EntityHandle<'world>,
+        world: &'world World,
+    ) -> Self::Field<'world>
+    {
+        assert_eq!(ComponentT::id().kind(), UidKind::Component);
+
+        Self::Field::from_entity_component_ref(
+            entity_handle
+                .get_matching_components(ComponentT::id())
+                .next(),
+            world,
+        )
+        .unwrap_or_else(|err| {
+            panic!(
+                "Creating handle to component {} failed: {err}",
+                type_name::<ComponentT>()
             );
         })
     }