[Merged by Bors] - Implement iter() for mutable Queries

crates/bevy_ecs/src/query/fetch.rs

@@ -43,9 +43,12 @@ use std::{
 pub trait WorldQuery {
     type Fetch: for<'world, 'state> Fetch<'world, 'state, State = Self::State>;
     type State: FetchState;
+    type ReadOnlyFetch: for<'world, 'state> Fetch<'world, 'state, State = Self::State>
+        + ReadOnlyFetch;
 }
 
 pub type QueryItem<'w, 's, Q> = <<Q as WorldQuery>::Fetch as Fetch<'w, 's>>::Item;
+pub type ReadOnlyQueryItem<'w, 's, Q> = <<Q as WorldQuery>::ReadOnlyFetch as Fetch<'w, 's>>::Item;
 
 pub trait Fetch<'world, 'state>: Sized {
     type Item;
@@ -136,6 +139,7 @@ pub unsafe trait ReadOnlyFetch {}
 impl WorldQuery for Entity {
     type Fetch = EntityFetch;
     type State = EntityState;
+    type ReadOnlyFetch = EntityFetch;
 }
 
 /// The [`Fetch`] of [`Entity`].
@@ -222,6 +226,7 @@ impl<'w, 's> Fetch<'w, 's> for EntityFetch {
 impl<T: Component> WorldQuery for &T {
     type Fetch = ReadFetch<T>;
     type State = ReadState<T>;
+    type ReadOnlyFetch = ReadFetch<T>;
 }
 
 /// The [`FetchState`] of `&T`.
@@ -376,6 +381,7 @@ impl<'w, 's, T: Component> Fetch<'w, 's> for ReadFetch<T> {
 impl<T: Component> WorldQuery for &mut T {
     type Fetch = WriteFetch<T>;
     type State = WriteState<T>;
+    type ReadOnlyFetch = ReadOnlyWriteFetch<T>;
 }
 
 /// The [`Fetch`] of `&mut T`.
@@ -403,6 +409,34 @@ impl<T> Clone for WriteFetch<T> {
     }
 }
 
+/// The [`ReadOnlyFetch`] of `&mut T`.
+pub struct ReadOnlyWriteFetch<T> {
+    table_components: NonNull<T>,
+    table_ticks: *const UnsafeCell<ComponentTicks>,
+    entities: *const Entity,
+    entity_table_rows: *const usize,
+    sparse_set: *const ComponentSparseSet,
+    last_change_tick: u32,
+    change_tick: u32,
+}
+
+/// SAFETY: access is read only
+unsafe impl<T> ReadOnlyFetch for ReadOnlyWriteFetch<T> {}
+
+impl<T> Clone for ReadOnlyWriteFetch<T> {
+    fn clone(&self) -> Self {
+        Self {
+            table_components: self.table_components,
+            table_ticks: self.table_ticks,
+            entities: self.entities,
+            entity_table_rows: self.entity_table_rows,
+            sparse_set: self.sparse_set,
+            last_change_tick: self.last_change_tick,
+            change_tick: self.change_tick,
+        }
+    }
+}
+
 /// The [`FetchState`] of `&mut T`.
 pub struct WriteState<T> {
     component_id: ComponentId,
@@ -555,9 +589,93 @@ impl<'w, 's, T: Component> Fetch<'w, 's> for WriteFetch<T> {
     }
 }
 
+impl<'w, 's, T: Component> Fetch<'w, 's> for ReadOnlyWriteFetch<T> {
+    type Item = &'w T;
+    type State = WriteState<T>;
+
+    const IS_DENSE: bool = {
+        match T::Storage::STORAGE_TYPE {
+            StorageType::Table => true,
+            StorageType::SparseSet => false,
+        }
+    };
+
+    unsafe fn init(
+        world: &World,
+        state: &Self::State,
+        last_change_tick: u32,
+        change_tick: u32,
+    ) -> Self {
+        let mut value = Self {
+            table_components: NonNull::dangling(),
+            entities: ptr::null::<Entity>(),
+            entity_table_rows: ptr::null::<usize>(),
+            sparse_set: ptr::null::<ComponentSparseSet>(),
+            table_ticks: ptr::null::<UnsafeCell<ComponentTicks>>(),
+            last_change_tick,
+            change_tick,
+        };
+        if T::Storage::STORAGE_TYPE == StorageType::SparseSet {
+            value.sparse_set = world
+                .storages()
+                .sparse_sets
+                .get(state.component_id)
+                .unwrap();
+        }
+        value
+    }
+
+    #[inline]
+    unsafe fn set_archetype(
+        &mut self,
+        state: &Self::State,
+        archetype: &Archetype,
+        tables: &Tables,
+    ) {
+        match T::Storage::STORAGE_TYPE {
+            StorageType::Table => {
+                self.entity_table_rows = archetype.entity_table_rows().as_ptr();
+                let column = tables[archetype.table_id()]
+                    .get_column(state.component_id)
+                    .unwrap();
+                self.table_components = column.get_data_ptr().cast::<T>();
+                self.table_ticks = column.get_ticks_ptr();
+            }
+            StorageType::SparseSet => self.entities = archetype.entities().as_ptr(),
+        }
+    }
+
+    #[inline]
+    unsafe fn set_table(&mut self, state: &Self::State, table: &Table) {
+        let column = table.get_column(state.component_id).unwrap();
+        self.table_components = column.get_data_ptr().cast::<T>();
+        self.table_ticks = column.get_ticks_ptr();
+    }
+
+    #[inline]
+    unsafe fn archetype_fetch(&mut self, archetype_index: usize) -> Self::Item {
+        match T::Storage::STORAGE_TYPE {
+            StorageType::Table => {
+                let table_row = *self.entity_table_rows.add(archetype_index);
+                &*self.table_components.as_ptr().add(table_row)
+            }
+            StorageType::SparseSet => {
+                let entity = *self.entities.add(archetype_index);
+                &*(*self.sparse_set).get(entity).unwrap().cast::<T>()
+            }
+        }
+    }
+
+    #[inline]
+    unsafe fn table_fetch(&mut self, table_row: usize) -> Self::Item {
+        &*self.table_components.as_ptr().add(table_row)
+    }
+}
+
 impl<T: WorldQuery> WorldQuery for Option<T> {
     type Fetch = OptionFetch<T::Fetch>;
     type State = OptionState<T::State>;
+    type ReadOnlyFetch = OptionFetch<T::ReadOnlyFetch>;
 }
 
 /// The [`Fetch`] of `Option<T>`.
@@ -733,6 +851,7 @@ impl<T: Component> ChangeTrackers<T> {
 impl<T: Component> WorldQuery for ChangeTrackers<T> {
     type Fetch = ChangeTrackersFetch<T>;
     type State = ChangeTrackersState<T>;
+    type ReadOnlyFetch = ChangeTrackersFetch<T>;
 }
 
 /// The [`FetchState`] of [`ChangeTrackers`].
@@ -983,6 +1102,7 @@ macro_rules! impl_tuple_fetch {
         impl<$($name: WorldQuery),*> WorldQuery for ($($name,)*) {
             type Fetch = ($($name::Fetch,)*);
             type State = ($($name::State,)*);
+            type ReadOnlyFetch = ($($name::ReadOnlyFetch,)*);
         }
 
         /// SAFETY: each item in the tuple is read only
@@ -992,3 +1112,45 @@ macro_rules! impl_tuple_fetch {
 }
 
 all_tuples!(impl_tuple_fetch, 0, 15, F, S);
+
+/// [`Fetch`] that does not actually fetch anything
+///
+/// Mostly useful when something is generic over the Fetch and you don't want to fetch as you will discard the result
+pub struct NopFetch<State> {
+    state: PhantomData<State>,
+}
+
+impl<'w, 's, State: FetchState> Fetch<'w, 's> for NopFetch<State> {
+    type Item = ();
+    type State = State;
+
+    const IS_DENSE: bool = true;
+
+    #[inline(always)]
+    unsafe fn init(
+        _world: &World,
+        _state: &Self::State,
+        _last_change_tick: u32,
+        _change_tick: u32,
+    ) -> Self {
+        Self { state: PhantomData }
+    }
+
+    #[inline(always)]
+    unsafe fn set_archetype(
+        &mut self,
+        _state: &Self::State,
+        _archetype: &Archetype,
+        _tables: &Tables,
+    ) {
+    }
+
+    #[inline(always)]
+    unsafe fn set_table(&mut self, _state: &Self::State, _table: &Table) {}
+
+    #[inline(always)]
+    unsafe fn archetype_fetch(&mut self, _archetype_index: usize) -> Self::Item {}
+
+    #[inline(always)]
+    unsafe fn table_fetch(&mut self, _table_row: usize) -> Self::Item {}
+}

crates/bevy_ecs/src/query/filter.rs

@@ -2,7 +2,7 @@ use crate::{
     archetype::{Archetype, ArchetypeComponentId},
     component::{Component, ComponentId, ComponentStorage, ComponentTicks, StorageType},
     entity::Entity,
-    query::{Access, Fetch, FetchState, FilteredAccess, WorldQuery},
+    query::{Access, Fetch, FetchState, FilteredAccess, ReadOnlyFetch, WorldQuery},
     storage::{ComponentSparseSet, Table, Tables},
     world::World,
 };
@@ -72,6 +72,7 @@ pub struct With<T>(PhantomData<T>);
 impl<T: Component> WorldQuery for With<T> {
     type Fetch = WithFetch<T>;
     type State = WithState<T>;
+    type ReadOnlyFetch = WithFetch<T>;
 }
 
 /// The [`Fetch`] of [`With`].
@@ -162,6 +163,9 @@ impl<'w, 's, T: Component> Fetch<'w, 's> for WithFetch<T> {
     }
 }
 
+// SAFETY: no component access or archetype component access
+unsafe impl<T> ReadOnlyFetch for WithFetch<T> {}
+
 /// Filter that selects entities without a component `T`.
 ///
 /// This is the negation of [`With`].
@@ -191,6 +195,7 @@ pub struct Without<T>(PhantomData<T>);
 impl<T: Component> WorldQuery for Without<T> {
     type Fetch = WithoutFetch<T>;
     type State = WithoutState<T>;
+    type ReadOnlyFetch = WithoutFetch<T>;
 }
 
 /// The [`Fetch`] of [`Without`].
@@ -281,6 +286,9 @@ impl<'w, 's, T: Component> Fetch<'w, 's> for WithoutFetch<T> {
     }
 }
 
+// SAFETY: no component access or archetype component access
+unsafe impl<T> ReadOnlyFetch for WithoutFetch<T> {}
+
 /// A filter that tests if any of the given filters apply.
 ///
 /// This is useful for example if a system with multiple components in a query only wants to run
@@ -342,8 +350,11 @@ macro_rules! impl_query_filter_tuple {
         {
             type Fetch = Or<($(OrFetch<$filter::Fetch>,)*)>;
             type State = Or<($($filter::State,)*)>;
+            type ReadOnlyFetch = Or<($(OrFetch<$filter::Fetch>,)*)>;
         }
 
+        /// SAFETY: this only works using the filter which doesn't write
+        unsafe impl<$($filter: FilterFetch),*> ReadOnlyFetch for Or<($(OrFetch<$filter>,)*)> {}
 
         #[allow(unused_variables)]
         #[allow(non_snake_case)]
@@ -464,9 +475,9 @@ macro_rules! impl_tick_filter {
         impl<T: Component> WorldQuery for $name<T> {
             type Fetch = $fetch_name<T>;
             type State = $state_name<T>;
+            type ReadOnlyFetch = $fetch_name<T>;
         }
 
-
         // SAFETY: this reads the T component. archetype component access and component access are updated to reflect that
         unsafe impl<T: Component> FetchState for $state_name<T> {
             fn init(world: &mut World) -> Self {
@@ -572,6 +583,9 @@ macro_rules! impl_tick_filter {
                 }
             }
         }
+
+        /// SAFETY: read-only access
+        unsafe impl<T: Component> ReadOnlyFetch for $fetch_name<T> {}
     };
 }