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 #5149 (although, if #5148 is merged, #5149 will re-open)

- #5149
- #5148

Author: nicopap

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+1` 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);