[Merged by Bors] - Speed up CommandQueue by storing commands more densely

crates/bevy_ecs/src/system/commands/command_queue.rs

@@ -1,24 +1,29 @@
-use std::mem::{ManuallyDrop, MaybeUninit};
+use std::mem::{self, MaybeUninit};
 
 use super::Command;
 use crate::world::World;
 
 struct CommandMeta {
-    offset: usize,
-    func: unsafe fn(value: *mut MaybeUninit<u8>, world: &mut World),
+    /// SAFETY: The `value` must point to a value of type `T: Command`,
+    /// where `T` is some specific type that was used to produce the function pointer `func`.
+    /// Returns the size of `T` in bytes.
+    func: unsafe fn(value: *mut MaybeUninit<u8>, world: &mut World) -> usize,
 }
 
 /// A queue of [`Command`]s
 //
-// NOTE: [`CommandQueue`] is implemented via a `Vec<MaybeUninit<u8>>` over a `Vec<Box<dyn Command>>`
+// NOTE: [`CommandQueue`] is implemented via a `Vec<MaybeUninit<u8>>` instead of a `Vec<Box<dyn Command>>`
 // as an optimization. Since commands are used frequently in systems as a way to spawn
 // entities/components/resources, and it's not currently possible to parallelize these
 // due to mutable [`World`] access, maximizing performance for [`CommandQueue`] is
 // preferred to simplicity of implementation.
 #[derive(Default)]
 pub struct CommandQueue {
+    // This contiguously stores a set of alternating objects:
+    // A `CommandMeta`, followed by a sequence of bytes with length specified by the metadata.
+    // These bytes hold the data for a type-erased `Command`, and must be passed to
+    // the corresponding `CommandMeta.func` to be used.
     bytes: Vec<MaybeUninit<u8>>,
-    metas: Vec<CommandMeta>,
 }
 
 // SAFETY: All commands [`Command`] implement [`Send`]
@@ -37,43 +42,55 @@ impl CommandQueue {
         /// SAFETY: This function is only every called when the `command` bytes is the associated
         /// [`Commands`] `T` type. Also this only reads the data via `read_unaligned` so unaligned
         /// accesses are safe.
-        unsafe fn write_command<T: Command>(command: *mut MaybeUninit<u8>, world: &mut World) {
+        unsafe fn write_command_and_get_size<T: Command>(
+            command: *mut MaybeUninit<u8>,
+            world: &mut World,
+        ) -> usize {
             let command = command.cast::<T>().read_unaligned();
             command.write(world);
+            mem::size_of::<T>()
         }
 
-        let size = std::mem::size_of::<C>();
+        let meta = CommandMeta {
+            func: write_command_and_get_size::<C>,
+        };
+
         let old_len = self.bytes.len();
 
-        self.metas.push(CommandMeta {
-            offset: old_len,
-            func: write_command::<C>,
-        });
-
-        // Use `ManuallyDrop` to forget `command` right away, avoiding
-        // any use of it after the `ptr::copy_nonoverlapping`.
-        let command = ManuallyDrop::new(command);
-
-        if size > 0 {
-            self.bytes.reserve(size);
-
-            // SAFETY: The internal `bytes` vector has enough storage for the
-            // command (see the call the `reserve` above), the vector has
-            // its length set appropriately and can contain any kind of bytes.
-            // In case we're writing a ZST and the `Vec` hasn't allocated yet
-            // then `as_mut_ptr` will be a dangling (non null) pointer, and
-            // thus valid for ZST writes.
-            // Also `command` is forgotten so that  when `apply` is called
-            // later, a double `drop` does not occur.
+        // Reserve enough bytes for both the metadata and the command itself.
+        self.bytes
+            .reserve(mem::size_of::<CommandMeta>() + mem::size_of::<C>());
+
+        // Pointer to the bytes at the end of the buffer.
+        // SAFETY: We know it is within bounds of the allocation, due to the call to `.reserve()`.
+        let ptr = unsafe { self.bytes.as_mut_ptr().add(old_len) };
+
+        // SAFETY: The end of the buffer has enough space for the metadata due to the `.reserve()` call,
+        // so we can cast it to a pointer and perform an unaligned write in order to fill the buffer.
+        // Since the buffer is of type `MaybeUninit<u8>`, any byte patterns are valid.
+        unsafe {
+            ptr.cast::<CommandMeta>().write_unaligned(meta);
+        }
+
+        if mem::size_of::<C>() > 0 {
+            // SAFETY: There is enough space after the metadata to store the command,
+            // due to the `.reserve()` call above.
+            // We will write to the buffer via an unaligned pointer write.
+            // Since the underlying buffer is of type `MaybeUninit<u8>`, any byte patterns are valid.
             unsafe {
-                std::ptr::copy_nonoverlapping(
-                    &*command as *const C as *const MaybeUninit<u8>,
-                    self.bytes.as_mut_ptr().add(old_len),
-                    size,
-                );
-                self.bytes.set_len(old_len + size);
+                ptr.add(mem::size_of::<CommandMeta>())
+                    .cast::<C>()
+                    .write_unaligned(command);
             }
         }
+
+        // Extend the length of the buffer to include the data we just wrote.
+        // SAFETY: The new length is guaranteed to fit in the vector's capacity,
+        // due to the call to `.reserve()` above.
+        unsafe {
+            self.bytes
+                .set_len(mem::size_of::<CommandMeta>() + mem::size_of::<C>() + old_len);
+        }
     }
 
     /// Execute the queued [`Command`]s in the world.
@@ -83,18 +100,36 @@ impl CommandQueue {
         // flush the previously queued entities
         world.flush();
 
+        // Cursor that will iterate over the entries in the buffer.
+        // It will alternate between values of type `CommandMeta` and a values of unknown types.
+        let mut cursor = self.bytes.as_mut_ptr();
+
+        // The address of the end of the buffer.
+        let end_addr = cursor as usize + self.bytes.len();
+
         // SAFETY: In the iteration below, `meta.func` will safely consume and drop each pushed command.
         // This operation is so that we can reuse the bytes `Vec<u8>`'s internal storage and prevent
         // unnecessary allocations.
         unsafe { self.bytes.set_len(0) };
 
-        for meta in self.metas.drain(..) {
-            // SAFETY: The implementation of `write_command` is safe for the according Command type.
-            // It's ok to read from `bytes.as_mut_ptr()` because we just wrote to it in `push`.
-            // The bytes are safely cast to their original type, safely read, and then dropped.
-            unsafe {
-                (meta.func)(self.bytes.as_mut_ptr().add(meta.offset), world);
-            }
+        while (cursor as usize) < end_addr {
+            // SAFETY: The bytes at `offset` are known to represent a value of type `CommandMeta`,
+            // since the buffer alternates between storing `CommandMeta` and unknown bytes.
+            // Its value will have been fully initialized during any calls to `push`.
+            let meta = unsafe { cursor.cast::<CommandMeta>().read_unaligned() };
+            // Advance to the bytes just after `meta`, which represent a type-erased command.
+            // SAFETY: For most types of `Command`, the pointer immediately following the metadata
+            // is guaranteed to be in bounds.
+            // The pointer might be out of bounds if the command is zero-sized,
+            // but it is okay to have a dangling pointer to a ZST.
+            cursor = unsafe { cursor.add(mem::size_of::<CommandMeta>()) };
+            // SAFETY: The type currently under the cursor must be the same type erased by `meta.func`.
+            // We know that they are the same type, since they were stored next to each other by `.push()`.
+            let size = unsafe { (meta.func)(cursor, world) };
+            // Advance the cursor past the command.
+            // SAFETY: At this point, it will either point to the next `CommandMeta`,
+            // or the cursor will be out of bounds and the loop will end.
+            cursor = unsafe { cursor.add(size) };
         }
     }
 }
@@ -203,7 +238,6 @@ mod test {
         // even though the first command panicking.
         // the `bytes`/`metas` vectors were cleared.
         assert_eq!(queue.bytes.len(), 0);
-        assert_eq!(queue.metas.len(), 0);
 
         // Even though the first command panicked, it's still ok to push
         // more commands.