diff --git a/crates/bevy_ecs/src/query/iter.rs b/crates/bevy_ecs/src/query/iter.rs index 8a693eabb02c4..336da8c6c0fd6 100644 --- a/crates/bevy_ecs/src/query/iter.rs +++ b/crates/bevy_ecs/src/query/iter.rs @@ -56,13 +56,7 @@ impl<'w, 's, Q: WorldQuery, F: ReadOnlyWorldQuery> Iterator for QueryIter<'w, 's } fn size_hint(&self) -> (usize, Option) { - let max_size = self - .query_state - .matched_archetype_ids - .iter() - .map(|id| self.archetypes[*id].len()) - .sum(); - + let max_size = self.cursor.max_remaining(self.tables, self.archetypes); let archetype_query = Q::IS_ARCHETYPAL && F::IS_ARCHETYPAL; let min_size = if archetype_query { max_size } else { 0 }; (min_size, Some(max_size)) @@ -348,11 +342,16 @@ impl<'w, 's, Q: WorldQuery, F: ReadOnlyWorldQuery, const K: usize> return None; } - // first, iterate from last to first until next item is found + // PERF: can speed up the following code using `cursor.remaining()` instead of `next_item.is_none()` + // when Q::IS_ARCHETYPAL && F::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_cursor(); match self.cursors[j].next(self.tables, self.archetypes, self.query_state) { @@ -406,36 +405,29 @@ impl<'w, 's, Q: ReadOnlyWorldQuery, F: ReadOnlyWorldQuery, const K: usize> Itera } fn size_hint(&self) -> (usize, Option) { - if K == 0 { - return (0, Some(0)); - } - - let max_size: usize = self - .query_state - .matched_archetype_ids - .iter() - .map(|id| self.archetypes[*id].len()) - .sum(); - - if max_size < K { - return (0, Some(0)); - } - if max_size == K { - return (1, Some(1)); - } - // binomial coefficient: (n ; k) = n! / k!(n-k)! = (n*n-1*...*n-k+1) / k! // See https://en.wikipedia.org/wiki/Binomial_coefficient // See https://blog.plover.com/math/choose.html for implementation // It was chosen to reduce overflow potential. fn choose(n: usize, k: usize) -> Option { + if k > n || n == 0 { + return Some(0); + } + let k = k.min(n - k); let ks = 1..=k; let ns = (n - k + 1..=n).rev(); ks.zip(ns) .try_fold(1_usize, |acc, (k, n)| Some(acc.checked_mul(n)? / k)) } - let smallest = K.min(max_size - K); - let max_combinations = choose(max_size, smallest); + // sum_i=0..k choose(cursors[i].remaining, k-i) + let max_combinations = self + .cursors + .iter() + .enumerate() + .try_fold(0, |acc, (i, cursor)| { + let n = cursor.max_remaining(self.tables, self.archetypes); + Some(acc + choose(n, K - i)?) + }); let archetype_query = F::IS_ARCHETYPAL && Q::IS_ARCHETYPAL; let known_max = max_combinations.unwrap_or(usize::MAX); @@ -449,11 +441,7 @@ where F: ArchetypeFilter, { fn len(&self) -> usize { - self.query_state - .matched_archetype_ids - .iter() - .map(|id| self.archetypes[*id].len()) - .sum() + self.size_hint().0 } } @@ -568,6 +556,21 @@ impl<'w, 's, Q: WorldQuery, F: ReadOnlyWorldQuery> QueryIterationCursor<'w, 's, } } + /// How many values will this cursor return at most? + /// + /// Note that if `Q::IS_ARCHETYPAL && F::IS_ARCHETYPAL`, the return value + /// will be **the exact count of remaining values**. + fn max_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].entity_count()).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: // QueryIter, QueryIterationCursor, QueryManyIter, QueryCombinationIter, QueryState::for_each_unchecked_manual, QueryState::par_for_each_unchecked_manual /// # Safety diff --git a/crates/bevy_ecs/src/query/mod.rs b/crates/bevy_ecs/src/query/mod.rs index 44147b98d55db..b1a610df91a5c 100644 --- a/crates/bevy_ecs/src/query/mod.rs +++ b/crates/bevy_ecs/src/query/mod.rs @@ -57,100 +57,6 @@ mod tests { #[test] fn query_filtered_exactsizeiterator_len() { - fn assert_all_sizes_iterator_equal( - iterator: impl ExactSizeIterator, - expected_size: usize, - query_type: &'static str, - ) { - let len = iterator.len(); - let size_hint_0 = iterator.size_hint().0; - let size_hint_1 = iterator.size_hint().1; - // `count` tests that not only it is the expected value, but also - // the value is accurate to what the query returns. - let count = iterator.count(); - // This will show up when one of the asserts in this function fails - println!( - r#"query declared sizes: - for query: {query_type} - expected: {expected_size} - len: {len} - size_hint().0: {size_hint_0} - size_hint().1: {size_hint_1:?} - count(): {count}"# - ); - assert_eq!(len, expected_size); - assert_eq!(size_hint_0, expected_size); - assert_eq!(size_hint_1, Some(expected_size)); - assert_eq!(count, expected_size); - } - fn assert_all_sizes_equal(world: &mut World, expected_size: usize) - where - Q: ReadOnlyWorldQuery, - F: ReadOnlyWorldQuery, - F::ReadOnly: ArchetypeFilter, - { - let mut query = world.query_filtered::(); - let iter = query.iter(world); - let query_type = type_name::>(); - assert_all_sizes_iterator_equal(iter, expected_size, query_type); - } - - let mut world = World::new(); - world.spawn((A(1), B(1))); - world.spawn(A(2)); - world.spawn(A(3)); - - assert_all_sizes_equal::<&A, With>(&mut world, 1); - assert_all_sizes_equal::<&A, Without>(&mut world, 2); - - let mut world = World::new(); - world.spawn((A(1), B(1), C(1))); - world.spawn((A(2), B(2))); - world.spawn((A(3), B(3))); - world.spawn((A(4), C(4))); - world.spawn((A(5), C(5))); - world.spawn((A(6), C(6))); - world.spawn(A(7)); - world.spawn(A(8)); - world.spawn(A(9)); - world.spawn(A(10)); - - // With/Without for B and C - assert_all_sizes_equal::<&A, With>(&mut world, 3); - assert_all_sizes_equal::<&A, With>(&mut world, 4); - assert_all_sizes_equal::<&A, Without>(&mut world, 7); - assert_all_sizes_equal::<&A, Without>(&mut world, 6); - - // With/Without (And) combinations - assert_all_sizes_equal::<&A, (With, With)>(&mut world, 1); - assert_all_sizes_equal::<&A, (With, Without)>(&mut world, 2); - assert_all_sizes_equal::<&A, (Without, With)>(&mut world, 3); - assert_all_sizes_equal::<&A, (Without, Without)>(&mut world, 4); - - // With/Without Or<()> combinations - assert_all_sizes_equal::<&A, Or<(With, With)>>(&mut world, 6); - assert_all_sizes_equal::<&A, Or<(With, Without)>>(&mut world, 7); - assert_all_sizes_equal::<&A, Or<(Without, With)>>(&mut world, 8); - assert_all_sizes_equal::<&A, Or<(Without, Without)>>(&mut world, 9); - assert_all_sizes_equal::<&A, (Or<(With,)>, Or<(With,)>)>(&mut world, 1); - assert_all_sizes_equal::<&A, Or<(Or<(With, With)>, With)>>(&mut world, 6); - - for i in 11..14 { - world.spawn((A(i), D(i))); - } - - assert_all_sizes_equal::<&A, Or<(Or<(With, With)>, With)>>(&mut world, 9); - assert_all_sizes_equal::<&A, Or<(Or<(With, With)>, Without)>>(&mut world, 10); - - // a fair amount of entities - for i in 14..20 { - world.spawn((C(i), D(i))); - } - assert_all_sizes_equal::, With)>(&mut world, 6); - } - - #[test] - fn query_filtered_combination_size() { fn choose(n: usize, k: usize) -> usize { if n == 0 || k == 0 || n < k { return 0; @@ -161,25 +67,30 @@ mod tests { } fn assert_combination(world: &mut World, expected_size: usize) where - Q: WorldQuery, + Q: ReadOnlyWorldQuery, F: ReadOnlyWorldQuery, F::ReadOnly: ArchetypeFilter, { let mut query = world.query_filtered::(); - let iter = query.iter_combinations::(world); let query_type = type_name::>(); - assert_all_sizes_iterator_equal(iter, expected_size, query_type); + let iter = query.iter_combinations::(world); + assert_all_sizes_iterator_equal(iter, expected_size, 0, query_type); + let iter = query.iter_combinations::(world); + assert_all_sizes_iterator_equal(iter, expected_size, 1, query_type); + let iter = query.iter_combinations::(world); + assert_all_sizes_iterator_equal(iter, expected_size, 5, query_type); } fn assert_all_sizes_equal(world: &mut World, expected_size: usize) where - Q: WorldQuery, + Q: ReadOnlyWorldQuery, F: ReadOnlyWorldQuery, F::ReadOnly: ArchetypeFilter, { let mut query = world.query_filtered::(); - let iter = query.iter(world); let query_type = type_name::>(); - assert_all_sizes_iterator_equal(iter, expected_size, query_type); + assert_all_exact_sizes_iterator_equal(query.iter(world), expected_size, 0, query_type); + assert_all_exact_sizes_iterator_equal(query.iter(world), expected_size, 1, query_type); + assert_all_exact_sizes_iterator_equal(query.iter(world), expected_size, 5, query_type); let expected = expected_size; assert_combination::(world, choose(expected, 0)); @@ -187,13 +98,29 @@ mod tests { assert_combination::(world, choose(expected, 2)); assert_combination::(world, choose(expected, 5)); assert_combination::(world, choose(expected, 43)); - assert_combination::(world, choose(expected, 128)); + assert_combination::(world, choose(expected, 64)); + } + fn assert_all_exact_sizes_iterator_equal( + iterator: impl ExactSizeIterator, + expected_size: usize, + skip: usize, + query_type: &'static str, + ) { + let len = iterator.len(); + println!("len: {len}"); + assert_all_sizes_iterator_equal(iterator, expected_size, skip, query_type); + assert_eq!(len, expected_size); } fn assert_all_sizes_iterator_equal( - iterator: impl Iterator, + mut iterator: impl Iterator, expected_size: usize, + skip: usize, query_type: &'static str, ) { + let expected_size = expected_size.saturating_sub(skip); + for _ in 0..skip { + iterator.next(); + } let size_hint_0 = iterator.size_hint().0; let size_hint_1 = iterator.size_hint().1; // `count` tests that not only it is the expected value, but also @@ -201,12 +128,12 @@ mod tests { let count = iterator.count(); // This will show up when one of the asserts in this function fails println!( - r#"query declared sizes: - for query: {query_type} - expected: {expected_size} - size_hint().0: {size_hint_0} - size_hint().1: {size_hint_1:?} - count(): {count}"# + "query declared sizes: \n\ + for query: {query_type} \n\ + expected: {expected_size} \n\ + size_hint().0: {size_hint_0} \n\ + size_hint().1: {size_hint_1:?} \n\ + count(): {count}" ); assert_eq!(size_hint_0, expected_size); assert_eq!(size_hint_1, Some(expected_size));