From cf46dd2e7ed2085bbd60d25cafca8219d62ef2ec Mon Sep 17 00:00:00 2001
From: Jakob Hellermann <jakob.hellermann@protonmail.com>
Date: Tue, 22 Nov 2022 15:31:18 +0000
Subject: [PATCH] fix mutable aliases for a very short time if `WorldCell` is
 already borrowed (#6639)

# Objective

Consider the test
```rust
let cell = world.cell();
let _value_a = cell.resource_mut::<A>();
let _value_b = cell.resource_mut::<A>();
```

Currently, this will roughly execute

```rust
// first call
let value = unsafe {
    self.world
    .get_non_send_unchecked_mut_with_id(component_id)?
};
return Some(WorldBorrowMut::new(value, archetype_component_id, self.access)))

// second call
let value = unsafe {
    self.world
    .get_non_send_unchecked_mut_with_id(component_id)?
};
return Some(WorldBorrowMut::new(value, archetype_component_id, self.access)))
```
where `WorldBorrowMut::new` will panic if the resource is already borrowed.

This means, that `_value_a` will be created, the access checked (OK), then `value_b` will be created, and the access checked (`panic`).
For a moment, both `_value_a` and `_value_b` existed as `&mut T` to the same location, which is insta-UB as far as I understand it.

## Solution
Flip the order so that `WorldBorrowMut::new` first checks the access, _then_ fetches creates the value. To do that, we pass a `impl FnOnce() -> Mut<T>` instead of the `Mut<T>` directly:

```rust
let get_value = || unsafe {
    self.world
    .get_non_send_unchecked_mut_with_id(component_id)?
};
return Some(WorldBorrowMut::new(get_value, archetype_component_id, self.access)))
```
---
 crates/bevy_ecs/src/world/world_cell.rs | 52 ++++++++++++-------------
 1 file changed, 24 insertions(+), 28 deletions(-)

diff --git a/crates/bevy_ecs/src/world/world_cell.rs b/crates/bevy_ecs/src/world/world_cell.rs
index 39770890e2..5f0fb5aaa5 100644
--- a/crates/bevy_ecs/src/world/world_cell.rs
+++ b/crates/bevy_ecs/src/world/world_cell.rs
@@ -88,21 +88,22 @@ pub struct WorldBorrow<'w, T> {
 }
 
 impl<'w, T> WorldBorrow<'w, T> {
-    fn new(
-        value: &'w T,
+    fn try_new(
+        value: impl FnOnce() -> Option<&'w T>,
         archetype_component_id: ArchetypeComponentId,
         access: Rc<RefCell<ArchetypeComponentAccess>>,
-    ) -> Self {
+    ) -> Option<Self> {
         assert!(
             access.borrow_mut().read(archetype_component_id),
             "Attempted to immutably access {}, but it is already mutably borrowed",
             std::any::type_name::<T>(),
         );
-        Self {
+        let value = value()?;
+        Some(Self {
             value,
             archetype_component_id,
             access,
-        }
+        })
     }
 }
 
@@ -129,21 +130,22 @@ pub struct WorldBorrowMut<'w, T> {
 }
 
 impl<'w, T> WorldBorrowMut<'w, T> {
-    fn new(
-        value: Mut<'w, T>,
+    fn try_new(
+        value: impl FnOnce() -> Option<Mut<'w, T>>,
         archetype_component_id: ArchetypeComponentId,
         access: Rc<RefCell<ArchetypeComponentAccess>>,
-    ) -> Self {
+    ) -> Option<Self> {
         assert!(
             access.borrow_mut().write(archetype_component_id),
             "Attempted to mutably access {}, but it is already mutably borrowed",
             std::any::type_name::<T>(),
         );
-        Self {
+        let value = value()?;
+        Some(Self {
             value,
             archetype_component_id,
             access,
-        }
+        })
     }
 }
 
@@ -189,12 +191,12 @@ impl<'w> WorldCell<'w> {
         let archetype_component_id = self
             .world
             .get_resource_archetype_component_id(component_id)?;
-        Some(WorldBorrow::new(
+        WorldBorrow::try_new(
             // SAFETY: ComponentId matches TypeId
-            unsafe { self.world.get_resource_with_id(component_id)? },
+            || unsafe { self.world.get_resource_with_id(component_id) },
             archetype_component_id,
             self.access.clone(),
-        ))
+        )
     }
 
     /// Gets a reference to the resource of the given type
@@ -222,15 +224,12 @@ impl<'w> WorldCell<'w> {
         let archetype_component_id = self
             .world
             .get_resource_archetype_component_id(component_id)?;
-        Some(WorldBorrowMut::new(
+        WorldBorrowMut::try_new(
             // SAFETY: ComponentId matches TypeId and access is checked by WorldBorrowMut
-            unsafe {
-                self.world
-                    .get_resource_unchecked_mut_with_id(component_id)?
-            },
+            || unsafe { self.world.get_resource_unchecked_mut_with_id(component_id) },
             archetype_component_id,
             self.access.clone(),
-        ))
+        )
     }
 
     /// Gets a mutable reference to the resource of the given type
@@ -258,12 +257,12 @@ impl<'w> WorldCell<'w> {
         let archetype_component_id = self
             .world
             .get_resource_archetype_component_id(component_id)?;
-        Some(WorldBorrow::new(
+        WorldBorrow::try_new(
             // SAFETY: ComponentId matches TypeId
-            unsafe { self.world.get_non_send_with_id(component_id)? },
+            || unsafe { self.world.get_non_send_with_id(component_id) },
             archetype_component_id,
             self.access.clone(),
-        ))
+        )
     }
 
     /// Gets an immutable reference to the non-send resource of the given type, if it exists.
@@ -291,15 +290,12 @@ impl<'w> WorldCell<'w> {
         let archetype_component_id = self
             .world
             .get_resource_archetype_component_id(component_id)?;
-        Some(WorldBorrowMut::new(
+        WorldBorrowMut::try_new(
             // SAFETY: ComponentId matches TypeId and access is checked by WorldBorrowMut
-            unsafe {
-                self.world
-                    .get_non_send_unchecked_mut_with_id(component_id)?
-            },
+            || unsafe { self.world.get_non_send_unchecked_mut_with_id(component_id) },
             archetype_component_id,
             self.access.clone(),
-        ))
+        )
     }
 
     /// Gets a mutable reference to the non-send resource of the given type, if it exists.