From c7c7acfb9e04a5fe57243c469e4e20fe7289ecf4 Mon Sep 17 00:00:00 2001
From: Robert Swain <robert.swain@gmail.com>
Date: Mon, 28 Aug 2023 23:25:48 +0200
Subject: [PATCH] Enhance bevymark

- Use `argh` for command line interface options
- Use seeded StdRng for reproducible random number generation
- Add a mode for testing 2D meshes that includes an option to uniquely vary the
  data of each material by setting a random flat colour on the ColorMaterial.
- Add a way of specifying the number of different textures to use for sprites
  or meshes. These are generated at the same resolution as the Bevy bird icon,
  but are just random flat colours for testing.
- Add a benchmark mode that spawns all entities during setup, and animates the
  entities using a fixed delta time for reproducible animation. The
  initially-spawned entities are still spawned in waves and animated as they
  would have been had they spawned at intervals.
---
 examples/stress_tests/bevymark.rs   | 429 +++++++++++++++++++++++-----
 examples/stress_tests/many_cubes.rs |   4 +-
 2 files changed, 360 insertions(+), 73 deletions(-)

diff --git a/examples/stress_tests/bevymark.rs b/examples/stress_tests/bevymark.rs
index d9cbe76a89315..bbf01661762a4 100644
--- a/examples/stress_tests/bevymark.rs
+++ b/examples/stress_tests/bevymark.rs
@@ -2,18 +2,24 @@
 //!
 //! Usage: spawn more entities by clicking on the screen.
 
+use std::str::FromStr;
+
+use argh::FromArgs;
 use bevy::{
     diagnostic::{DiagnosticsStore, FrameTimeDiagnosticsPlugin, LogDiagnosticsPlugin},
     prelude::*,
+    render::render_resource::{Extent3d, TextureDimension, TextureFormat},
+    sprite::{MaterialMesh2dBundle, Mesh2dHandle},
     window::{PresentMode, WindowResolution},
 };
-use rand::{rngs::StdRng, thread_rng, Rng, SeedableRng};
+use rand::{rngs::StdRng, seq::SliceRandom, Rng, SeedableRng};
 
 const BIRDS_PER_SECOND: u32 = 10000;
 const GRAVITY: f32 = -9.8 * 100.0;
 const MAX_VELOCITY: f32 = 750.;
 const BIRD_SCALE: f32 = 0.15;
-const HALF_BIRD_SIZE: f32 = 256. * BIRD_SCALE * 0.5;
+const BIRD_TEXTURE_SIZE: usize = 256;
+const HALF_BIRD_SIZE: f32 = BIRD_TEXTURE_SIZE as f32 * BIRD_SCALE * 0.5;
 
 #[derive(Resource)]
 struct BevyCounter {
@@ -26,7 +32,62 @@ struct Bird {
     velocity: Vec3,
 }
 
+#[derive(FromArgs, Resource)]
+/// `bevymark` sprite / 2D mesh stress test
+struct Args {
+    /// whether to use sprite or mesh2d
+    #[argh(option, default = "Mode::Sprite")]
+    mode: Mode,
+
+    /// whether to step animations by a fixed amount such that each frame is the same across runs.
+    /// If spawning waves, all are spawned up-front to immediately start rendering at the heaviest
+    /// load.
+    #[argh(switch)]
+    benchmark: bool,
+
+    /// how many birds to spawn per wave.
+    #[argh(option, default = "0")]
+    per_wave: usize,
+
+    /// the number of waves to spawn.
+    #[argh(option, default = "0")]
+    waves: usize,
+
+    /// whether to vary the material data in each instance.
+    #[argh(switch)]
+    vary_per_instance: bool,
+
+    /// the number of different textures from which to randomly select the material color. 0 means no textures.
+    #[argh(option, default = "1")]
+    material_texture_count: usize,
+}
+
+#[derive(Default, Clone)]
+enum Mode {
+    #[default]
+    Sprite,
+    Mesh2d,
+}
+
+impl FromStr for Mode {
+    type Err = String;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s {
+            "sprite" => Ok(Self::Sprite),
+            "mesh2d" => Ok(Self::Mesh2d),
+            _ => Err(format!(
+                "Unknown mode: '{s}', valid modes: 'sprite', 'mesh2d'"
+            )),
+        }
+    }
+}
+
+const FIXED_TIMESTEP: f32 = 0.2;
+
 fn main() {
+    let args: Args = argh::from_env();
+
     App::new()
         .add_plugins((
             DefaultPlugins.set(WindowPlugin {
@@ -41,6 +102,7 @@ fn main() {
             FrameTimeDiagnosticsPlugin,
             LogDiagnosticsPlugin::default(),
         ))
+        .insert_resource(args)
         .insert_resource(BevyCounter {
             count: 0,
             color: Color::WHITE,
@@ -56,50 +118,93 @@ fn main() {
                 counter_system,
             ),
         )
-        .insert_resource(FixedTime::new_from_secs(0.2))
+        .insert_resource(FixedTime::new_from_secs(FIXED_TIMESTEP))
         .run();
 }
 
 #[derive(Resource)]
 struct BirdScheduled {
-    wave: usize,
+    waves: usize,
     per_wave: usize,
 }
 
 fn scheduled_spawner(
     mut commands: Commands,
+    args: Res<Args>,
     windows: Query<&Window>,
     mut scheduled: ResMut<BirdScheduled>,
     mut counter: ResMut<BevyCounter>,
-    bird_texture: Res<BirdTexture>,
+    bird_resources: ResMut<BirdResources>,
 ) {
     let window = windows.single();
 
-    if scheduled.wave > 0 {
+    if scheduled.waves > 0 {
+        let bird_resources = bird_resources.into_inner();
         spawn_birds(
             &mut commands,
+            args.into_inner(),
             &window.resolution,
             &mut counter,
             scheduled.per_wave,
-            bird_texture.clone_weak(),
+            bird_resources,
+            None,
+            scheduled.waves - 1,
         );
 
-        let mut rng = thread_rng();
-        counter.color = Color::rgb_linear(rng.gen(), rng.gen(), rng.gen());
-        scheduled.wave -= 1;
+        scheduled.waves -= 1;
     }
 }
 
-#[derive(Resource, Deref)]
-struct BirdTexture(Handle<Image>);
+#[derive(Resource)]
+struct BirdResources {
+    textures: Vec<Handle<Image>>,
+    materials: Vec<Handle<ColorMaterial>>,
+    quad: Mesh2dHandle,
+    color_rng: StdRng,
+    material_rng: StdRng,
+    velocity_rng: StdRng,
+}
 
 #[derive(Component)]
 struct StatsText;
 
-fn setup(mut commands: Commands, asset_server: Res<AssetServer>) {
+#[allow(clippy::too_many_arguments)]
+fn setup(
+    mut commands: Commands,
+    args: Res<Args>,
+    asset_server: Res<AssetServer>,
+    mut meshes: ResMut<Assets<Mesh>>,
+    material_assets: ResMut<Assets<ColorMaterial>>,
+    images: ResMut<Assets<Image>>,
+    windows: Query<&Window>,
+    counter: ResMut<BevyCounter>,
+) {
     warn!(include_str!("warning_string.txt"));
 
-    let texture = asset_server.load("branding/icon.png");
+    let args = args.into_inner();
+    let images = images.into_inner();
+
+    let mut textures = Vec::with_capacity(args.material_texture_count.max(1));
+    if matches!(args.mode, Mode::Sprite) || args.material_texture_count > 0 {
+        textures.push(asset_server.load("branding/icon.png"));
+    }
+    init_textures(&mut textures, args, images);
+
+    let material_assets = material_assets.into_inner();
+    let materials = init_materials(args, &textures, material_assets);
+
+    let mut bird_resources = BirdResources {
+        textures,
+        materials,
+        quad: meshes
+            .add(Mesh::from(shape::Quad::new(Vec2::splat(
+                BIRD_TEXTURE_SIZE as f32,
+            ))))
+            .into(),
+        color_rng: StdRng::seed_from_u64(42),
+        material_rng: StdRng::seed_from_u64(42),
+        velocity_rng: StdRng::seed_from_u64(42),
+    };
 
     let text_section = move |color, value: &str| {
         TextSection::new(
@@ -133,31 +238,50 @@ fn setup(mut commands: Commands, asset_server: Res<AssetServer>) {
         StatsText,
     ));
 
-    commands.insert_resource(BirdTexture(texture));
-    commands.insert_resource(BirdScheduled {
-        per_wave: std::env::args()
-            .nth(1)
-            .and_then(|arg| arg.parse::<usize>().ok())
-            .unwrap_or_default(),
-        wave: std::env::args()
-            .nth(2)
-            .and_then(|arg| arg.parse::<usize>().ok())
-            .unwrap_or(1),
-    });
+    let mut scheduled = BirdScheduled {
+        per_wave: args.per_wave,
+        waves: args.waves,
+    };
+
+    if args.benchmark {
+        let counter = counter.into_inner();
+        for wave in (0..scheduled.waves).rev() {
+            spawn_birds(
+                &mut commands,
+                args,
+                &windows.single().resolution,
+                counter,
+                scheduled.per_wave,
+                &mut bird_resources,
+                Some(wave),
+                wave,
+            );
+        }
+        scheduled.waves = 0;
+    }
+    commands.insert_resource(bird_resources);
+    commands.insert_resource(scheduled);
 }
 
+#[allow(clippy::too_many_arguments)]
 fn mouse_handler(
     mut commands: Commands,
+    args: Res<Args>,
     time: Res<Time>,
     mouse_button_input: Res<Input<MouseButton>>,
     windows: Query<&Window>,
-    bird_texture: Res<BirdTexture>,
+    bird_resources: ResMut<BirdResources>,
     mut counter: ResMut<BevyCounter>,
+    mut rng: Local<Option<StdRng>>,
+    mut wave: Local<usize>,
 ) {
+    if rng.is_none() {
+        *rng = Some(StdRng::seed_from_u64(42));
+    }
+    let rng = rng.as_mut().unwrap();
     let window = windows.single();
 
     if mouse_button_input.just_released(MouseButton::Left) {
-        let mut rng = thread_rng();
         counter.color = Color::rgb_linear(rng.gen(), rng.gen(), rng.gen());
     }
 
@@ -165,83 +289,196 @@ fn mouse_handler(
         let spawn_count = (BIRDS_PER_SECOND as f64 * time.delta_seconds_f64()) as usize;
         spawn_birds(
             &mut commands,
+            args.into_inner(),
             &window.resolution,
             &mut counter,
             spawn_count,
-            bird_texture.clone_weak(),
+            bird_resources.into_inner(),
+            None,
+            *wave,
         );
+        *wave += 1;
     }
 }
 
+fn bird_velocity_transform(
+    half_extents: Vec2,
+    mut translation: Vec3,
+    velocity_rng: &mut StdRng,
+    waves: Option<usize>,
+    dt: f32,
+) -> (Transform, Vec3) {
+    let mut velocity = Vec3::new(MAX_VELOCITY * (velocity_rng.gen::<f32>() - 0.5), 0., 0.);
+
+    if let Some(waves) = waves {
+        // Step the movement and handle collisions as if the wave had been spawned at fixed time intervals
+        // and with dt-spaced frames of simulation
+        for _ in 0..(waves * (FIXED_TIMESTEP / dt).round() as usize) {
+            step_movement(&mut translation, &mut velocity, dt);
+            handle_collision(half_extents, &translation, &mut velocity);
+        }
+    }
+    (
+        Transform {
+            translation,
+            scale: Vec3::splat(BIRD_SCALE),
+            ..default()
+        },
+        velocity,
+    )
+}
+
+const FIXED_DELTA_TIME: f32 = 1.0 / 60.0;
+
+#[allow(clippy::too_many_arguments)]
 fn spawn_birds(
     commands: &mut Commands,
+    args: &Args,
     primary_window_resolution: &WindowResolution,
     counter: &mut BevyCounter,
     spawn_count: usize,
-    texture: Handle<Image>,
+    bird_resources: &mut BirdResources,
+    waves_to_simulate: Option<usize>,
+    wave: usize,
 ) {
     let bird_x = (primary_window_resolution.width() / -2.) + HALF_BIRD_SIZE;
     let bird_y = (primary_window_resolution.height() / 2.) - HALF_BIRD_SIZE;
 
-    let mut rng = StdRng::from_entropy();
+    let half_extents = 0.5
+        * Vec2::new(
+            primary_window_resolution.width(),
+            primary_window_resolution.height(),
+        );
 
     let color = counter.color;
     let current_count = counter.count;
 
-    commands.spawn_batch((0..spawn_count).map(move |count| {
-        let velocity_x = rng.gen::<f32>() * MAX_VELOCITY - (MAX_VELOCITY * 0.5);
-        let bird_z = (current_count + count) as f32 * 0.00001;
-        (
-            SpriteBundle {
-                texture: texture.clone(),
-                transform: Transform {
-                    translation: Vec3::new(bird_x, bird_y, bird_z),
-                    scale: Vec3::splat(BIRD_SCALE),
-                    ..default()
-                },
-                sprite: Sprite { color, ..default() },
-                ..default()
-            },
-            Bird {
-                velocity: Vec3::new(velocity_x, 0., 0.),
-            },
-        )
-    }));
+    match args.mode {
+        Mode::Sprite => {
+            let batch = (0..spawn_count)
+                .map(|count| {
+                    let bird_z = (current_count + count) as f32 * 0.00001;
+                    let (transform, velocity) = bird_velocity_transform(
+                        half_extents,
+                        Vec3::new(bird_x, bird_y, bird_z),
+                        &mut bird_resources.velocity_rng,
+                        waves_to_simulate,
+                        FIXED_DELTA_TIME,
+                    );
+
+                    let color = if args.vary_per_instance {
+                        Color::rgb_linear(
+                            bird_resources.color_rng.gen(),
+                            bird_resources.color_rng.gen(),
+                            bird_resources.color_rng.gen(),
+                        )
+                    } else {
+                        color
+                    };
+                    (
+                        SpriteBundle {
+                            texture: bird_resources
+                                .textures
+                                .choose(&mut bird_resources.material_rng)
+                                .unwrap()
+                                .clone(),
+                            transform,
+                            sprite: Sprite { color, ..default() },
+                            ..default()
+                        },
+                        Bird { velocity },
+                    )
+                })
+                .collect::<Vec<_>>();
+            commands.spawn_batch(batch);
+        }
+        Mode::Mesh2d => {
+            let batch = (0..spawn_count)
+                .map(|count| {
+                    let bird_z = (current_count + count) as f32 * 0.00001;
+                    let (transform, velocity) = bird_velocity_transform(
+                        half_extents,
+                        Vec3::new(bird_x, bird_y, bird_z),
+                        &mut bird_resources.velocity_rng,
+                        waves_to_simulate,
+                        FIXED_DELTA_TIME,
+                    );
+
+                    let material =
+                        if args.vary_per_instance || args.material_texture_count > args.waves {
+                            bird_resources
+                                .materials
+                                .choose(&mut bird_resources.material_rng)
+                                .unwrap()
+                                .clone()
+                        } else {
+                            bird_resources.materials[wave % bird_resources.materials.len()].clone()
+                        };
+                    (
+                        MaterialMesh2dBundle {
+                            mesh: bird_resources.quad.clone(),
+                            material,
+                            transform,
+                            ..default()
+                        },
+                        Bird { velocity },
+                    )
+                })
+                .collect::<Vec<_>>();
+            commands.spawn_batch(batch);
+        }
+    }
 
     counter.count += spawn_count;
+    counter.color = Color::rgb_linear(
+        bird_resources.color_rng.gen(),
+        bird_resources.color_rng.gen(),
+        bird_resources.color_rng.gen(),
+    );
+}
+
+fn step_movement(translation: &mut Vec3, velocity: &mut Vec3, dt: f32) {
+    translation.x += velocity.x * dt;
+    translation.y += velocity.y * dt;
+    velocity.y += GRAVITY * dt;
 }
 
-fn movement_system(time: Res<Time>, mut bird_query: Query<(&mut Bird, &mut Transform)>) {
+fn movement_system(
+    args: Res<Args>,
+    time: Res<Time>,
+    mut bird_query: Query<(&mut Bird, &mut Transform)>,
+) {
+    let dt = if args.benchmark {
+        FIXED_DELTA_TIME
+    } else {
+        time.delta_seconds()
+    };
     for (mut bird, mut transform) in &mut bird_query {
-        transform.translation.x += bird.velocity.x * time.delta_seconds();
-        transform.translation.y += bird.velocity.y * time.delta_seconds();
-        bird.velocity.y += GRAVITY * time.delta_seconds();
+        step_movement(&mut transform.translation, &mut bird.velocity, dt);
     }
 }
 
+fn handle_collision(half_extents: Vec2, translation: &Vec3, velocity: &mut Vec3) {
+    if (velocity.x > 0. && translation.x + HALF_BIRD_SIZE > half_extents.x)
+        || (velocity.x <= 0. && translation.x - HALF_BIRD_SIZE < -(half_extents.x))
+    {
+        velocity.x = -velocity.x;
+    }
+    let velocity_y = velocity.y;
+    if velocity_y < 0. && translation.y - HALF_BIRD_SIZE < -half_extents.y {
+        velocity.y = -velocity_y;
+    }
+    if translation.y + HALF_BIRD_SIZE > half_extents.y && velocity_y > 0.0 {
+        velocity.y = 0.0;
+    }
+}
 fn collision_system(windows: Query<&Window>, mut bird_query: Query<(&mut Bird, &Transform)>) {
     let window = windows.single();
 
-    let half_width = window.width() * 0.5;
-    let half_height = window.height() * 0.5;
+    let half_extents = 0.5 * Vec2::new(window.width(), window.height());
 
     for (mut bird, transform) in &mut bird_query {
-        let x_vel = bird.velocity.x;
-        let y_vel = bird.velocity.y;
-        let x_pos = transform.translation.x;
-        let y_pos = transform.translation.y;
-
-        if (x_vel > 0. && x_pos + HALF_BIRD_SIZE > half_width)
-            || (x_vel <= 0. && x_pos - HALF_BIRD_SIZE < -(half_width))
-        {
-            bird.velocity.x = -x_vel;
-        }
-        if y_vel < 0. && y_pos - HALF_BIRD_SIZE < -half_height {
-            bird.velocity.y = -y_vel;
-        }
-        if y_pos + HALF_BIRD_SIZE > half_height && y_vel > 0.0 {
-            bird.velocity.y = 0.0;
-        }
+        handle_collision(half_extents, &transform.translation, &mut bird.velocity);
     }
 }
 
@@ -268,3 +505,53 @@ fn counter_system(
         }
     };
 }
+
+fn init_textures(textures: &mut Vec<Handle<Image>>, args: &Args, images: &mut Assets<Image>) {
+    let mut color_rng = StdRng::seed_from_u64(42);
+    while textures.len() < args.material_texture_count {
+        let pixel = [color_rng.gen(), color_rng.gen(), color_rng.gen(), 255];
+        textures.push(images.add(Image::new_fill(
+            Extent3d {
+                width: BIRD_TEXTURE_SIZE as u32,
+                height: BIRD_TEXTURE_SIZE as u32,
+                depth_or_array_layers: 1,
+            },
+            TextureDimension::D2,
+            &pixel,
+            TextureFormat::Rgba8UnormSrgb,
+        )));
+    }
+}
+
+fn init_materials(
+    args: &Args,
+    textures: &[Handle<Image>],
+    assets: &mut Assets<ColorMaterial>,
+) -> Vec<Handle<ColorMaterial>> {
+    let capacity = if args.vary_per_instance {
+        args.per_wave * args.waves
+    } else {
+        args.material_texture_count.max(args.waves)
+    }
+    .max(1);
+
+    let mut materials = Vec::with_capacity(capacity);
+    materials.push(assets.add(ColorMaterial {
+        color: Color::WHITE,
+        texture: textures.get(0).cloned(),
+    }));
+
+    let mut color_rng = StdRng::seed_from_u64(42);
+    let mut texture_rng = StdRng::seed_from_u64(42);
+    materials.extend(
+        std::iter::repeat_with(|| {
+            assets.add(ColorMaterial {
+                color: Color::rgb_u8(color_rng.gen(), color_rng.gen(), color_rng.gen()),
+                texture: textures.choose(&mut texture_rng).cloned(),
+            })
+        })
+        .take(capacity - materials.len()),
+    );
+
+    materials
+}
diff --git a/examples/stress_tests/many_cubes.rs b/examples/stress_tests/many_cubes.rs
index 29efbfda2c2f8..218757e471c79 100644
--- a/examples/stress_tests/many_cubes.rs
+++ b/examples/stress_tests/many_cubes.rs
@@ -33,7 +33,7 @@ struct Args {
 
     /// whether to vary the material data in each instance.
     #[argh(switch)]
-    vary_material_data: bool,
+    vary_per_instance: bool,
 
     /// the number of different textures from which to randomly select the material base color. 0 means no textures.
     #[argh(option, default = "0")]
@@ -206,7 +206,7 @@ fn init_materials(
     textures: &[Handle<Image>],
     assets: &mut Assets<StandardMaterial>,
 ) -> Vec<Handle<StandardMaterial>> {
-    let capacity = if args.vary_material_data {
+    let capacity = if args.vary_per_instance {
         match args.layout {
             Layout::Cube => (WIDTH - WIDTH / 10) * (HEIGHT - HEIGHT / 10),
             Layout::Sphere => WIDTH * HEIGHT * 4,