From a06512362ce29f93be2a9ccce54950b6e8441b1e Mon Sep 17 00:00:00 2001
From: TheRawMeatball <therawmeatball@gmail.com>
Date: Mon, 24 Oct 2022 19:24:49 +0000
Subject: [PATCH] Optimize rendering slow-down at high entity counts (#5509)

# Objective

- Improve #3953

## Solution

- The very specific circumstances under which the render world is reset meant that the flush_as_invalid function could be replaced with one that had a noop as its init method.
- This removes a double-writing issue leading to greatly increased performance.

Running the reproduction code in the linked issue, this change nearly doubles the framerate.

Co-authored-by: Carter Anderson <mcanders1@gmail.com>
---
 crates/bevy_ecs/src/archetype.rs  |  4 ++++
 crates/bevy_ecs/src/entity/mod.rs | 21 +++++++++++++++++++++
 crates/bevy_render/src/lib.rs     | 23 ++++++++++++++---------
 3 files changed, 39 insertions(+), 9 deletions(-)

diff --git a/crates/bevy_ecs/src/archetype.rs b/crates/bevy_ecs/src/archetype.rs
index 9b9377e7ce2b9..b52d4baad44a3 100644
--- a/crates/bevy_ecs/src/archetype.rs
+++ b/crates/bevy_ecs/src/archetype.rs
@@ -14,10 +14,14 @@ use std::{
 };
 
 #[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
+#[repr(transparent)]
 pub struct ArchetypeId(usize);
 
 impl ArchetypeId {
     pub const EMPTY: ArchetypeId = ArchetypeId(0);
+    /// # Safety:
+    ///
+    /// This must always have an all-1s bit pattern to ensure soundness in fast entity id space allocation.
     pub const INVALID: ArchetypeId = ArchetypeId(usize::MAX);
 
     #[inline]
diff --git a/crates/bevy_ecs/src/entity/mod.rs b/crates/bevy_ecs/src/entity/mod.rs
index 1c56e1b40e15b..14071b50c11a4 100644
--- a/crates/bevy_ecs/src/entity/mod.rs
+++ b/crates/bevy_ecs/src/entity/mod.rs
@@ -563,6 +563,9 @@ impl Entities {
     /// Flush _must_ set the entity location to the correct [`ArchetypeId`] for the given [`Entity`]
     /// each time init is called. This _can_ be [`ArchetypeId::INVALID`], provided the [`Entity`]
     /// has not been assigned to an [`Archetype`][crate::archetype::Archetype].
+    ///
+    /// Note: freshly-allocated entities (ones which don't come from the pending list) are guaranteed
+    /// to be initialized with the invalid archetype.
     pub unsafe fn flush(&mut self, mut init: impl FnMut(Entity, &mut EntityLocation)) {
         let free_cursor = self.free_cursor.get_mut();
         let current_free_cursor = *free_cursor;
@@ -613,6 +616,20 @@ impl Entities {
         }
     }
 
+    /// # Safety
+    ///
+    /// This function is safe if and only if the world this Entities is on has no entities.
+    pub unsafe fn flush_and_reserve_invalid_assuming_no_entities(&mut self, count: usize) {
+        let free_cursor = self.free_cursor.get_mut();
+        *free_cursor = 0;
+        self.meta.reserve(count);
+        // the EntityMeta struct only contains integers, and it is valid to have all bytes set to u8::MAX
+        self.meta.as_mut_ptr().write_bytes(u8::MAX, count);
+        self.meta.set_len(count);
+
+        self.len = count as u32;
+    }
+
     /// Accessor for getting the length of the vec in `self.meta`
     #[inline]
     pub fn meta_len(&self) -> usize {
@@ -630,7 +647,10 @@ impl Entities {
     }
 }
 
+// Safety:
+// This type must not contain any pointers at any level, and be safe to fully fill with u8::MAX.
 #[derive(Copy, Clone, Debug)]
+#[repr(C)]
 pub struct EntityMeta {
     pub generation: u32,
     pub location: EntityLocation,
@@ -648,6 +668,7 @@ impl EntityMeta {
 
 /// A location of an entity in an archetype.
 #[derive(Copy, Clone, Debug)]
+#[repr(C)]
 pub struct EntityLocation {
     /// The archetype index
     pub archetype_id: ArchetypeId,
diff --git a/crates/bevy_render/src/lib.rs b/crates/bevy_render/src/lib.rs
index 9919ac9da87f3..d9b56cd9f2f1c 100644
--- a/crates/bevy_render/src/lib.rs
+++ b/crates/bevy_render/src/lib.rs
@@ -244,15 +244,20 @@ impl Plugin for RenderPlugin {
                     // reserve all existing app entities for use in render_app
                     // they can only be spawned using `get_or_spawn()`
                     let meta_len = app_world.entities().meta_len();
-                    render_app
-                        .world
-                        .entities()
-                        .reserve_entities(meta_len as u32);
-
-                    // flushing as "invalid" ensures that app world entities aren't added as "empty archetype" entities by default
-                    // these entities cannot be accessed without spawning directly onto them
-                    // this _only_ works as expected because clear_entities() is called at the end of every frame.
-                    unsafe { render_app.world.entities_mut() }.flush_as_invalid();
+
+                    assert_eq!(
+                        render_app.world.entities().len(),
+                        0,
+                        "An entity was spawned after the entity list was cleared last frame and before the extract stage began. This is not supported",
+                    );
+
+                    // This is safe given the clear_entities call in the past frame and the assert above
+                    unsafe {
+                        render_app
+                            .world
+                            .entities_mut()
+                            .flush_and_reserve_invalid_assuming_no_entities(meta_len);
+                    }
                 }
 
                 {