[Merged by Bors] - REMOVE unsound lifetime annotations on EntityMut

crates/bevy_ecs/src/world/entity_ref.rs

@@ -70,9 +70,16 @@ impl<'w> EntityRef<'w> {
         }
     }
 
+    /// Gets a mutable reference to the component of type `T` associated with
+    /// this entity without ensuring there are no other borrows active and without
+    /// ensuring that the returned reference will stay valid.
+    ///
     /// # Safety
-    /// This allows aliased mutability. You must make sure this call does not result in multiple
-    /// mutable references to the same component
+    ///
+    /// - The returned reference must never alias a mutable borrow of this component.
+    /// - The returned reference must not be used after this component is moved which
+    ///   may happen from **any** `insert_component`, `remove_component` or `despawn`
+    ///   operation on this world (non-exhaustive list).
     #[inline]
     pub unsafe fn get_unchecked_mut<T: Component>(
         &self,
@@ -145,39 +152,44 @@ impl<'w> EntityMut<'w> {
     }
 
     #[inline]
-    pub fn get<T: Component>(&self) -> Option<&'w T> {
-        // SAFE: entity location is valid and returned component is of type T
-        unsafe {
-            get_component_with_type(self.world, TypeId::of::<T>(), self.entity, self.location)
-                .map(|value| &*value.cast::<T>())
-        }
+    pub fn get<T: Component>(&self) -> Option<&'_ T> {
+        // SAFE: lifetimes enforce correct usage of returned borrow
+        unsafe { self.get_unchecked::<T>() }
     }
 
     #[inline]
-    pub fn get_mut<T: Component>(&mut self) -> Option<Mut<'w, T>> {
-        // SAFE: world access is unique, entity location is valid, and returned component is of type
-        // T
-        unsafe {
-            get_component_and_ticks_with_type(
-                self.world,
-                TypeId::of::<T>(),
-                self.entity,
-                self.location,
-            )
-            .map(|(value, ticks)| Mut {
-                value: &mut *value.cast::<T>(),
-                ticks: Ticks {
-                    component_ticks: &mut *ticks,
-                    last_change_tick: self.world.last_change_tick(),
-                    change_tick: self.world.change_tick(),
-                },
-            })
-        }
+    pub fn get_mut<T: Component>(&mut self) -> Option<Mut<'_, T>> {
+        // SAFE: world access is unique, and lifetimes enforce correct usage of returned borrow
+        unsafe { self.get_unchecked_mut::<T>() }
     }
 
+    /// Gets an immutable reference to the component of type `T` associated with
+    /// this entity without ensuring there are no unique borrows active and without
+    /// ensuring that the returned reference will stay valid.
+    ///
     /// # Safety
-    /// This allows aliased mutability. You must make sure this call does not result in multiple
-    /// mutable references to the same component
+    ///
+    /// - The returned reference must never alias a mutable borrow of this component.
+    /// - The returned reference must not be used after this component is moved which
+    ///   may happen from **any** `insert_component`, `remove_component` or `despawn`
+    ///   operation on this world (non-exhaustive list).
+    #[inline]
+    pub unsafe fn get_unchecked<T: Component>(&self) -> Option<&'w T> {
+        // SAFE: entity location is valid and returned component is of type T
+        get_component_with_type(self.world, TypeId::of::<T>(), self.entity, self.location)
+            .map(|value| &*value.cast::<T>())
+    }
+
+    /// Gets a mutable reference to the component of type `T` associated with
+    /// this entity without ensuring there are no other borrows active and without
+    /// ensuring that the returned reference will stay valid.
+    ///
+    /// # Safety
+    ///
+    /// - The returned reference must never alias a mutable borrow of this component.
+    /// - The returned reference must not be used after this component is moved which
+    ///   may happen from **any** `insert_component`, `remove_component` or `despawn`
+    ///   operation on this world (non-exhaustive list).
     #[inline]
     pub unsafe fn get_unchecked_mut<T: Component>(&self) -> Option<Mut<'w, T>> {
         get_component_and_ticks_with_type(self.world, TypeId::of::<T>(), self.entity, self.location)

crates/bevy_ecs/src/world/mod.rs

@@ -224,8 +224,8 @@ impl World {
     /// let entity = world.spawn()
     ///     .insert(Position { x: 0.0, y: 0.0 })
     ///     .id();
-    ///
-    /// let mut position = world.entity_mut(entity).get_mut::<Position>().unwrap();
+    /// let mut entity_mut = world.entity_mut(entity);
+    /// let mut position = entity_mut.get_mut::<Position>().unwrap();
     /// position.x = 1.0;
     /// ```
     #[inline]
@@ -437,7 +437,8 @@ impl World {
     /// ```
     #[inline]
     pub fn get_mut<T: Component>(&mut self, entity: Entity) -> Option<Mut<T>> {
-        self.get_entity_mut(entity)?.get_mut()
+        // SAFE: lifetimes enforce correct usage of returned borrow
+        unsafe { self.get_entity_mut(entity)?.get_unchecked_mut::<T>() }
     }
 
     /// Despawns the given `entity`, if it exists. This will also remove all of the entity's

crates/bevy_ecs_compile_fail_tests/tests/ui/entity_ref_mut_lifetime_safety.rs

@@ -0,0 +1,64 @@
+use bevy_ecs::prelude::*;
+
+#[derive(Component, Eq, PartialEq, Debug)]
+struct A(Box<usize>);
+
+#[derive(Component)]
+struct B;
+
+fn main() {
+    let mut world = World::default();
+    let e = world.spawn().insert(A(Box::new(10_usize))).id();
+
+    let mut e_mut = world.entity_mut(e);
+
+    {
+        let gotten: &A = e_mut.get::<A>().unwrap();
+        let gotten2: A = e_mut.remove::<A>().unwrap();
+        assert_eq!(gotten, &gotten2); // oops UB
+    }
+
+    e_mut.insert(A(Box::new(12_usize)));
+
+    {
+        let mut gotten: Mut<A> = e_mut.get_mut::<A>().unwrap();
+        let mut gotten2: A = e_mut.remove::<A>().unwrap();
+        assert_eq!(&mut *gotten, &mut gotten2); // oops UB
+    }
+
+    e_mut.insert(A(Box::new(14_usize)));
+
+    {
+        let gotten: &A = e_mut.get::<A>().unwrap();
+        e_mut.despawn();
+        assert_eq!(gotten, &A(Box::new(14_usize))); // oops UB
+    }
+
+    let e = world.spawn().insert(A(Box::new(16_usize))).id();
+    let mut e_mut = world.entity_mut(e);
+
+    {
+        let gotten: &A = e_mut.get::<A>().unwrap();
+        let gotten_mut: Mut<A> = e_mut.get_mut::<A>().unwrap();
+        assert_eq!(gotten, &*gotten_mut); // oops UB
+    }
+
+    {
+        let gotten_mut: Mut<A> = e_mut.get_mut::<A>().unwrap();
+        let gotten: &A = e_mut.get::<A>().unwrap();
+        assert_eq!(gotten, &*gotten_mut); // oops UB
+    }
+
+    {
+        let gotten: &A = e_mut.get::<A>().unwrap();
+        e_mut.insert::<B>(B);
+        assert_eq!(gotten, &A(Box::new(16_usize))); // oops UB
+        e_mut.remove::<B>();
+    }
+
+    {
+        let mut gotten_mut: Mut<A> = e_mut.get_mut::<A>().unwrap();
+        e_mut.insert::<B>(B);
+        assert_eq!(&mut *gotten_mut, &mut A(Box::new(16_usize))); // oops UB
+    }
+}

crates/bevy_ecs_compile_fail_tests/tests/ui/entity_ref_mut_lifetime_safety.stderr

@@ -0,0 +1,69 @@
+error[E0502]: cannot borrow `e_mut` as mutable because it is also borrowed as immutable
+  --> tests/ui/entity_ref_mut_lifetime_safety.rs:17:26
+   |
+16 |         let gotten: &A = e_mut.get::<A>().unwrap();
+   |                          ---------------- immutable borrow occurs here
+17 |         let gotten2: A = e_mut.remove::<A>().unwrap();
+   |                          ^^^^^^^^^^^^^^^^^^^ mutable borrow occurs here
+18 |         assert_eq!(gotten, &gotten2); // oops UB
+   |         ---------------------------- immutable borrow later used here
+
+error[E0499]: cannot borrow `e_mut` as mutable more than once at a time
+  --> tests/ui/entity_ref_mut_lifetime_safety.rs:25:30
+   |
+24 |         let mut gotten: Mut<A> = e_mut.get_mut::<A>().unwrap();
+   |                                  -------------------- first mutable borrow occurs here
+25 |         let mut gotten2: A = e_mut.remove::<A>().unwrap();
+   |                              ^^^^^^^^^^^^^^^^^^^ second mutable borrow occurs here
+26 |         assert_eq!(&mut *gotten, &mut gotten2); // oops UB
+   |                          ------ first borrow later used here
+
+error[E0505]: cannot move out of `e_mut` because it is borrowed
+  --> tests/ui/entity_ref_mut_lifetime_safety.rs:33:9
+   |
+32 |         let gotten: &A = e_mut.get::<A>().unwrap();
+   |                          ---------------- borrow of `e_mut` occurs here
+33 |         e_mut.despawn();
+   |         ^^^^^ move out of `e_mut` occurs here
+34 |         assert_eq!(gotten, &A(Box::new(14_usize))); // oops UB
+   |         ------------------------------------------ borrow later used here
+
+error[E0502]: cannot borrow `e_mut` as mutable because it is also borrowed as immutable
+  --> tests/ui/entity_ref_mut_lifetime_safety.rs:42:34
+   |
+41 |         let gotten: &A = e_mut.get::<A>().unwrap();
+   |                          ---------------- immutable borrow occurs here
+42 |         let gotten_mut: Mut<A> = e_mut.get_mut::<A>().unwrap();
+   |                                  ^^^^^^^^^^^^^^^^^^^^ mutable borrow occurs here
+43 |         assert_eq!(gotten, &*gotten_mut); // oops UB
+   |         -------------------------------- immutable borrow later used here
+
+error[E0502]: cannot borrow `e_mut` as immutable because it is also borrowed as mutable
+  --> tests/ui/entity_ref_mut_lifetime_safety.rs:48:26
+   |
+47 |         let gotten_mut: Mut<A> = e_mut.get_mut::<A>().unwrap();
+   |                                  -------------------- mutable borrow occurs here
+48 |         let gotten: &A = e_mut.get::<A>().unwrap();
+   |                          ^^^^^^^^^^^^^^^^ immutable borrow occurs here
+49 |         assert_eq!(gotten, &*gotten_mut); // oops UB
+   |                              ---------- mutable borrow later used here
+
+error[E0502]: cannot borrow `e_mut` as mutable because it is also borrowed as immutable
+  --> tests/ui/entity_ref_mut_lifetime_safety.rs:54:9
+   |
+53 |         let gotten: &A = e_mut.get::<A>().unwrap();
+   |                          ---------------- immutable borrow occurs here
+54 |         e_mut.insert::<B>(B);
+   |         ^^^^^^^^^^^^^^^^^^^^ mutable borrow occurs here
+55 |         assert_eq!(gotten, &A(Box::new(16_usize))); // oops UB
+   |         ------------------------------------------ immutable borrow later used here
+
+error[E0499]: cannot borrow `e_mut` as mutable more than once at a time
+  --> tests/ui/entity_ref_mut_lifetime_safety.rs:61:9
+   |
+60 |         let mut gotten_mut: Mut<A> = e_mut.get_mut::<A>().unwrap();
+   |                                      -------------------- first mutable borrow occurs here
+61 |         e_mut.insert::<B>(B);
+   |         ^^^^^^^^^^^^^^^^^^^^ second mutable borrow occurs here
+62 |         assert_eq!(&mut *gotten_mut, &mut A(Box::new(16_usize))); // oops UB
+   |                          ---------- first borrow later used here