wip: QueryOne

ten3roberts · Oct 11, 2023 · fa04900 · fa04900
1 parent dacb5e1
commit fa04900
Show file tree

Hide file tree

Showing 5 changed files with 68 additions and 35 deletions.
diff --git a/src/entity_ref.rs b/src/entity_ref.rs
@@ -15,7 +15,7 @@ use crate::{
     name,
     query::QueryOne,
     writer::{EntityWriter, FnWriter, Missing, Replace, SingleComponentWriter, WriteDedup},
-    Component, ComponentKey, ComponentValue, Entity, Fetch, RelationExt, World,
+    ArchetypeId, Component, ComponentKey, ComponentValue, Entity, Fetch, RelationExt, World,
 };
 use crate::{RelationIter, RelationIterMut};
 
@@ -266,7 +266,7 @@ impl<'a> EntityRefMut<'a> {
         let loc = self.loc();
         EntityRef {
             arch: self.world.archetypes.get(loc.arch_id),
-            slot: loc.slot,
+            loc,
             id: self.id,
             world: self.world,
         }
@@ -278,7 +278,7 @@ impl<'a> EntityRefMut<'a> {
         let loc = self.loc();
         EntityRef {
             arch: self.world.archetypes.get(loc.arch_id),
-            slot: loc.slot,
+            loc,
             id: self.id,
             world: self.world,
         }
@@ -303,7 +303,7 @@ impl<'a> EntityRefMut<'a> {
 pub struct EntityRef<'a> {
     pub(crate) world: &'a World,
     pub(crate) arch: &'a Archetype,
-    pub(crate) slot: Slot,
+    pub(crate) loc: EntityLocation,
     pub(crate) id: Entity,
 }
 
@@ -314,7 +314,7 @@ impl<'a> EntityRef<'a> {
         component: Component<T>,
     ) -> Result<AtomicRef<'a, T>, MissingComponent> {
         self.arch
-            .get(self.slot, component)
+            .get(self.loc.slot, component)
             .ok_or_else(|| MissingComponent {
                 id: self.id,
                 desc: component.desc(),
@@ -327,7 +327,7 @@ impl<'a> EntityRef<'a> {
         component: Component<T>,
     ) -> Result<RefMut<'a, T>, MissingComponent> {
         self.arch
-            .get_mut(self.slot, component, self.world.advance_change_tick())
+            .get_mut(self.loc.slot, component, self.world.advance_change_tick())
             .ok_or_else(|| MissingComponent {
                 id: self.id,
                 desc: component.desc(),
@@ -357,7 +357,7 @@ impl<'a> EntityRef<'a> {
         let change_tick = self.world.advance_change_tick();
 
         self.arch
-            .update(self.slot, component, FnWriter::new(f), change_tick)
+            .update(self.loc.slot, component, FnWriter::new(f), change_tick)
     }
 
     /// Updates a component in place
@@ -369,15 +369,19 @@ impl<'a> EntityRef<'a> {
         let tick = self.world.advance_change_tick();
 
         self.arch
-            .update(self.slot, component, WriteDedup::new(value), tick)
+            .update(self.loc.slot, component, WriteDedup::new(value), tick)
+    }
+
+    pub fn query_one<'q, Q: Fetch<'q>>(&'q self, query: &'q Q) -> QueryOne<'q, Q> {
+        QueryOne::new(query, self.world, self.arch, self.loc)
     }
 
     /// Attempt concurrently access a component mutably using and fail if the component is already borrowed
     pub fn try_get<T: ComponentValue>(
         &self,
         component: Component<T>,
     ) -> core::result::Result<Option<AtomicRef<T>>, BorrowError> {
-        self.arch.try_get(self.slot, component)
+        self.arch.try_get(self.loc.slot, component)
     }
 
     /// Attempt to concurrently access a component mutably using and fail if the component is already borrowed
@@ -386,7 +390,7 @@ impl<'a> EntityRef<'a> {
         component: Component<T>,
     ) -> core::result::Result<Option<RefMut<T>>, BorrowMutError> {
         self.arch
-            .try_get_mut(self.slot, component, self.world.advance_change_tick())
+            .try_get_mut(self.loc.slot, component, self.world.advance_change_tick())
     }
 
     /// Returns all relations to other entities of the specified kind
@@ -395,7 +399,7 @@ impl<'a> EntityRef<'a> {
         &self,
         relation: impl RelationExt<T>,
     ) -> RelationIter<'a, T> {
-        RelationIter::new(relation, self.arch, self.slot)
+        RelationIter::new(relation, self.arch, self.loc.slot)
     }
 
     /// Returns all relations to other entities of the specified kind
@@ -407,7 +411,7 @@ impl<'a> EntityRef<'a> {
         RelationIterMut::new(
             relation,
             self.arch,
-            self.slot,
+            self.loc.slot,
             self.world.advance_change_tick(),
         )
     }
@@ -423,7 +427,7 @@ impl<'a> Debug for EntityRef<'a> {
         EntityFormatter {
             world: self.world,
             arch: self.arch,
-            slot: self.slot,
+            slot: self.loc.slot,
             id: self.id,
         }
         .fmt(f)

diff --git a/src/fetch/entity_ref.rs b/src/fetch/entity_ref.rs
@@ -2,8 +2,9 @@ use alloc::vec::Vec;
 
 use crate::{
     archetype::{Archetype, Slot},
+    entity::EntityLocation,
     system::{Access, AccessKind},
-    EntityRef, Fetch, FetchItem, World,
+    ArchetypeId, EntityRef, Fetch, FetchItem, World,
 };
 
 use super::{FetchAccessData, PreparedFetch};
@@ -32,6 +33,7 @@ impl<'w> Fetch<'w> for EntityRefs {
         Some(PreparedEntityRef {
             arch: data.arch,
             world: data.world,
+            arch_id: data.arch_id,
         })
     }
 
@@ -57,12 +59,14 @@ impl<'w> Fetch<'w> for EntityRefs {
 pub struct PreparedEntityRef<'a> {
     world: &'a World,
     arch: &'a Archetype,
+    arch_id: ArchetypeId,
 }
 
 #[doc(hidden)]
 pub struct Batch<'a> {
     pub(crate) world: &'a World,
     pub(crate) arch: &'a Archetype,
+    pub(crate) arch_id: ArchetypeId,
     slot: Slot,
 }
 
@@ -75,6 +79,7 @@ impl<'w, 'q> PreparedFetch<'q> for PreparedEntityRef<'w> {
             world: self.world,
             arch: self.arch,
             slot: slice.start,
+            arch_id: self.arch_id,
         }
     }
 
@@ -86,7 +91,10 @@ impl<'w, 'q> PreparedFetch<'q> for PreparedEntityRef<'w> {
         EntityRef {
             arch: chunk.arch,
             world: chunk.world,
-            slot,
+            loc: EntityLocation {
+                arch_id: chunk.arch_id,
+                slot,
+            },
             id: *chunk.arch.entities.get_unchecked(slot),
         }
     }

diff --git a/src/query/one.rs b/src/query/one.rs
@@ -1,15 +1,16 @@
-use core::mem;
+use core::{mem, ops::Deref};
 
 use crate::{
     archetype::{Archetype, Slice},
     entity::EntityLocation,
     fetch::{FetchPrepareData, PreparedFetch},
-    EntityRef, Fetch, World,
+    EntityRef, Fetch, FetchItem, World,
 };
 
 pub struct QueryOne<'w, Q: Fetch<'w>> {
-    prepared: Q::Prepared,
-    item: <Q::Prepared as PreparedFetch<'static>>::Item,
+    prepared: Option<Q::Prepared>,
+    loc: EntityLocation,
+    // item: <Q::Prepared as PreparedFetch<'static>>::Item,
 }
 
 impl<'w, Q: Fetch<'w>> QueryOne<'w, Q> {
@@ -18,28 +19,31 @@ impl<'w, Q: Fetch<'w>> QueryOne<'w, Q> {
         world: &'w World,
         arch: &'w Archetype,
         loc: EntityLocation,
-    ) -> Option<Self> {
-        let mut prepared = fetch.prepare(FetchPrepareData {
+    ) -> Self {
+        let prepared = fetch.prepare(FetchPrepareData {
             world,
             arch,
             arch_id: loc.arch_id,
             old_tick: 0,
             new_tick: world.advance_change_tick(),
-        })?;
+        });
 
-        let item = {
-            let mut chunk = unsafe { prepared.create_chunk(Slice::single(loc.slot)) };
+        Self { prepared, loc }
+    }
 
-            unsafe { <Q::Prepared as PreparedFetch<'_>>::fetch_next(&mut chunk) }
-        };
+    /// Fetches the query item from the entity, or `None` if the entity does not match the query
+    pub fn get(&mut self) -> Option<<Q as FetchItem<'_>>::Item> {
+        match &mut self.prepared {
+            Some(prepared) => {
+                let item = {
+                    let mut chunk = unsafe { prepared.create_chunk(Slice::single(self.loc.slot)) };
 
-        let item = unsafe {
-            mem::transmute::<
-                <Q::Prepared as PreparedFetch<'_>>::Item,
-                <Q::Prepared as PreparedFetch<'static>>::Item,
-            >(item)
-        };
+                    unsafe { <Q::Prepared as PreparedFetch<'_>>::fetch_next(&mut chunk) }
+                };
 
-        Some(Self { prepared, item })
+                Some(item)
+            }
+            None => None,
+        }
     }
 }
diff --git a/src/world.rs b/src/world.rs
@@ -1066,7 +1066,7 @@ impl World {
         Ok(EntityRef {
             world: self,
             arch,
-            slot: loc.slot,
+            loc,
             id,
         })
     }

diff --git a/tests/entity_access.rs b/tests/entity_access.rs
@@ -1,4 +1,4 @@
-use flax::{component, name, Entity, World};
+use flax::{component, name, Entity, FetchExt, World};
 
 #[test]
 fn entity_access() {
@@ -16,4 +16,21 @@ fn entity_access() {
         .spawn(&mut world);
 
     let entity = world.entity(id).unwrap();
+
+    let query = &(name().cloned(), a());
+    let query2 = &(name().cloned(), a().as_mut());
+    {
+        let mut query = entity.query_one(query);
+        assert_eq!(query.get(), Some(("a".into(), &5)));
+    }
+
+    {
+        let mut query = entity.query_one(query2);
+        *query.get().unwrap().1 += 1;
+
+        assert_eq!(query.get(), Some(("a".into(), &mut 6)));
+    }
+
+    let mut query = entity.query_one(query);
+    assert_eq!(query.get(), Some(("a".into(), &6)));
 }
-Original file line number
+Diff line change
@@ Expand Up / @@ -1066,7 +1066,7 @@ impl World { @@
             Ok(EntityRef {
                 world: self,
                 arch,
-                slot: loc.slot,
+                loc,
                 id,
             })
         }
@@ Expand Down @@