Remove duplicate lookups from Resource initialization (#7174)

# Objective

* `World::init_resource` and `World::get_resource_or_insert_with` are implemented naively, and as such they perform duplicate `TypeId -> ComponentId` lookups.
* `World::get_resource_or_insert_with` contains an additional duplicate `ComponentId -> ResourceData` lookup.
    * This function also contains an unnecessary panic branch, which we rely on the optimizer to be able to remove.

## Solution

Implement the functions using engine-internal code, instead of combining high-level functions. This allows computed variables to persist across different branches, instead of being recomputed.

crates/bevy_ecs/src/storage/resource.rs

@@ -1,4 +1,5 @@
 use crate::archetype::ArchetypeComponentId;
+use crate::change_detection::{MutUntyped, TicksMut};
 use crate::component::{ComponentId, ComponentTicks, Components, TickCells};
 use crate::storage::{Column, SparseSet, TableRow};
 use bevy_ptr::{OwningPtr, Ptr, UnsafeCellDeref};
@@ -99,6 +100,20 @@ impl<const SEND: bool> ResourceData<SEND> {
         })
     }
 
+    pub(crate) fn get_mut(
+        &mut self,
+        last_change_tick: u32,
+        change_tick: u32,
+    ) -> Option<MutUntyped<'_>> {
+        let (ptr, ticks) = self.get_with_ticks()?;
+        Some(MutUntyped {
+            // SAFETY: We have exclusive access to the underlying storage.
+            value: unsafe { ptr.assert_unique() },
+            // SAFETY: We have exclusive access to the underlying storage.
+            ticks: unsafe { TicksMut::from_tick_cells(ticks, last_change_tick, change_tick) },
+        })
+    }
+
     /// Inserts a value into the resource. If a value is already present
     /// it will be replaced.
     ///

crates/bevy_ecs/src/world/mod.rs

@@ -17,7 +17,7 @@ use crate::{
     entity::{AllocAtWithoutReplacement, Entities, Entity, EntityLocation},
     event::{Event, Events},
     ptr::UnsafeCellDeref,
-    query::{QueryState, ReadOnlyWorldQuery, WorldQuery},
+    query::{DebugCheckedUnwrap, QueryState, ReadOnlyWorldQuery, WorldQuery},
     storage::{ResourceData, SparseSet, Storages},
     system::Resource,
 };
@@ -767,9 +767,20 @@ impl World {
     /// and those default values will be here instead.
     #[inline]
     pub fn init_resource<R: Resource + FromWorld>(&mut self) {
-        if !self.contains_resource::<R>() {
+        let component_id = self.components.init_resource::<R>();
-            let resource = R::from_world(self);
+        if self
-            self.insert_resource(resource);
+            .storages
+            .resources
+            .get(component_id)
+            .map_or(true, |data| !data.is_present())
+        {
+            let value = R::from_world(self);
+            OwningPtr::make(value, |ptr| {
+                // SAFETY: component_id was just initialized and corresponds to resource of type R.
+                unsafe {
+                    self.insert_resource_by_id(component_id, ptr);
+                }
+            });
         }
     }
 
@@ -782,7 +793,7 @@ impl World {
     pub fn insert_resource<R: Resource>(&mut self, value: R) {
         let component_id = self.components.init_resource::<R>();
         OwningPtr::make(value, |ptr| {
-            // SAFETY: component_id was just initialized and corresponds to resource of type R
+            // SAFETY: component_id was just initialized and corresponds to resource of type R.
             unsafe {
                 self.insert_resource_by_id(component_id, ptr);
             }
@@ -802,9 +813,20 @@ impl World {
     /// Panics if called from a thread other than the main thread.
     #[inline]
     pub fn init_non_send_resource<R: 'static + FromWorld>(&mut self) {
-        if !self.contains_non_send::<R>() {
+        let component_id = self.components.init_non_send::<R>();
-            let resource = R::from_world(self);
+        if self
-            self.insert_non_send_resource(resource);
+            .storages
+            .resources
+            .get(component_id)
+            .map_or(true, |data| !data.is_present())
+        {
+            let value = R::from_world(self);
+            OwningPtr::make(value, |ptr| {
+                // SAFETY: component_id was just initialized and corresponds to resource of type R.
+                unsafe {
+                    self.insert_non_send_by_id(component_id, ptr);
+                }
+            });
         }
     }
 
@@ -821,7 +843,7 @@ impl World {
     pub fn insert_non_send_resource<R: 'static>(&mut self, value: R) {
         let component_id = self.components.init_non_send::<R>();
         OwningPtr::make(value, |ptr| {
-            // SAFETY: component_id was just initialized and corresponds to resource of type R
+            // SAFETY: component_id was just initialized and corresponds to resource of type R.
             unsafe {
                 self.insert_non_send_by_id(component_id, ptr);
             }
@@ -959,7 +981,6 @@ impl World {
         unsafe { self.get_resource_unchecked_mut() }
     }
 
-    // PERF: optimize this to avoid redundant lookups
     /// Gets a mutable reference to the resource of type `T` if it exists,
     /// otherwise inserts the resource using the result of calling `func`.
     #[inline]
@@ -967,10 +988,27 @@ impl World {
         &mut self,
         func: impl FnOnce() -> R,
     ) -> Mut<'_, R> {
-        if !self.contains_resource::<R>() {
+        let change_tick = self.change_tick();
-            self.insert_resource(func());
+        let last_change_tick = self.last_change_tick();
+
+        let component_id = self.components.init_resource::<R>();
+        let data = self.initialize_resource_internal(component_id);
+        if !data.is_present() {
+            OwningPtr::make(func(), |ptr| {
+                // SAFETY: component_id was just initialized and corresponds to resource of type R.
+                unsafe {
+                    data.insert(ptr, change_tick);
+                }
+            });
         }
-        self.resource_mut()
+
+        // SAFETY: The resource must be present, as we would have inserted it if it was empty.
+        let data = unsafe {
+            data.get_mut(last_change_tick, change_tick)
+                .debug_checked_unwrap()
+        };
+        // SAFETY: The underlying type of the resource is `R`.
+        unsafe { data.with_type::<R>() }
     }
 
     /// Gets a mutable reference to the resource of the given type, if it exists