Merge pull request #1390 from godot-rust/qol/bench-setup

`#[bench(manual)]` for more control, including setup

Author: Bromeon

godot-macros/src/bench.rs

@@ -19,47 +19,79 @@ pub fn attribute_bench(input_decl: venial::Item) -> ParseResult<TokenStream> {
         _ => return bail!(&input_decl, "#[bench] can only be applied to functions"),
     };
 
-    // Note: allow attributes for things like #[rustfmt] or #[clippy]
-    if func.generic_params.is_some() || !func.params.is_empty() || func.where_clause.is_some() {
+    // Disallow generics, but allow attributes for things like #[rustfmt] or #[clippy].
+    if func.generic_params.is_some() || func.where_clause.is_some() {
         return bad_signature(&func);
     }
 
-    // Ignore -> (), as no one does that by accident.
-    // We need `ret` to make sure the type is correct and to avoid unused imports (by IDEs).
-    let Some(ret) = func.return_ty else {
+    let mut attr = KvParser::parse_required(&func.attributes, "bench", &func.name)?;
+    let manual = attr.handle_alone("manual")?;
+    let repetitions = attr.handle_usize("repeat")?;
+    attr.finish()?;
+
+    // Validate attribute combinations.
+    if manual && repetitions.is_some() {
         return bail!(
             func,
-            "#[bench] function must return a value from its computation, to prevent optimizing the operation away"
+            "#[bench(manual)] cannot be combined with `repeat` -- pass repetitions to bench_measure() instead"
         );
-    };
+    }
 
-    let mut attr = KvParser::parse_required(&func.attributes, "bench", &func.name)?;
-    let repetitions = attr.handle_usize("repeat")?.unwrap_or(DEFAULT_REPETITIONS);
-    attr.finish()?;
+    let repetitions = repetitions.unwrap_or(DEFAULT_REPETITIONS);
+
+    // Validate parameter count.
+    if !func.params.is_empty() {
+        return bad_signature(&func);
+    }
 
     let bench_name = &func.name;
     let bench_name_str = func.name.to_string();
 
     let body = &func.body;
 
     // Filter out #[bench] itself, but preserve other attributes like #[allow], #[expect], etc.
-    let other_attributes = retain_attributes_except(&func.attributes, "bench");
+    let other_attributes: Vec<_> = retain_attributes_except(&func.attributes, "bench").collect();
 
-    Ok(quote! {
-        #(#other_attributes)*
-        pub fn #bench_name() {
-            for _ in 0..#repetitions {
-                let __ret: #ret = #body;
-                crate::common::bench_used(__ret);
+    let generated_fn = if manual {
+        // Manual mode: user calls bench_measure() directly.
+        let ret = func.return_ty;
+        quote! {
+            #(#other_attributes)*
+            // Don't return crate::framework::BenchResult here, to keep user imports working naturally.
+            pub fn #bench_name() -> #ret {
+                #body
+            }
+        }
+    } else {
+        // Automatic mode: framework controls timing.
+        // Ignore `-> ()`, as no one does that by accident.
+        // We need `ret` to make sure the type is correct and to avoid unused imports (by IDEs).
+        let Some(ret) = func.return_ty else {
+            return bail!(
+                func,
+                "#[bench] function must return a value from its computation, to prevent optimizing the operation away"
+            );
+        };
+
+        quote! {
+            #(#other_attributes)*
+            pub fn #bench_name() -> crate::framework::BenchResult {
+                crate::framework::bench_measure(#repetitions, || {
+                    let __ret: #ret = #body;
+                    __ret // passed onto bench_used() by caller.
+                })
             }
         }
+    };
+
+    Ok(quote! {
+        #generated_fn
 
         ::godot::sys::plugin_add!(crate::framework::__GODOT_BENCH; crate::framework::RustBenchmark {
             name: #bench_name_str,
             file: std::file!(),
             line: std::line!(),
             function: #bench_name,
-            repetitions: #repetitions,
         });
     })
 }
@@ -68,7 +100,12 @@ fn bad_signature(func: &venial::Function) -> Result<TokenStream, venial::Error>
     bail!(
         func,
         "#[bench] function must have one of these signatures:\
-        \n  fn {f}() {{ ... }}",
+        \n\
+        \n(1)  #[bench]\
+        \n     fn {f}() -> T {{ ... }}\
+        \n\
+        \n(2)  #[bench(manual)]\
+        \n     fn {f}() -> BenchResult {{ ... /* call to bench_measure() */ }}",
         f = func.name,
     )
 }

itest/rust/src/benchmarks/mod.rs

@@ -16,7 +16,7 @@ use godot::obj::{Gd, InstanceId, NewAlloc, NewGd, Singleton};
 use godot::prelude::{varray, Callable, RustCallable, Variant};
 use godot::register::GodotClass;
 
-use crate::framework::bench;
+use crate::framework::{bench, bench_measure, BenchResult};
 
 mod color;
 
@@ -114,16 +114,18 @@ fn packed_array_from_iter_unknown_size() -> PackedInt32Array {
     }))
 }
 
-#[bench(repeat = 25)]
-fn call_callv_rust_fn() -> Variant {
+#[bench(manual)]
+fn call_callv_rust_fn() -> BenchResult {
     let callable = Callable::from_fn("RustFunction", |_| ());
-    callable.callv(&varray![])
+
+    bench_measure(25, || callable.callv(&varray![]))
 }
 
-#[bench(repeat = 25)]
-fn call_callv_custom() -> Variant {
+#[bench(manual)]
+fn call_callv_custom() -> BenchResult {
     let callable = Callable::from_custom(MyRustCallable {});
-    callable.callv(&varray![])
+
+    bench_measure(25, || callable.callv(&varray![]))
 }
 
 // ----------------------------------------------------------------------------------------------------------------------------------------------

itest/rust/src/common.rs

@@ -21,13 +21,3 @@ where
 
     assert_eq!(value, back);
 }
-
-/// Signal to the compiler that a value is used (to avoid optimization).
-pub fn bench_used<T: Sized>(value: T) {
-    // The following check would be used to prevent `()` arguments, ensuring that a value from the bench is actually going into the blackbox.
-    // However, we run into this issue, despite no array being used: https://github.com/rust-lang/rust/issues/43408.
-    //   error[E0401]: can't use generic parameters from outer function
-    // sys::static_assert!(std::mem::size_of::<T>() != 0, "returned unit value in benchmark; make sure to use a real value");
-
-    std::hint::black_box(value);
-}

itest/rust/src/framework/bencher.rs

@@ -17,39 +17,59 @@
 // Instead, we focus on min (fastest run) and median -- even median may vary quite a bit between runs; but it gives an idea of the distribution.
 // See also https://easyperf.net/blog/2019/12/30/Comparing-performance-measurements#average-median-minimum.
 
+use std::any::TypeId;
 use std::time::{Duration, Instant};
 
 const WARMUP_RUNS: usize = 200;
 const TEST_RUNS: usize = 501; // uneven, so median need not be interpolated.
 const METRIC_COUNT: usize = 2;
 
-pub struct BenchResult {
+pub type BenchResult = Result<BenchMeasurement, String>;
+
+pub struct BenchMeasurement {
     pub stats: [Duration; METRIC_COUNT],
 }
 
 pub fn metrics() -> [&'static str; METRIC_COUNT] {
     ["min", "median"]
 }
 
-pub fn run_benchmark(code: fn(), inner_repetitions: usize) -> BenchResult {
+/// Measures the timing of the passed closure (repeated `inner_repetitions` times).
+///
+/// Used by both `#[bench]` automatic mode (generated by macro) and `#[bench(manual)]` (called explicitly).
+///
+/// Returns `Err(String)` if there is something wrong with the benchmark setup.
+pub fn bench_measure<R: 'static>(inner_repetitions: usize, work: impl Fn() -> R) -> BenchResult {
+    // Runtime check: ensure the closure doesn't return ().
+    if TypeId::of::<R>() == TypeId::of::<()>() {
+        return Err("bench_measure(): closure must return non-unit type to prevent the computation from being optimized away".to_string());
+    }
+
+    // Warmup phase.
     for _ in 0..WARMUP_RUNS {
-        code();
+        for _ in 0..inner_repetitions {
+            let _ = work();
+        }
     }
 
+    // Measurement phase.
     let mut times = Vec::with_capacity(TEST_RUNS);
     for _ in 0..TEST_RUNS {
         let start = Instant::now();
-        code();
+        for _ in 0..inner_repetitions {
+            let fruits_of_labor = work();
+            bench_used(fruits_of_labor);
+        }
         let duration = start.elapsed();
 
         times.push(duration / inner_repetitions as u32);
     }
     times.sort();
 
-    calculate_stats(times)
+    Ok(calculate_stats(times))
 }
 
-fn calculate_stats(times: Vec<Duration>) -> BenchResult {
+fn calculate_stats(times: Vec<Duration>) -> BenchMeasurement {
     // See top of file for rationale.
 
     /*let mean = {
@@ -72,7 +92,17 @@ fn calculate_stats(times: Vec<Duration>) -> BenchResult {
     let min = times[0];
     let median = times[TEST_RUNS / 2];
 
-    BenchResult {
+    BenchMeasurement {
         stats: [min, median],
     }
 }
+
+/// Signal to the compiler that a value is used (to avoid optimization).
+fn bench_used<T: Sized>(value: T) {
+    // The following check would be used to prevent `()` arguments, ensuring that a value from the bench is actually going into the blackbox.
+    // However, we run into this issue, despite no array being used: https://github.com/rust-lang/rust/issues/43408.
+    //   error[E0401]: can't use generic parameters from outer function
+    // sys::static_assert!(std::mem::size_of::<T>() != 0, "returned unit value in benchmark; make sure to use a real value");
+
+    std::hint::black_box(value);
+}

itest/rust/src/framework/mod.rs

@@ -158,8 +158,7 @@ pub struct RustBenchmark {
     pub file: &'static str,
     #[allow(dead_code)]
     pub line: u32,
-    pub function: fn(),
-    pub repetitions: usize,
+    pub function: fn() -> BenchResult,
 }
 
 pub fn passes_filter(filters: &[String], test_name: &str) -> bool {

itest/rust/src/framework/runner.rs

@@ -376,7 +376,7 @@ impl IntegrationTests {
             print_bench_pre(bench.name, bench.file, last_file.as_deref());
             last_file = Some(bench.file.to_string());
 
-            let result = bencher::run_benchmark(bench.function, bench.repetitions);
+            let result = (bench.function)();
             print_bench_post(result);
         }
     }
@@ -546,9 +546,17 @@ fn print_bench_pre(benchmark: &str, bench_file: &str, last_file: Option<&str>) {
 }
 
 fn print_bench_post(result: BenchResult) {
-    for stat in result.stats.iter() {
-        print!(" {:>10.3}μs", stat.as_nanos() as f64 / 1000.0);
+    match result {
+        Ok(measured) => {
+            for stat in measured.stats.iter() {
+                print!(" {:>10.3}μs", stat.as_nanos() as f64 / 1000.0);
+            }
+        }
+        Err(msg) => {
+            print!("   {FMT_RED}ERROR: {msg}{FMT_END}");
+        }
     }
+
     println!();
 }