feat: make entity ids a filter

src/filter/constant.rs

@@ -1,8 +1,8 @@
 use crate::{
-    archetype::Slice,
+    archetype::{Slice, Slot},
     fetch::{FetchAccessData, FetchPrepareData, PreparedFetch},
     system::Access,
-    Fetch, FetchItem,
+    Entity, Fetch, FetchItem,
 };
 use alloc::vec::Vec;
 use core::fmt::{self, Formatter};
@@ -106,6 +106,72 @@ impl<'q> PreparedFetch<'q> for All {
     unsafe fn fetch_next(_: &mut Self::Chunk) -> Self::Item {}
 }
 
+impl<'w> FetchItem<'w> for Entity {
+    type Item = Entity;
+}
+
+impl<'w> Fetch<'w> for Entity {
+    const MUTABLE: bool = false;
+    type Prepared = PreparedEntity;
+
+    fn prepare(&'w self, data: FetchPrepareData<'w>) -> Option<Self::Prepared> {
+        let loc = data.world.location(*self).ok()?;
+
+        if data.arch_id == loc.arch_id {
+            Some(PreparedEntity {
+                slot: loc.slot,
+                id: *self,
+            })
+        } else {
+            None
+        }
+    }
+
+    fn describe(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "entity {:?}", self)
+    }
+
+    fn filter_arch(&self, data: FetchAccessData) -> bool {
+        if let Ok(loc) = data.world.location(*self) {
+            data.arch_id == loc.arch_id
+        } else {
+            false
+        }
+    }
+
+    #[inline]
+    fn access(&self, _: FetchAccessData, _: &mut Vec<Access>) {}
+}
+
+#[doc(hidden)]
+pub struct PreparedEntity {
+    slot: Slot,
+    id: Entity,
+}
+
+impl<'w> PreparedFetch<'w> for PreparedEntity {
+    type Item = Entity;
+
+    type Chunk = Entity;
+
+    unsafe fn create_chunk(&'w mut self, slots: Slice) -> Self::Chunk {
+        assert!(slots.start == self.slot && slots.end == self.slot);
+        self.id
+    }
+
+    unsafe fn fetch_next(chunk: &mut Self::Chunk) -> Self::Item {
+        *chunk
+    }
+
+    unsafe fn filter_slots(&mut self, slots: Slice) -> Slice {
+        if slots.contains(self.slot) {
+            Slice::single(self.slot)
+        } else {
+            Slice::new(slots.end, slots.end)
+        }
+    }
+}
+
 impl<'q> FetchItem<'q> for Slice {
     type Item = ();
 }

src/filter/mod.rs

@@ -21,8 +21,8 @@ use crate::{
 
 pub use change::{ChangeFilter, RemovedFilter};
 pub use cmp::{Cmp, Equal, Greater, GreaterEq, Less, LessEq};
-pub use constant::*;
-pub use set::*;
+pub use constant::{All, Nothing};
+pub use set::{And, Not, Or, Union};
 
 macro_rules! gen_bitops {
     ($ty:ident[$($p: tt),*]) => {

tests/filter.rs

@@ -1,6 +1,5 @@
-use flax::*;
+use flax::{filter::Or, *};
 use itertools::Itertools;
-use pretty_assertions::assert_eq;
 use std::sync::Arc;
 
 component! {
@@ -329,3 +328,47 @@ fn sparse_and() {
     let batches = schedule.batch_info(&world);
     assert_eq!(batches.len(), 1);
 }
+
+#[test]
+fn entity_filter() {
+    component! {
+        index: usize,
+    }
+
+    let mut world = World::new();
+
+    let ids = (0..10)
+        .map(|i| {
+            Entity::builder()
+                .set(a(), 5.4)
+                .set(index(), i)
+                .set(b(), "Foo".into())
+                .spawn(&mut world)
+        })
+        .collect_vec();
+
+    let ids2 = (10..20)
+        .map(|i| {
+            Entity::builder()
+                .set(a(), 5.4)
+                .set(index(), i)
+                .spawn(&mut world)
+        })
+        .collect_vec();
+
+    let ids3 = (20..40)
+        .map(|i| Entity::builder().set(index(), i).spawn(&mut world))
+        .collect_vec();
+
+    let mut query = Query::new(index().copied()).filter(Or((ids[5], ids2[7])));
+
+    assert_eq!(query.borrow(&world).iter().sorted().collect_vec(), &[5, 17]);
+    let mut query = Query::new((entity_ids(), Or((ids[5], ids2[7])).satisfied()));
+
+    let all_ids = ids.iter().chain(ids2.iter()).chain(ids3.iter());
+    let expected = all_ids
+        .map(|&id| (id, id == ids[5] || id == ids2[7]))
+        .collect_vec();
+
+    assert_eq!(query.borrow(&world).iter().sorted().collect_vec(), expected);
+}