many_components stress test improvements

examples/stress_tests/many_components.rs

@@ -29,11 +29,14 @@ use bevy::{
 
 use rand::prelude::{Rng, SeedableRng, SliceRandom};
 use rand_chacha::ChaCha8Rng;
-use std::{alloc::Layout, num::Wrapping};
+use std::{alloc::Layout, mem::ManuallyDrop, num::Wrapping, ptr::NonNull};
 
 // A simple system that matches against several components and does some menial calculation to create
 // some non-trivial load.
 fn base_system(access_components: In<Vec<ComponentId>>, mut query: Query<FilteredEntityMut>) {
+    #[cfg(feature = "trace")]
+    let _span = bevy::utils::tracing::info_span!("base_system", components = ?access_components.0, count = query.iter().len()).entered();
+
     for mut filtered_entity in &mut query {
         // We calculate Faulhaber's formula mod 256 with n = value and p = exponent.
         // See https://en.wikipedia.org/wiki/Faulhaber%27s_formula
@@ -84,8 +87,8 @@ fn stress_test(num_entities: u32, num_components: u32, num_systems: u32) {
             world.register_component_with_descriptor(
                 #[allow(unsafe_code)]
                 // SAFETY:
-                // we don't implement a drop function
-                // u8 is Sync and Send
+                // * We don't implement a drop function
+                // * u8 is Sync and Send
                 unsafe {
                     ComponentDescriptor::new_with_layout(
                         format!("Component{}", i).to_string(),
@@ -124,23 +127,51 @@ fn stress_test(num_entities: u32, num_components: u32, num_systems: u32) {
     // spawn a bunch of entities
     for _ in 1..=num_entities {
         let num_components = rng.gen_range(1..10);
-        let components = component_ids.choose_multiple(&mut rng, num_components);
+        let components: Vec<ComponentId> = component_ids
+            .choose_multiple(&mut rng, num_components)
+            .copied()
+            .collect();
 
         let mut entity = world.spawn_empty();
-        for &component_id in components {
-            let value: u8 = rng.gen_range(0..255);
-            OwningPtr::make(value, |ptr| {
+        // We use `ManuallyDrop` here as we need to avoid dropping the u8's when `values` is dropped
+        // since ownership of the values is passed to the world in `insert_by_ids`. 
+        // But we do want to deallocate the memory when values is dropped.
+        let mut values: Vec<ManuallyDrop<u8>> = components
+            .iter()
+            .map(|_id| ManuallyDrop::new(rng.gen_range(0..255)))
+            .collect();
+        let ptrs: Vec<OwningPtr> = values
+            .iter_mut()
+            .map(|value| {
+                // ManuallyDrop is repr(transparent) so casting it to a *mut u8 is valid.
+                let ptr = value as *mut ManuallyDrop<u8> as *mut u8;
+                // SAFETY: value is non null since it was initialized from an &mut above
                 #[allow(unsafe_code)]
+                let ptr = unsafe { NonNull::new_unchecked(ptr) };
                 // SAFETY:
-                // component_id is from the same world
-                // value is u8, so ptr is a valid reference for component_id
+                // * values is valid until the end of the for block and this pointer is consumed by `insert_by_ids`.
+                // * ptr was created from a u8, so is aligned, and has correct provenance.
+                // * ptr was created from an exclusive reference and nothing else reads or writes to `values`.
+                #[allow(unsafe_code)]
                 unsafe {
-                    entity.insert_by_id(component_id, ptr);
+                    OwningPtr::new(ptr)
                 }
-            });
+            })
+            .collect();
+        // SAFETY:
+        // * component_id's are from the same world
+        // * `values` was initialized above, so references are valid
+        #[allow(unsafe_code)]
+        unsafe {
+            entity.insert_by_ids(&components, ptrs.into_iter());
         }
     }
 
+    println!(
+        "Number of Archetype-Components: {}",
+        world.archetypes().archetype_components_len()
+    );
+
     // overwrite Update schedule in the app
     app.add_schedule(schedule);
     app.add_plugins(MinimalPlugins)