Allow disjoint mutable world access via EntityMut

crates/bevy_ecs/src/lib.rs

@@ -50,7 +50,7 @@ pub mod prelude {
             Commands, Deferred, In, IntoSystem, Local, NonSend, NonSendMut, ParallelCommands,
             ParamSet, Query, ReadOnlySystem, Res, ResMut, Resource, System, SystemParamFunction,
         },
-        world::{EntityRef, FromWorld, World},
+        world::{EntityBorrow, EntityBorrowMut, EntityRef, FromWorld, World},
     };
 }
 
@@ -70,7 +70,7 @@ mod tests {
         entity::Entity,
         query::{Added, Changed, FilteredAccess, ReadOnlyWorldQuery, With, Without},
         system::Resource,
-        world::{EntityRef, Mut, World},
+        world::{Mut, World},
     };
     use bevy_tasks::{ComputeTaskPool, TaskPool};
     use std::{
@@ -1323,27 +1323,13 @@ mod tests {
         world.query::<(&A, &mut A)>();
     }
 
-    #[test]
-    #[should_panic]
-    fn entity_ref_and_mut_query_panic() {
-        let mut world = World::new();
-        world.query::<(EntityRef, &mut A)>();
-    }
-
     #[test]
     #[should_panic]
     fn mut_and_ref_query_panic() {
         let mut world = World::new();
         world.query::<(&mut A, &A)>();
     }
 
-    #[test]
-    #[should_panic]
-    fn mut_and_entity_ref_query_panic() {
-        let mut world = World::new();
-        world.query::<(&mut A, EntityRef)>();
-    }
-
     #[test]
     #[should_panic]
     fn mut_and_mut_query_panic() {

crates/bevy_ecs/src/query/access.rs

@@ -53,9 +53,11 @@ pub struct Access<T: SparseSetIndex> {
     reads_and_writes: FixedBitSet,
     /// The exclusively-accessed elements.
     writes: FixedBitSet,
-    /// Is `true` if this has access to all elements in the collection?
+    /// Is `true` if this has access to all elements in the collection.
     /// This field is a performance optimization for `&World` (also harder to mess up for soundness).
     reads_all: bool,
+    /// Is `true` if this has mutable access to all elements in the collection.
+    writes_all: bool,
     marker: PhantomData<T>,
 }
 
@@ -83,6 +85,7 @@ impl<T: SparseSetIndex> Access<T> {
     pub const fn new() -> Self {
         Self {
             reads_all: false,
+            writes_all: false,
             reads_and_writes: FixedBitSet::new(),
             writes: FixedBitSet::new(),
             marker: PhantomData,
@@ -113,39 +116,58 @@ impl<T: SparseSetIndex> Access<T> {
 
     /// Returns `true` if this can access the element given by `index`.
     pub fn has_read(&self, index: T) -> bool {
-        self.reads_all || self.reads_and_writes.contains(index.sparse_set_index())
+        self.reads_all
+            || self.writes_all
+            || self.reads_and_writes.contains(index.sparse_set_index())
+    }
+
+    /// Returns `true` if this can access anything.
+    pub fn has_any_read(&self) -> bool {
+        self.reads_all || self.writes_all || !self.reads_and_writes.is_clear()
     }
 
     /// Returns `true` if this can exclusively access the element given by `index`.
     pub fn has_write(&self, index: T) -> bool {
-        self.writes.contains(index.sparse_set_index())
+        self.writes_all || self.writes.contains(index.sparse_set_index())
     }
 
     /// Returns `true` if this accesses anything mutably.
     pub fn has_any_write(&self) -> bool {
-        !self.writes.is_clear()
+        self.writes_all || !self.writes.is_clear()
     }
 
     /// Sets this as having access to all indexed elements (i.e. `&World`).
     pub fn read_all(&mut self) {
         self.reads_all = true;
     }
 
+    /// Sets this as having mutable access to all indexed elements (i.e. `EntityBorrowMut`).
+    pub fn write_all(&mut self) {
+        self.writes_all = true;
+    }
+
     /// Returns `true` if this has access to all indexed elements (i.e. `&World`).
     pub fn has_read_all(&self) -> bool {
-        self.reads_all
+        self.reads_all || self.writes_all
+    }
+
+    /// Returns `true` if this has write access to all indexed elements (i.e. `EntityBorrowMut`).
+    pub fn has_write_all(&self) -> bool {
+        self.writes_all
     }
 
     /// Removes all accesses.
     pub fn clear(&mut self) {
         self.reads_all = false;
+        self.writes_all = false;
         self.reads_and_writes.clear();
         self.writes.clear();
     }
 
     /// Adds all access from `other`.
     pub fn extend(&mut self, other: &Access<T>) {
         self.reads_all = self.reads_all || other.reads_all;
+        self.writes_all = self.writes_all || other.writes_all;
         self.reads_and_writes.union_with(&other.reads_and_writes);
         self.writes.union_with(&other.writes);
     }
@@ -157,11 +179,19 @@ impl<T: SparseSetIndex> Access<T> {
     pub fn is_compatible(&self, other: &Access<T>) -> bool {
         // Only systems that do not write data are compatible with systems that operate on `&World`.
         if self.reads_all {
-            return other.writes.count_ones(..) == 0;
+            return !other.has_any_write();
         }
 
         if other.reads_all {
-            return self.writes.count_ones(..) == 0;
+            return !self.has_any_write();
+        }
+
+        if self.writes_all {
+            return !other.has_any_read();
+        }
+
+        if other.writes_all {
+            return !self.has_any_read();
         }
 
         self.writes.is_disjoint(&other.reads_and_writes)
@@ -172,12 +202,23 @@ impl<T: SparseSetIndex> Access<T> {
     pub fn get_conflicts(&self, other: &Access<T>) -> Vec<T> {
         let mut conflicts = FixedBitSet::default();
         if self.reads_all {
+            // QUESTION: How to handle `other.writes_all`?
             conflicts.extend(other.writes.ones());
         }
 
         if other.reads_all {
+            // QUESTION: How to handle `self.writes_all`.
             conflicts.extend(self.writes.ones());
         }
+
+        if self.writes_all {
+            conflicts.extend(other.reads_and_writes.ones());
+        }
+
+        if other.writes_all {
+            conflicts.extend(self.reads_and_writes.ones());
+        }
+
         conflicts.extend(self.writes.intersection(&other.reads_and_writes));
         conflicts.extend(self.reads_and_writes.intersection(&other.writes));
         conflicts
@@ -377,6 +418,11 @@ impl<T: SparseSetIndex> FilteredAccess<T> {
     pub fn read_all(&mut self) {
         self.access.read_all();
     }
+
+    /// Sets the underlying unfiltered access as having mutable access to all indexed elements.
+    pub fn write_all(&mut self) {
+        self.access.write_all();
+    }
 }
 
 #[derive(Clone, Eq, PartialEq)]

crates/bevy_ecs/src/query/fetch.rs

@@ -5,7 +5,7 @@ use crate::{
     entity::Entity,
     query::{Access, DebugCheckedUnwrap, FilteredAccess},
     storage::{ComponentSparseSet, Table, TableRow},
-    world::{unsafe_world_cell::UnsafeWorldCell, EntityRef, Mut, Ref, World},
+    world::{unsafe_world_cell::UnsafeWorldCell, EntityBorrow, EntityBorrowMut, Mut, Ref, World},
 };
 pub use bevy_ecs_macros::WorldQuery;
 use bevy_ptr::{ThinSlicePtr, UnsafeCellDeref};
@@ -524,9 +524,9 @@ unsafe impl WorldQuery for Entity {
 unsafe impl ReadOnlyWorldQuery for Entity {}
 
 /// SAFETY: `Self` is the same as `Self::ReadOnly`
-unsafe impl<'a> WorldQuery for EntityRef<'a> {
-    type Fetch<'w> = &'w World;
-    type Item<'w> = EntityRef<'w>;
+unsafe impl WorldQuery for EntityBorrow<'_> {
+    type Fetch<'w> = UnsafeWorldCell<'w>;
+    type Item<'w> = EntityBorrow<'w>;
     type ReadOnly = Self;
     type State = ();
 
@@ -544,8 +544,7 @@ unsafe impl<'a> WorldQuery for EntityRef<'a> {
         _last_run: Tick,
         _this_run: Tick,
     ) -> Self::Fetch<'w> {
-        // SAFE: EntityRef has permission to access the whole world immutably thanks to update_component_access and update_archetype_component_access
-        world.world()
+        world
     }
 
     #[inline]
@@ -568,13 +567,15 @@ unsafe impl<'a> WorldQuery for EntityRef<'a> {
         _table_row: TableRow,
     ) -> Self::Item<'w> {
         // SAFETY: `fetch` must be called with an entity that exists in the world
-        unsafe { world.get_entity(entity).debug_checked_unwrap() }
+        let cell = world.get_entity(entity).debug_checked_unwrap();
+        // SAFETY: Read-only access to every component has been registered.
+        EntityBorrow::new(cell)
     }
 
     fn update_component_access(_state: &Self::State, access: &mut FilteredAccess<ComponentId>) {
         assert!(
             !access.access().has_any_write(),
-            "EntityRef conflicts with a previous access in this query. Shared access cannot coincide with exclusive access.",
+            "EntityBorrow conflicts with a previous access in this query. Shared access cannot coincide with exclusive access.",
         );
         access.read_all();
     }
@@ -599,8 +600,85 @@ unsafe impl<'a> WorldQuery for EntityRef<'a> {
     }
 }
 
-/// SAFETY: access is read only
-unsafe impl<'a> ReadOnlyWorldQuery for EntityRef<'a> {}
+/// SAFETY: Access is read-only.
+unsafe impl ReadOnlyWorldQuery for EntityBorrow<'_> {}
+
+/// SAFETY: The accesses of `Self::ReadOnly` are a subset of the accesses of `Self`
+unsafe impl<'a> WorldQuery for EntityBorrowMut<'a> {
+    type Fetch<'w> = UnsafeWorldCell<'w>;
+    type Item<'w> = EntityBorrowMut<'w>;
+    type ReadOnly = EntityBorrow<'a>;
+    type State = ();
+
+    fn shrink<'wlong: 'wshort, 'wshort>(item: Self::Item<'wlong>) -> Self::Item<'wshort> {
+        item
+    }
+
+    const IS_DENSE: bool = true;
+
+    const IS_ARCHETYPAL: bool = true;
+
+    unsafe fn init_fetch<'w>(
+        world: UnsafeWorldCell<'w>,
+        _state: &Self::State,
+        _last_run: Tick,
+        _this_run: Tick,
+    ) -> Self::Fetch<'w> {
+        world
+    }
+
+    #[inline]
+    unsafe fn set_archetype<'w>(
+        _fetch: &mut Self::Fetch<'w>,
+        _state: &Self::State,
+        _archetype: &'w Archetype,
+        _table: &Table,
+    ) {
+    }
+
+    #[inline]
+    unsafe fn set_table<'w>(_fetch: &mut Self::Fetch<'w>, _state: &Self::State, _table: &'w Table) {
+    }
+
+    #[inline(always)]
+    unsafe fn fetch<'w>(
+        world: &mut Self::Fetch<'w>,
+        entity: Entity,
+        _table_row: TableRow,
+    ) -> Self::Item<'w> {
+        // SAFETY: `fetch` must be called with an entity that exists in the world
+        let cell = world.get_entity(entity).debug_checked_unwrap();
+        // SAFETY: mutable access to every component has been registered.
+        EntityBorrowMut::new(cell)
+    }
+
+    fn update_component_access(_state: &Self::State, access: &mut FilteredAccess<ComponentId>) {
+        assert!(
+            !access.access().has_any_read(),
+            "EntityBorrowMut conflicts with a previous access in this query. Exclusive access cannot coincide with any other accesses.",
+        );
+        access.write_all();
+    }
+
+    fn update_archetype_component_access(
+        _state: &Self::State,
+        archetype: &Archetype,
+        access: &mut Access<ArchetypeComponentId>,
+    ) {
+        for component_id in archetype.components() {
+            access.add_write(archetype.get_archetype_component_id(component_id).unwrap());
+        }
+    }
+
+    fn init_state(_world: &mut World) {}
+
+    fn matches_component_set(
+        _state: &Self::State,
+        _set_contains_id: &impl Fn(ComponentId) -> bool,
+    ) -> bool {
+        true
+    }
+}
 
 #[doc(hidden)]
 pub struct ReadFetch<'w, T> {