Introduce QueryBorrow::iter_combinations

src/query.rs

@@ -6,6 +6,7 @@
 // copied, modified, or distributed except according to those terms.
 
 use core::any::TypeId;
+use core::array;
 use core::marker::PhantomData;
 use core::ptr::NonNull;
 use core::slice::Iter as SliceIter;
@@ -624,6 +625,13 @@ impl<'w, Q: Query> QueryBorrow<'w, Q> {
         unsafe { QueryIter::new(self.world) }
     }
 
+    /// Execute the query for a combination of entities
+    // The lifetime narrowing here is required for soundness.
+    pub fn iter_combinations<const K: usize>(&mut self) -> QueryCombinationIter<'_, Q, K> {
+        self.borrow();
+        unsafe { QueryCombinationIter::new(self.world) }
+    }
+
     /// Provide random access to the query results
     pub fn view(&mut self) -> View<'_, Q> {
         self.borrow();
@@ -833,6 +841,161 @@ impl<'q, Q: Query> ExactSizeIterator for QueryIter<'q, Q> {
     }
 }
 
+/// Struct providing iteration over a combination of the set of entities with the components in `Q`
+pub struct QueryCombinationIter<'q, Q: Query, const K: usize> {
+    world: &'q World,
+    max_archetypes: usize,
+    archetypes: [usize; K],
+    entity_ids: [Entity; K],
+    chunks: [ChunkIter<Q>; K],
+}
+
+impl<'q, Q: Query, const K: usize> QueryCombinationIter<'q, Q, K> {
+    /// # Safety
+    ///
+    /// `'q` must be sufficient to guarantee that `Q` cannot violate borrow safety, either with
+    /// dynamic borrow checks or by representing exclusive access to the `World`.
+    unsafe fn new(world: &'q World) -> Self {
+        let max_archetypes = world.archetypes().len();
+
+        let mut iter = Self {
+            world,
+            max_archetypes,
+            archetypes: [0; K],
+            entity_ids: [Entity::DANGLING; K],
+            chunks: array::from_fn(|_| ChunkIter::empty()),
+        };
+
+        for i in 0..K {
+            if i > 0 {
+                let previous_chunk = &iter.chunks[i-1];
+                iter.chunks[i] = ChunkIter {
+                    entities: previous_chunk.entities,
+                    fetch: previous_chunk.fetch.clone(),
+                    position: previous_chunk.position,
+                    len: previous_chunk.len,
+                };
+                iter.archetypes[i] = iter.archetypes[i-1];
+            }
+            if i < K-1 {
+                loop {
+                    match unsafe { iter.chunks[i].next_entity() } {
+                        None => match iter.next_archetype(i) {
+                            true => continue, // Keep getting entities for i
+                            false => {
+                                // Not enough elements, return dummy
+                                return Self {
+                                    world,
+                                    max_archetypes,
+                                    archetypes: [max_archetypes; K],
+                                    entity_ids: [Entity::DANGLING; K],
+                                    chunks: array::from_fn(|_| ChunkIter::empty()),
+                                };
+                            }
+                        }
+                        Some(id) => {
+                            iter.entity_ids[i] = Entity {
+                                id,
+                                generation: unsafe {
+                                    iter.world
+                                        .entities_meta()
+                                        .get_unchecked(id as usize)
+                                        .generation
+                                },
+                            };
+                            break;
+                        }
+                    }
+                }
+            }
+        }
+
+        iter
+    }
+
+    /// Advance query to the next archetype
+    fn next_archetype(&mut self, index: usize) -> bool {
+        let archetype = self.archetypes[index];
+        if archetype >= self.max_archetypes {
+            return false;
+        }
+
+        self.archetypes[index] += 1;
+
+        let archetype = unsafe { self.world.archetypes_inner().get_unchecked(archetype) };
+        let state = Q::Fetch::prepare(archetype);
+        let fetch = state.map(|state| Q::Fetch::execute(archetype, state));
+        self.chunks[index] = fetch.map_or(ChunkIter::empty(), |fetch| ChunkIter::new(archetype, fetch));
+        true
+    }
+
+    /// Get next combination of queried components
+    // This method narrows the lifetime of the returned item to self
+    pub fn next(&mut self) -> Option<[(Entity, Q::Item<'_>); K]> {
+        'outer: for i in (0..K).rev() {
+            'entity_i: loop {
+                match unsafe { self.chunks[i].next_entity() } {
+                    None => match self.next_archetype(i) {
+                        true => continue 'entity_i, // Keep getting entities for i
+                        false => continue 'outer, // No more entities for i, move on to parent
+                    }
+                    Some(id) => {
+                        self.entity_ids[i] = Entity {
+                            id,
+                            generation: unsafe {
+                                self.world
+                                    .entities_meta()
+                                    .get_unchecked(id as usize)
+                                    .generation
+                            },
+                        };
+                        // Update children to next entities after i
+                        for j in (i + 1)..K {
+                            let previous_chunk = &self.chunks[j-1];
+                            self.chunks[j] = ChunkIter {
+                                entities: previous_chunk.entities,
+                                fetch: previous_chunk.fetch.clone(),
+                                position: previous_chunk.position,
+                                len: previous_chunk.len,
+                            };
+                            self.archetypes[j] = self.archetypes[j-1];
+                            loop {
+                                match unsafe { self.chunks[j].next_entity() } {
+                                    None => match self.next_archetype(j) {
+                                        true => continue, // Keep getting entities for j
+                                        false => continue 'entity_i // No more entities for j, keep getting entities for i
+                                    }
+                                    Some(id) => {
+                                        self.entity_ids[j] = Entity {
+                                            id,
+                                            generation: unsafe {
+                                                self.world
+                                                    .entities_meta()
+                                                    .get_unchecked(id as usize)
+                                                    .generation
+                                            },
+                                        };
+                                        break;
+                                    }
+                                }
+                            }
+                        }
+                        // Found combination
+                        return Some(array::from_fn(|i| {
+                            (
+                                self.entity_ids[i],
+                                unsafe { self.chunks[i].current_item() }
+                            )
+                        }));
+                    },
+                }
+            }
+        }
+
+        None
+    }
+}
+
 /// A query builder that's convertible directly into an iterator
 pub struct QueryMut<'q, Q: Query> {
     iter: QueryIter<'q, Q>,
@@ -958,6 +1121,21 @@ impl<Q: Query> ChunkIter<Q> {
         Some((*entity, item))
     }
 
+    #[inline]
+    unsafe fn next_entity(&mut self) -> Option<u32> {
+        if self.position == self.len {
+            return None;
+        }
+        let entity = self.entities.as_ptr().add(self.position);
+        self.position += 1;
+        Some(*entity)
+    }
+
+    #[inline]
+    unsafe fn current_item<'a>(&mut self) -> Q::Item<'a> {
+        Q::get(&self.fetch, self.position - 1)
+    }
+
     fn remaining(&self) -> usize {
         self.len - self.position
     }

tests/tests.rs

@@ -985,3 +985,69 @@ fn query_many_duplicate() {
     let e = world.spawn(());
     _ = world.query_many_mut::<(), 2>([e, e]);
 }
+
+#[test]
+fn query_combination() {
+    let mut world = World::new();
+    let a = world.spawn((42, true));
+    let b = world.spawn((17,));
+    let c = world.spawn((21,));
+    let mut query = world.query::<&mut i32>();
+    let mut iter = query.iter_combinations::<2>();
+    assert_eq!(iter.next(), Some([(a, &mut 42), (b, &mut 17)]));
+    assert_eq!(iter.next(), Some([(a, &mut 42), (c, &mut 21)]));
+    assert_eq!(iter.next(), Some([(b, &mut 17), (c, &mut 21)]));
+    assert_eq!(iter.next(), None);
+}
+
+#[test]
+fn query_combination_exact() {
+    let mut world = World::new();
+    let a = world.spawn((42, true));
+    let b = world.spawn((17,));
+    let c = world.spawn((21,));
+    let mut query = world.query::<&mut i32>();
+    let mut iter = query.iter_combinations::<3>();
+    assert_eq!(iter.next(), Some([(a, &mut 42), (b, &mut 17), (c, &mut 21)]));
+    assert_eq!(iter.next(), None);
+}
+
+#[test]
+fn query_combination_zero() {
+    let mut world = World::new();
+    let _ = world.spawn((42, true));
+    let _ = world.spawn((17,));
+    let _ = world.spawn((21,));
+    let mut query = world.query::<&mut i32>();
+    let mut iter = query.iter_combinations::<0>();
+    assert_eq!(iter.next(), None);
+}
+
+#[test]
+fn query_combination_not_enough_entities() {
+    let mut world = World::new();
+    let _ = world.spawn((42, true));
+    let _ = world.spawn((17,));
+    let _ = world.spawn((21,));
+    let mut query = world.query::<&mut i32>();
+
+    let mut iter = query.iter_combinations::<4>();
+    assert_eq!(iter.next(), None);
+
+    let mut iter = query.iter_combinations::<5>();
+    assert_eq!(iter.next(), None);
+
+    let mut iter = query.iter_combinations::<100>();
+    assert_eq!(iter.next(), None);
+}
+
+#[test]
+fn query_combination_no_match() {
+    let mut world = World::new();
+    let _ = world.spawn((42, true));
+    let _ = world.spawn((17,));
+    let _ = world.spawn((21,));
+    let mut query = world.query::<&String>();
+    let mut iter = query.iter_combinations::<2>();
+    assert_eq!(iter.next(), None);
+}