flax/query/
iter.rs

1use crate::{
2    archetype::{Archetype, Slice, Slot},
3    fetch::PreparedFetch,
4    filter::{next_slice, Filtered},
5    Entity,
6};
7
8/// Iterates over a chunk of entities, specified by a predicate.
9/// In essence, this is the unflattened version of [crate::QueryIter].
10pub struct Chunk<'q, Q: PreparedFetch<'q>> {
11    arch: &'q Archetype,
12    fetch: Q::Chunk,
13    pos: Slot,
14    end: Slot,
15}
16
17impl<'q, Q: PreparedFetch<'q>> core::fmt::Debug for Chunk<'q, Q> {
18    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
19        f.debug_struct("Batch")
20            .field("pos", &self.pos)
21            .field("end", &self.end)
22            .finish()
23    }
24}
25
26impl<'q, Q: PreparedFetch<'q>> Chunk<'q, Q> {
27    pub(crate) fn new(arch: &'q Archetype, chunk: Q::Chunk, slice: Slice) -> Self {
28        Self {
29            arch,
30            fetch: chunk,
31            pos: slice.start,
32            end: slice.end,
33        }
34    }
35
36    pub(crate) fn slots(&self) -> Slice {
37        Slice::new(self.pos, self.end)
38    }
39
40    /// Returns the archetype for this batch.
41    /// **Note**: The borrow of the fetch is still held and may result in borrow
42    /// errors.
43    pub fn arch(&self) -> &Archetype {
44        self.arch
45    }
46
47    /// Returns the number of items which would be yielded by this batch
48    pub fn len(&self) -> usize {
49        self.slots().len()
50    }
51
52    /// Returns true if the batch is empty
53    pub fn is_empty(&self) -> bool {
54        self.slots().is_empty()
55    }
56}
57
58impl<'q, Q> Iterator for Chunk<'q, Q>
59where
60    Q: PreparedFetch<'q>,
61{
62    type Item = Q::Item;
63
64    fn next(&mut self) -> Option<Q::Item> {
65        if self.pos == self.end {
66            None
67        } else {
68            // let fetch = unsafe { &mut *(self.fetch as *mut Q::Batch) };
69            let item = unsafe { Q::fetch_next(&mut self.fetch) };
70            self.pos += 1;
71            Some(item)
72        }
73    }
74}
75
76impl<'q, Q> Chunk<'q, Q>
77where
78    Q: PreparedFetch<'q>,
79{
80    pub(crate) fn next_with_id(&mut self) -> Option<(Entity, Q::Item)> {
81        if self.pos == self.end {
82            None
83        } else {
84            let item = unsafe { Q::fetch_next(&mut self.fetch) };
85            let id = self.arch.entities[self.pos];
86            self.pos += 1;
87            Some((id, item))
88        }
89    }
90
91    pub(crate) fn next_full(&mut self) -> Option<(Slot, Entity, Q::Item)> {
92        if self.pos == self.end {
93            None
94        } else {
95            let slot = self.pos;
96            let item = unsafe { Q::fetch_next(&mut self.fetch) };
97            let id = self.arch.entities[slot];
98            self.pos += 1;
99
100            Some((slot, id, item))
101        }
102    }
103}
104
105/// An iterator over a single archetype which returns chunks.
106/// The chunk size is determined by the largest continuous matched entities for
107/// filters.
108pub struct ArchetypeChunks<'q, Q, F> {
109    pub(crate) arch: &'q Archetype,
110    pub(crate) fetch: *mut Filtered<Q, F>,
111    pub(crate) slots: Slice,
112}
113
114unsafe impl<'q, Q: 'q, F: 'q> Sync for ArchetypeChunks<'q, Q, F> where &'q mut Filtered<Q, F>: Sync {}
115unsafe impl<'q, Q: 'q, F: 'q> Send for ArchetypeChunks<'q, Q, F> where &'q mut Filtered<Q, F>: Send {}
116
117impl<'q, Q, F> Iterator for ArchetypeChunks<'q, Q, F>
118where
119    Q: 'q + PreparedFetch<'q>,
120    F: 'q + PreparedFetch<'q>,
121{
122    type Item = Chunk<'q, Q>;
123
124    #[inline]
125    fn next(&mut self) -> Option<Self::Item> {
126        // Fetch will never change and all calls are disjoint
127        let fetch = unsafe { &mut *self.fetch };
128
129        // Get the next chunk
130        let slots = next_slice(&mut self.slots, fetch)?;
131
132        // Safety: Disjoint chunk
133        let chunk = unsafe { fetch.create_chunk(slots) };
134        let chunk = Chunk::new(self.arch, chunk, slots);
135
136        Some(chunk)
137    }
138}
flax/query/iter.rs

flax/query/
iter.rs
