Fix size_hint for partially consumed QueryIter

Instead of returning the total count of elements in the `QueryIter` in
`size_hint`, we return the count of remaining elements in it. This
Fixes bevyengine#5149 (although, if bevyengine#5148 is merged, bevyengine#5149 will re-open)

- bevyengine#5149
- bevyengine#5148

crates/bevy_ecs/src/query/iter.rs

@@ -62,13 +62,7 @@ where
  }
 
  fn size_hint(&self) -> (usize, Option<usize>) {
- let max_size = self
- .query_state
- .matched_archetype_ids
- .iter()
- .map(|id| self.archetypes[*id].len())
- .sum();
-
+ let max_size = self.cursor.remaining(self.tables, self.archetypes);
  let archetype_query = F::Fetch::IS_ARCHETYPAL && QF::IS_ARCHETYPAL;
  let min_size = if archetype_query { max_size } else { 0 };
  (min_size, Some(max_size))
@@ -264,11 +258,16 @@ impl<'w, 's, Q: WorldQuery, F: WorldQuery, const K: usize> QueryCombinationIter<
  return None;
  }
 
- // first, iterate from last to first until next item is found
+ // TODO: can speed up the following code using `cursor.remaining()` instead of `next_item.is_none()`
+ // when Q::Fetch::IS_ARCHETYPAL && F::Fetch::IS_ARCHETYPAL
+ //
+ // let `i` be the index of `c`, the last cursor in `self.cursors` that
+ // returns `K-i` or more elements.
+ // Make cursor in index `j` for all `j` in `(i, K)` a copy of `c` advanced `j-i` times.
+ // If no such `c` exists, return `None`
  'outer: for i in (0..K).rev() {
  match self.cursors[i].next(self.tables, self.archetypes, self.query_state) {
  Some(_) => {
- // walk forward up to last element, propagating cursor state forward
  for j in (i + 1)..K {
  self.cursors[j] = self.cursors[j - 1].clone();
  match self.cursors[j].next(self.tables, self.archetypes, self.query_state) {
@@ -364,11 +363,7 @@ where
  F: WorldQuery + ArchetypeFilter,
 {
  fn len(&self) -> usize {
- self.query_state
- .matched_archetype_ids
- .iter()
- .map(|id| self.archetypes[*id].len())
- .sum()
+ self.size_hint().0
  }
 }
 
@@ -473,6 +468,18 @@ where
  }
  }
 
+ /// How many values will this cursor return?
+ fn remaining(&self, tables: &'w Tables, archetypes: &'w Archetypes) -> usize {
+ let remaining_matched: usize = if Self::IS_DENSE {
+ let ids = self.table_id_iter.clone();
+ ids.map(|id| tables[*id].len()).sum()
+ } else {
+ let ids = self.archetype_id_iter.clone();
+ ids.map(|id| archetypes[*id].len()).sum()
+ };
+ remaining_matched + self.current_len - self.current_index
+ }
+
  // NOTE: If you are changing query iteration code, remember to update the following places, where relevant:
  // QueryIterationCursor, QueryState::for_each_unchecked_manual, QueryState::par_for_each_unchecked_manual
  /// # Safety

crates/bevy_ecs/src/query/mod.rs

@@ -144,18 +144,30 @@ mod tests {
  assert_eq!(values.iter(&world).size_hint().1.unwrap(), n);
  assert_eq!(values.iter(&world).len(), n);
  assert_eq!(values.iter(&world).count(), n);
+ let mut iterator = values.iter(&world);
+ let _ = iterator.next();
+ assert_eq!(iterator.len(), n - 1);
+
  let mut values = world.query_filtered::<&A, Or<(With<B>, Without<C>)>>();
  let n = 7;
  assert_eq!(values.iter(&world).size_hint().0, n);
  assert_eq!(values.iter(&world).size_hint().1.unwrap(), n);
  assert_eq!(values.iter(&world).len(), n);
  assert_eq!(values.iter(&world).count(), n);
+ let mut iterator = values.iter(&world);
+ let _ = iterator.next();
+ assert_eq!(iterator.len(), n - 1);
+
  let mut values = world.query_filtered::<&A, Or<(Without<B>, With<C>)>>();
  let n = 8;
  assert_eq!(values.iter(&world).size_hint().0, n);
  assert_eq!(values.iter(&world).size_hint().1.unwrap(), n);
  assert_eq!(values.iter(&world).len(), n);
  assert_eq!(values.iter(&world).count(), n);
+ let mut iterator = values.iter(&world);
+ let _ = iterator.next();
+ assert_eq!(iterator.len(), n - 1);
+
  let mut values = world.query_filtered::<&A, Or<(Without<B>, Without<C>)>>();
  let n = 9;
  assert_eq!(values.iter(&world).size_hint().0, n);
@@ -169,6 +181,12 @@ mod tests {
  assert_eq!(values.iter(&world).size_hint().1.unwrap(), n);
  assert_eq!(values.iter(&world).len(), n);
  assert_eq!(values.iter(&world).count(), n);
+ let mut iterator = values.iter(&world);
+ let _ = iterator.next();
+ assert_eq!(iterator.len(), 0);
+ let _ = iterator.next();
+ assert_eq!(iterator.len(), 0);
+
  let mut values = world.query_filtered::<&A, Or<(Or<(With<B>, With<C>)>, With<D>)>>();
  let n = 6;
  assert_eq!(values.iter(&world).size_hint().0, n);