From e9b3aeb38f3d41b89b58a2f6b739f7aea9a51e5b Mon Sep 17 00:00:00 2001
From: Robert Swain <robert.swain@gmail.com>
Date: Sat, 2 Sep 2023 21:16:44 +0200
Subject: [PATCH] Enhance bevymark (#9674)

# Objective

- In preparation for an initial 2D/3D mesh batching/instancing PR,
enhance `bevymark` to support some different test modes that enable
comparison and optimisation of performance

## Solution

- 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.

---------

Co-authored-by: IceSentry <IceSentry@users.noreply.github.com>
---
 examples/stress_tests/bevymark.rs   | 470 ++++++++++++++++++++++------
 examples/stress_tests/many_cubes.rs |   4 +-
 2 files changed, 382 insertions(+), 92 deletions(-)

diff --git a/examples/stress_tests/bevymark.rs b/examples/stress_tests/bevymark.rs
index d9cbe76a89315..4e2670ecabd9d 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(
@@ -113,51 +218,77 @@ fn setup(mut commands: Commands, asset_server: Res<AssetServer>) {
     };
 
     commands.spawn(Camera2dBundle::default());
-    commands.spawn((
-        TextBundle::from_sections([
-            text_section(Color::GREEN, "Bird Count"),
-            text_section(Color::CYAN, ""),
-            text_section(Color::GREEN, "\nFPS (raw): "),
-            text_section(Color::CYAN, ""),
-            text_section(Color::GREEN, "\nFPS (SMA): "),
-            text_section(Color::CYAN, ""),
-            text_section(Color::GREEN, "\nFPS (EMA): "),
-            text_section(Color::CYAN, ""),
-        ])
-        .with_style(Style {
-            position_type: PositionType::Absolute,
-            top: Val::Px(5.0),
-            left: Val::Px(5.0),
+    commands
+        .spawn(NodeBundle {
+            style: Style {
+                position_type: PositionType::Absolute,
+                padding: UiRect::all(Val::Px(5.0)),
+                ..default()
+            },
+            z_index: ZIndex::Global(i32::MAX),
+            background_color: Color::BLACK.with_a(0.75).into(),
             ..default()
-        }),
-        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),
-    });
+        })
+        .with_children(|c| {
+            c.spawn((
+                TextBundle::from_sections([
+                    text_section(Color::GREEN, "Bird Count: "),
+                    text_section(Color::CYAN, ""),
+                    text_section(Color::GREEN, "\nFPS (raw): "),
+                    text_section(Color::CYAN, ""),
+                    text_section(Color::GREEN, "\nFPS (SMA): "),
+                    text_section(Color::CYAN, ""),
+                    text_section(Color::GREEN, "\nFPS (EMA): "),
+                    text_section(Color::CYAN, ""),
+                ]),
+                StatsText,
+            ));
+        });
+
+    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 +296,192 @@ 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::from_translation(translation).with_scale(Vec3::splat(BIRD_SCALE)),
+        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 +508,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,