Fix query transmute from table to archetype iteration unsoundness (#1…

…4615)

# Objective

- Fixes #14348 
- Fixes #14528
- Less complex (but also likely less performant) alternative to #14611

## Solution

- Add a `is_dense` field flag to `QueryIter` indicating whether it is
dense or not, that is whether it can perform dense iteration or not;
- Check this flag any time iteration over a query is performed.

---

It would be nice if someone could try benching this change to see if it
actually matters.

~Note that this not 100% ready for mergin, since there are a bunch of
safety comments on the use of the various `IS_DENSE` for checks that
still need to be updated.~ This is ready modulo benchmarks

---------

Co-authored-by: Alice Cecile <alice.i.cecile@gmail.com>

crates/bevy_ecs/src/query/builder.rs

@@ -1,5 +1,6 @@
 use std::marker::PhantomData;
 
+use crate::component::StorageType;
 use crate::{component::ComponentId, prelude::*};
 
 use super::{FilteredAccess, QueryData, QueryFilter};
@@ -68,6 +69,26 @@ impl<'w, D: QueryData, F: QueryFilter> QueryBuilder<'w, D, F> {
  }
  }
 
+ pub(super) fn is_dense(&self) -> bool {
+ // Note: `component_id` comes from the user in safe code, so we cannot trust it to
+ // exist. If it doesn't exist we pessimistically assume it's sparse.
+ let is_dense = |component_id| {
+ self.world()
+ .components()
+ .get_info(component_id)
+ .map_or(false, |info| info.storage_type() == StorageType::Table)
+ };
+
+ self.access
+ .access()
+ .component_reads_and_writes()
+ .all(is_dense)
+ && self.access.access().archetypal().all(is_dense)
+ && !self.access.access().has_read_all_components()
+ && self.access.with_filters().all(is_dense)
+ && self.access.without_filters().all(is_dense)
+ }
+
  /// Returns a reference to the world passed to [`Self::new`].
  pub fn world(&self) -> &World {
  self.world
@@ -396,4 +417,27 @@ mod tests {
  assert_eq!(1, b.deref::<B>().0);
  }
  }
+
+ /// Regression test for issue #14348
+ #[test]
+ fn builder_static_dense_dynamic_sparse() {
+ #[derive(Component)]
+ struct Dense;
+
+ #[derive(Component)]
+ #[component(storage = "SparseSet")]
+ struct Sparse;
+
+ let mut world = World::new();
+
+ world.spawn(Dense);
+ world.spawn((Dense, Sparse));
+
+ let mut query = QueryBuilder::<&Dense>::new(&mut world)
+ .with::<Sparse>()
+ .build();
+
+ let matched = query.iter(&world).count();
+ assert_eq!(matched, 1);
+ }
 }

crates/bevy_ecs/src/query/iter.rs

@@ -128,7 +128,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
  /// # Safety
  /// - all `rows` must be in `[0, table.entity_count)`.
  /// - `table` must match D and F
- /// - Both `D::IS_DENSE` and `F::IS_DENSE` must be true.
+ /// - The query iteration must be dense (i.e. `self.query_state.is_dense` must be true).
  #[inline]
  pub(super) unsafe fn fold_over_table_range<B, Func>(
  &mut self,
@@ -183,7 +183,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
  /// # Safety
  /// - all `indices` must be in `[0, archetype.len())`.
  /// - `archetype` must match D and F
- /// - Either `D::IS_DENSE` or `F::IS_DENSE` must be false.
+ /// - The query iteration must not be dense (i.e. `self.query_state.is_dense` must be false).
  #[inline]
  pub(super) unsafe fn fold_over_archetype_range<B, Func>(
  &mut self,
@@ -252,7 +252,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
  /// - all `indices` must be in `[0, archetype.len())`.
  /// - `archetype` must match D and F
  /// - `archetype` must have the same length with it's table.
- /// - Either `D::IS_DENSE` or `F::IS_DENSE` must be false.
+ /// - The query iteration must not be dense (i.e. `self.query_state.is_dense` must be false).
  #[inline]
  pub(super) unsafe fn fold_over_dense_archetype_range<B, Func>(
  &mut self,
@@ -1031,40 +1031,47 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Iterator for QueryIter<'w, 's, D, F>
  let Some(item) = self.next() else { break };
  accum = func(accum, item);
  }
- for id in self.cursor.storage_id_iter.clone() {
- if D::IS_DENSE && F::IS_DENSE {
+
+ if self.cursor.is_dense {
+ for id in self.cursor.storage_id_iter.clone() {
+ // SAFETY: `self.cursor.is_dense` is true, so storage ids are guaranteed to be table ids.
+ let table_id = unsafe { id.table_id };
  // SAFETY: Matched table IDs are guaranteed to still exist.
- let table = unsafe { self.tables.get(id.table_id).debug_checked_unwrap() };
+ let table = unsafe { self.tables.get(table_id).debug_checked_unwrap() };
+
  accum =
  // SAFETY: 
  // - The fetched table matches both D and F
  // - The provided range is equivalent to [0, table.entity_count)
- // - The if block ensures that D::IS_DENSE and F::IS_DENSE are both true
+ // - The if block ensures that the query iteration is dense
  unsafe { self.fold_over_table_range(accum, &mut func, table, 0..table.entity_count()) };
- } else {
- let archetype =
- // SAFETY: Matched archetype IDs are guaranteed to still exist.
- unsafe { self.archetypes.get(id.archetype_id).debug_checked_unwrap() };
+ }
+ } else {
+ for id in self.cursor.storage_id_iter.clone() {
+ // SAFETY: `self.cursor.is_dense` is false, so storage ids are guaranteed to be archetype ids.
+ let archetype_id = unsafe { id.archetype_id };
+ // SAFETY: Matched archetype IDs are guaranteed to still exist.
+ let archetype = unsafe { self.archetypes.get(archetype_id).debug_checked_unwrap() };
  // SAFETY: Matched table IDs are guaranteed to still exist.
  let table = unsafe { self.tables.get(archetype.table_id()).debug_checked_unwrap() };
 
  // When an archetype and its table have equal entity counts, dense iteration can be safely used.
  // this leverages cache locality to optimize performance.
  if table.entity_count() == archetype.len() {
  accum =
- // SAFETY:
- // - The fetched archetype matches both D and F
- // - The provided archetype and its' table have the same length.
- // - The provided range is equivalent to [0, archetype.len)
- // - The if block ensures that ether D::IS_DENSE or F::IS_DENSE are false
- unsafe { self.fold_over_dense_archetype_range(accum, &mut func, archetype,0..archetype.len()) };
+  // SAFETY:
+  // - The fetched archetype matches both D and F
+  // - The provided archetype and its' table have the same length.
+  // - The provided range is equivalent to [0, archetype.len)
+  // - The if block ensures that the query iteration is not dense.
+  unsafe { self.fold_over_dense_archetype_range(accum, &mut func, archetype, 0..archetype.len()) };
  } else {
  accum =
- // SAFETY:
- // - The fetched archetype matches both D and F
- // - The provided range is equivalent to [0, archetype.len)
- // - The if block ensures that ether D::IS_DENSE or F::IS_DENSE are false
- unsafe { self.fold_over_archetype_range(accum, &mut func, archetype,0..archetype.len()) };
+  // SAFETY:
+  // - The fetched archetype matches both D and F
+  // - The provided range is equivalent to [0, archetype.len)
+  // - The if block ensures that the query iteration is not dense.
+  unsafe { self.fold_over_archetype_range(accum, &mut func, archetype, 0..archetype.len()) };
  }
  }
  }
@@ -1675,6 +1682,8 @@ impl<'w, 's, D: QueryData, F: QueryFilter, const K: usize> Debug
 }
 
 struct QueryIterationCursor<'w, 's, D: QueryData, F: QueryFilter> {
+ // whether the query iteration is dense or not. Mirrors QueryState's `is_dense` field.
+ is_dense: bool,
  storage_id_iter: std::slice::Iter<'s, StorageId>,
  table_entities: &'w [Entity],
  archetype_entities: &'w [ArchetypeEntity],
@@ -1689,6 +1698,7 @@ struct QueryIterationCursor<'w, 's, D: QueryData, F: QueryFilter> {
 impl<D: QueryData, F: QueryFilter> Clone for QueryIterationCursor<'_, '_, D, F> {
  fn clone(&self) -> Self {
  Self {
+ is_dense: self.is_dense,
  storage_id_iter: self.storage_id_iter.clone(),
  table_entities: self.table_entities,
  archetype_entities: self.archetype_entities,
@@ -1701,8 +1711,6 @@ impl<D: QueryData, F: QueryFilter> Clone for QueryIterationCursor<'_, '_, D, F>
 }
 
 impl<'w, 's, D: QueryData, F: QueryFilter> QueryIterationCursor<'w, 's, D, F> {
- const IS_DENSE: bool = D::IS_DENSE && F::IS_DENSE;
-
  unsafe fn init_empty(
  world: UnsafeWorldCell<'w>,
  query_state: &'s QueryState<D, F>,
@@ -1732,13 +1740,15 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIterationCursor<'w, 's, D, F> {
  table_entities: &[],
  archetype_entities: &[],
  storage_id_iter: query_state.matched_storage_ids.iter(),
+ is_dense: query_state.is_dense,
  current_len: 0,
  current_row: 0,
  }
  }
 
  fn reborrow(&mut self) -> QueryIterationCursor<'_, 's, D, F> {
  QueryIterationCursor {
+ is_dense: self.is_dense,
  fetch: D::shrink_fetch(self.fetch.clone()),
  filter: F::shrink_fetch(self.filter.clone()),
  table_entities: self.table_entities,
@@ -1754,7 +1764,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIterationCursor<'w, 's, D, F> {
  unsafe fn peek_last(&mut self) -> Option<D::Item<'w>> {
  if self.current_row > 0 {
  let index = self.current_row - 1;
- if Self::IS_DENSE {
+ if self.is_dense {
  let entity = self.table_entities.get_unchecked(index);
  Some(D::fetch(
  &mut self.fetch,
@@ -1780,7 +1790,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIterationCursor<'w, 's, D, F> {
  /// will be **the exact count of remaining values**.
  fn max_remaining(&self, tables: &'w Tables, archetypes: &'w Archetypes) -> usize {
  let ids = self.storage_id_iter.clone();
- let remaining_matched: usize = if Self::IS_DENSE {
+ let remaining_matched: usize = if self.is_dense {
  // SAFETY: The if check ensures that storage_id_iter stores TableIds
  unsafe { ids.map(|id| tables[id.table_id].entity_count()).sum() }
  } else {
@@ -1803,7 +1813,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIterationCursor<'w, 's, D, F> {
  archetypes: &'w Archetypes,
  query_state: &'s QueryState<D, F>,
  ) -> Option<D::Item<'w>> {
- if Self::IS_DENSE {
+ if self.is_dense {
  loop {
  // we are on the beginning of the query, or finished processing a table, so skip to the next
  if self.current_row == self.current_len {

crates/bevy_ecs/src/query/par_iter.rs

@@ -126,7 +126,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryParIter<'w, 's, D, F> {
  fn get_batch_size(&self, thread_count: usize) -> usize {
  let max_items = || {
  let id_iter = self.state.matched_storage_ids.iter();
- if D::IS_DENSE && F::IS_DENSE {
+ if self.state.is_dense {
  // SAFETY: We only access table metadata.
  let tables = unsafe { &self.world.world_metadata().storages().tables };
  id_iter