feat: remove fetch slot indexing and allow acces to world in filters

ten3roberts · Sep 10, 2023 · 2e3e582 · 2e3e582
1 parent 2278763
commit 2e3e582
Show file tree

Hide file tree

Showing 29 changed files with 390 additions and 199 deletions.
diff --git a/flax-derive/src/lib.rs b/flax-derive/src/lib.rs
@@ -166,8 +166,8 @@ fn derive_fetch_struct(params: &Params) -> TokenStream {
             }
 
             #[inline]
-            fn filter_arch(&self, arch: &#crate_name::archetype::Archetype) -> bool {
-                #(#crate_name::Fetch::filter_arch(&self.#field_names, arch))&&*
+            fn filter_arch(&self, data: #crate_name::fetch::FetchAccessData) -> bool {
+                #(#crate_name::Fetch::filter_arch(&self.#field_names, data))&&*
             }
 
             fn describe(&self, f: &mut ::core::fmt::Formatter<'_>) -> ::core::fmt::Result {
@@ -387,9 +387,9 @@ fn derive_prepared_struct(params: &Params) -> TokenStream {
             type Chunk = (#(<<#field_types as #crate_name::fetch::Fetch<'w>>::Prepared as #crate_name::fetch::PreparedFetch<'q>>::Chunk,)*);
 
             #[inline]
-            unsafe fn fetch_next(chunk: &mut Self::Chunk, slot: #crate_name::archetype::Slot) -> Self::Item {
+            unsafe fn fetch_next(chunk: &mut Self::Chunk) -> Self::Item {
                 Self::Item {
-                    #(#field_names: <<#field_types as #crate_name::fetch::Fetch<'w>>::Prepared as #crate_name::fetch::PreparedFetch<'q>>::fetch_next(&mut chunk.#field_idx, slot),)*
+                    #(#field_names: <<#field_types as #crate_name::fetch::Fetch<'w>>::Prepared as #crate_name::fetch::PreparedFetch<'q>>::fetch_next(&mut chunk.#field_idx),)*
                 }
             }
 

diff --git a/src/component.rs b/src/component.rs
@@ -36,7 +36,7 @@ impl<T> ComponentValue for T where T: Send + Sync + 'static {}
 #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
 pub struct ComponentKey {
     pub(crate) id: Entity,
-    /// The object entity if the component is a pair
+    /// The object entity if the component is a relation
     pub(crate) object: Option<Entity>,
 }
 

diff --git a/src/fetch/as_deref.rs b/src/fetch/as_deref.rs
@@ -1,5 +1,5 @@
 use super::{FetchAccessData, FmtQuery, PreparedFetch, ReadOnlyFetch};
-use crate::{archetype::Slot, query::ArchetypeSearcher, system::Access, Fetch, FetchItem};
+use crate::{query::ArchetypeSearcher, system::Access, Fetch, FetchItem};
 use alloc::vec::Vec;
 use core::{fmt, ops::Deref};
 
@@ -30,8 +30,8 @@ where
     }
 
     #[inline]
-    fn filter_arch(&self, arch: &crate::archetype::Archetype) -> bool {
-        self.0.filter_arch(arch)
+    fn filter_arch(&self, data: FetchAccessData) -> bool {
+        self.0.filter_arch(data)
     }
 
     #[inline]
@@ -68,8 +68,8 @@ where
     }
 
     #[inline]
-    unsafe fn fetch_next(chunk: &mut Self::Chunk, slot: Slot) -> Self::Item {
-        F::fetch_next(chunk, slot)
+    unsafe fn fetch_next(chunk: &mut Self::Chunk) -> Self::Item {
+        F::fetch_next(chunk)
     }
 }
 

diff --git a/src/fetch/cloned.rs b/src/fetch/cloned.rs
@@ -6,7 +6,7 @@ use core::{
 use alloc::vec::Vec;
 
 use crate::{
-    archetype::{Archetype, Slice, Slot},
+    archetype::{Slice, Slot},
     system::Access,
     Fetch, FetchItem,
 };
@@ -45,8 +45,8 @@ where
         Some(Cloned(self.0.prepare(data)?))
     }
 
-    fn filter_arch(&self, arch: &Archetype) -> bool {
-        self.0.filter_arch(arch)
+    fn filter_arch(&self, data: FetchAccessData) -> bool {
+        self.0.filter_arch(data)
     }
 
     fn access(&self, data: FetchAccessData, dst: &mut Vec<Access>) {
@@ -77,8 +77,8 @@ where
         self.0.create_chunk(slots)
     }
 
-    unsafe fn fetch_next(chunk: &mut Self::Chunk, slot: Slot) -> Self::Item {
-        F::fetch_next(chunk, slot).clone()
+    unsafe fn fetch_next(chunk: &mut Self::Chunk) -> Self::Item {
+        F::fetch_next(chunk).clone()
     }
 
     unsafe fn filter_slots(&mut self, slots: Slice) -> Slice {

diff --git a/src/fetch/component.rs b/src/fetch/component.rs
@@ -21,7 +21,7 @@ impl<'w, 'q, T: 'q> PreparedFetch<'q> for ReadComponent<'w, T> {
 
     #[inline]
     // See: <https://godbolt.org/z/8fWa136b9>
-    unsafe fn fetch_next(chunk: &mut Self::Chunk, _: Slot) -> Self::Item {
+    unsafe fn fetch_next(chunk: &mut Self::Chunk) -> Self::Item {
         let old = chunk.as_ptr();
         chunk.advance(1);
         &*old
@@ -57,8 +57,8 @@ where
     }
 
     #[inline]
-    fn filter_arch(&self, arch: &Archetype) -> bool {
-        arch.has(self.key())
+    fn filter_arch(&self, data: FetchAccessData) -> bool {
+        data.arch.has(self.key())
     }
 
     fn access(&self, data: FetchAccessData, dst: &mut Vec<Access>) {

diff --git a/src/fetch/component_mut.rs b/src/fetch/component_mut.rs
@@ -3,7 +3,7 @@ use alloc::vec::Vec;
 use core::fmt::{self, Formatter};
 
 use crate::{
-    archetype::{Archetype, CellMutGuard, Slice, Slot},
+    archetype::{Archetype, CellMutGuard, Slice},
     system::{Access, AccessKind},
     util::PtrMut,
     Component, ComponentValue, Fetch, FetchItem,
@@ -36,8 +36,8 @@ where
     }
 
     #[inline]
-    fn filter_arch(&self, arch: &Archetype) -> bool {
-        arch.has(self.0.key())
+    fn filter_arch(&self, data: FetchAccessData) -> bool {
+        data.arch.has(self.0.key())
     }
 
     #[inline]
@@ -88,7 +88,7 @@ impl<'w, 'q, T: 'q + ComponentValue> PreparedFetch<'q> for WriteComponent<'w, T>
 
     #[inline]
     // See: <https://godbolt.org/z/8fWa136b9>
-    unsafe fn fetch_next(chunk: &mut Self::Chunk, _: Slot) -> Self::Item {
+    unsafe fn fetch_next(chunk: &mut Self::Chunk) -> Self::Item {
         let old = chunk.as_ptr();
         chunk.advance(1);
         &mut *old

diff --git a/src/fetch/copied.rs b/src/fetch/copied.rs
@@ -5,11 +5,7 @@ use core::{
 
 use alloc::vec::Vec;
 
-use crate::{
-    archetype::{Archetype, Slice, Slot},
-    system::Access,
-    Fetch, FetchItem,
-};
+use crate::{archetype::Slice, system::Access, Fetch, FetchItem};
 
 use super::{FetchAccessData, FetchPrepareData, PreparedFetch, ReadOnlyFetch, TransformFetch};
 
@@ -45,8 +41,8 @@ where
         Some(Copied(self.0.prepare(data)?))
     }
 
-    fn filter_arch(&self, arch: &Archetype) -> bool {
-        self.0.filter_arch(arch)
+    fn filter_arch(&self, data: FetchAccessData) -> bool {
+        self.0.filter_arch(data)
     }
 
     #[inline]
@@ -78,8 +74,8 @@ where
         self.0.create_chunk(slots)
     }
 
-    unsafe fn fetch_next(chunk: &mut Self::Chunk, slot: Slot) -> Self::Item {
-        *F::fetch_next(chunk, slot)
+    unsafe fn fetch_next(chunk: &mut Self::Chunk) -> Self::Item {
+        *F::fetch_next(chunk)
     }
 
     unsafe fn filter_slots(&mut self, slots: Slice) -> Slice {

diff --git a/src/fetch/entity_ref.rs b/src/fetch/entity_ref.rs
@@ -35,7 +35,7 @@ impl<'w> Fetch<'w> for EntityRefs {
         })
     }
 
-    fn filter_arch(&self, _: &crate::archetype::Archetype) -> bool {
+    fn filter_arch(&self, _: FetchAccessData) -> bool {
         true
     }
 
@@ -63,21 +63,26 @@ pub struct PreparedEntityRef<'a> {
 pub struct Batch<'a> {
     pub(crate) world: &'a World,
     pub(crate) arch: &'a Archetype,
+    slot: Slot,
 }
 
 impl<'w, 'q> PreparedFetch<'q> for PreparedEntityRef<'w> {
     type Item = EntityRef<'q>;
     type Chunk = Batch<'q>;
 
-    unsafe fn create_chunk(&'q mut self, _: crate::archetype::Slice) -> Self::Chunk {
+    unsafe fn create_chunk(&'q mut self, slice: crate::archetype::Slice) -> Self::Chunk {
         Batch {
             world: self.world,
             arch: self.arch,
+            slot: slice.start,
         }
     }
 
     #[inline]
-    unsafe fn fetch_next(chunk: &mut Self::Chunk, slot: Slot) -> Self::Item {
+    unsafe fn fetch_next(chunk: &mut Self::Chunk) -> Self::Item {
+        let slot = chunk.slot;
+        chunk.slot += 1;
+
         EntityRef {
             arch: chunk.arch,
             world: chunk.world,
@@ -91,7 +96,29 @@ impl<'w, 'q> PreparedFetch<'q> for PreparedEntityRef<'w> {
 mod test {
     use itertools::Itertools;
 
-    use crate::{component, name, Entity, EntityIds, FetchExt, Query, World};
+    use crate::{component, name, BatchSpawn, Entity, EntityIds, FetchExt, Query, World};
+
+    #[test]
+    fn entity_refs_chunks() {
+        component! {
+            a: i32,
+        }
+
+        let mut batch = BatchSpawn::new(32);
+        batch.set(a(), (0..).map(|v| (v % 8) - 4)).unwrap();
+
+        let mut world = World::new();
+        batch.spawn(&mut world);
+
+        let mut query = Query::new(super::EntityRefs).filter(a().ge(0));
+        let res = query
+            .borrow(&world)
+            .iter()
+            .map(|v| v.get_copy(a()).unwrap())
+            .collect_vec();
+
+        assert_eq!(res, &[0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3]);
+    }
 
     #[test]
     fn entity_refs() {

diff --git a/src/fetch/map.rs b/src/fetch/map.rs
@@ -1,8 +1,8 @@
 use alloc::vec::Vec;
 
-use crate::{archetype::Slot, Fetch, FetchItem};
+use crate::{Fetch, FetchItem};
 
-use super::{FmtQuery, PreparedFetch};
+use super::{FetchAccessData, FmtQuery, PreparedFetch};
 
 /// Maps the result of a query to another type on the query level.
 ///
@@ -41,8 +41,8 @@ where
         })
     }
 
-    fn filter_arch(&self, arch: &crate::archetype::Archetype) -> bool {
-        self.query.filter_arch(arch)
+    fn filter_arch(&self, data: FetchAccessData) -> bool {
+        self.query.filter_arch(data)
     }
 
     fn access(&self, data: super::FetchAccessData, dst: &mut Vec<crate::system::Access>) {
@@ -69,8 +69,8 @@ where
         (self.func, self.query.create_chunk(slots))
     }
 
-    unsafe fn fetch_next(chunk: &mut Self::Chunk, slot: Slot) -> Self::Item {
-        (chunk.0)(Q::fetch_next(&mut chunk.1, slot))
+    unsafe fn fetch_next(chunk: &mut Self::Chunk) -> Self::Item {
+        (chunk.0)(Q::fetch_next(&mut chunk.1))
     }
 
     unsafe fn filter_slots(&mut self, slots: crate::archetype::Slice) -> crate::archetype::Slice {

diff --git a/src/fetch/maybe_mut.rs b/src/fetch/maybe_mut.rs
@@ -38,8 +38,8 @@ impl<'w, T: ComponentValue> Fetch<'w> for MaybeMut<T> {
         })
     }
 
-    fn filter_arch(&self, arch: &crate::archetype::Archetype) -> bool {
-        arch.has(self.0.key())
+    fn filter_arch(&self, data: FetchAccessData) -> bool {
+        data.arch.has(self.0.key())
     }
 
     fn access(&self, data: FetchAccessData, dst: &mut Vec<Access>) {
@@ -82,21 +82,26 @@ pub struct Batch<'a> {
     cell: &'a Cell,
     new_tick: u32,
     ids: &'a [Entity],
+    slot: Slot,
 }
 
 impl<'w, 'q, T: ComponentValue> PreparedFetch<'q> for PreparedMaybeMut<'w, T> {
     type Item = MutGuard<'q, T>;
     type Chunk = Batch<'q>;
 
-    unsafe fn create_chunk(&'q mut self, _: crate::archetype::Slice) -> Self::Chunk {
+    unsafe fn create_chunk(&'q mut self, slice: crate::archetype::Slice) -> Self::Chunk {
         Batch {
             cell: self.cell,
             new_tick: self.new_tick,
             ids: self.entities,
+            slot: slice.start,
         }
     }
 
-    unsafe fn fetch_next(chunk: &mut Self::Chunk, slot: Slot) -> Self::Item {
+    unsafe fn fetch_next(chunk: &mut Self::Chunk) -> Self::Item {
+        let slot = chunk.slot;
+        chunk.slot += 1;
+
         MutGuard {
             slot,
             cell: chunk.cell,