feat: Add criterion to revme bench command (#2295)

nemnem7 · web-flow · commit 6f383227cad2 · 2025-03-26T00:00:11.000+01:00
* temp code preview

* add criterion to benchmarks

* added args to fn run() calls in bins/revme/

* modifed the code inside the benchmark tests

* left only .replay function inside bench_functions and added minor revisions to the rest of the code

* edited comment annotation in burntpix.rs file

* added comments and set different default value for warm_up_time
diff --git a/bins/revme/Cargo.toml b/bins/revme/Cargo.toml
@@ -27,6 +27,7 @@ context-interface.workspace = true
 handler.workspace = true
 inspector = { workspace = true, features = ["std", "serde-json"] }
 statetest-types.workspace = true
+criterion.workspace = true
 
 # alloy
 alloy-rlp = { workspace = true, features = ["arrayvec", "derive"] }
@@ -45,9 +46,6 @@ triehash.workspace = true
 walkdir.workspace = true
 k256 = { workspace = true, features = ["ecdsa"] }
 
-[dev-dependencies]
-criterion.workspace = true
-
 [[bench]]
 name = "evm"
 harness = false
diff --git a/bins/revme/benches/evm.rs b/bins/revme/benches/evm.rs
@@ -1,20 +1,11 @@
 use criterion::{criterion_group, criterion_main, Criterion};
-use revme::cmd::{
-    bench::{self, BenchName},
-    MainCmd,
-};
+use revme::cmd::bench;
 
 fn evm(c: &mut Criterion) {
-    // call analysis to init static data.
-    revme::cmd::bench::analysis::run();
-
-    for &bench_name in BenchName::ALL {
-        let cmd = MainCmd::Bench(bench::Cmd { name: bench_name });
-        c.bench_function(bench_name.as_str(), |b| {
-            b.iter(|| cmd.run().unwrap());
-        });
-    }
+    bench::analysis::run(c);
+    bench::burntpix::run(c);
+    bench::snailtracer::run(c);
+    bench::transfer::run(c);
 }
-
 criterion_group!(benches, evm);
 criterion_main!(benches);
diff --git a/bins/revme/src/cmd/bench.rs b/bins/revme/src/cmd/bench.rs
@@ -36,16 +36,41 @@ impl BenchName {
 pub struct Cmd {
     #[arg(value_enum)]
     pub name: BenchName,
+    /// Warmup represents warm up time for benchmarks ran
+    #[arg(short = 'w', long)]
+    pub warmup: Option<f64>,
+    /// Samples represents default measurement time for benchmarks ran
+    #[arg(short = 'm', long)]
+    pub time: Option<f64>,
+    /// Samples represents size of the sample for benchmarks ran
+    #[arg(short = 's', long)]
+    pub samples: Option<usize>,
 }
 
 impl Cmd {
     /// Runs bench command.
     pub fn run(&self) {
+        let mut criterion = criterion::Criterion::default()
+            .warm_up_time(std::time::Duration::from_secs_f64(
+                self.warmup.unwrap_or(0.5),
+            ))
+            // Measurement_time of 0.1 will get 500+ iterations for analysis and transfer and will be extended if needed in order to test the given sample size (minimum sample size is 10 per criterion documentation) as is the case with burntpix and snailtracer benchmark tests
+            .measurement_time(std::time::Duration::from_secs_f64(self.time.unwrap_or(0.1)))
+            .sample_size(self.samples.unwrap_or(10));
+
         match self.name {
-            BenchName::Analysis => analysis::run(),
-            BenchName::Burntpix => burntpix::run(),
-            BenchName::Snailtracer => snailtracer::run(),
-            BenchName::Transfer => transfer::run(),
+            BenchName::Analysis => {
+                analysis::run(&mut criterion);
+            }
+            BenchName::Burntpix => {
+                burntpix::run(&mut criterion);
+            }
+            BenchName::Snailtracer => {
+                snailtracer::run(&mut criterion);
+            }
+            BenchName::Transfer => {
+                transfer::run(&mut criterion);
+            }
         }
     }
 }
diff --git a/bins/revme/src/cmd/bench/analysis.rs b/bins/revme/src/cmd/bench/analysis.rs
@@ -1,5 +1,4 @@
-use std::time::Instant;
-
+use criterion::Criterion;
 use database::{BenchmarkDB, BENCH_CALLER, BENCH_TARGET};
 use revm::{
     bytecode::Bytecode,
@@ -9,9 +8,8 @@ use revm::{
 
 const BYTES: &str = include_str!("analysis.hex");
 
-pub fn run() {
+pub fn run(criterion: &mut Criterion) {
     let bytecode = Bytecode::new_raw(Bytes::from(hex::decode(BYTES).unwrap()));
-
     // BenchmarkDB is dummy state that implements Database trait.
     let context = Context::mainnet()
         .with_db(BenchmarkDB::new_bytecode(bytecode))
@@ -21,14 +19,10 @@ pub fn run() {
             tx.kind = TxKind::Call(BENCH_TARGET);
             tx.data = bytes!("8035F0CE");
         });
-
     let mut evm = context.build_mainnet();
-
-    let time = Instant::now();
-    let _ = evm.replay();
-    println!("First init: {:?}", time.elapsed());
-
-    let time = Instant::now();
-    let _ = evm.replay();
-    println!("Run: {:?}", time.elapsed());
+    criterion.bench_function("analysis", |b| {
+        b.iter(|| {
+            let _ = evm.replay();
+        });
+    });
 }
diff --git a/bins/revme/src/cmd/bench/burntpix.rs b/bins/revme/src/cmd/bench/burntpix.rs
@@ -1,5 +1,6 @@
 pub mod static_data;
 
+use criterion::Criterion;
 use static_data::{
     BURNTPIX_ADDRESS_ONE, BURNTPIX_ADDRESS_THREE, BURNTPIX_ADDRESS_TWO, BURNTPIX_BYTECODE_FOUR,
     BURNTPIX_BYTECODE_ONE, BURNTPIX_BYTECODE_THREE, BURNTPIX_BYTECODE_TWO, BURNTPIX_MAIN_ADDRESS,
@@ -9,17 +10,15 @@ use static_data::{
 use alloy_sol_types::{sol, SolCall};
 use database::{CacheDB, BENCH_CALLER};
 use revm::{
-    context_interface::result::{ExecutionResult, Output},
     database_interface::EmptyDB,
     primitives::{hex, keccak256, Address, Bytes, TxKind, B256, U256},
     state::{AccountInfo, Bytecode},
     Context, ExecuteEvm, MainBuilder, MainContext,
 };
 
-use std::fs::File;
-use std::{error::Error, time::Instant};
+use std::error::Error;
 
-use std::{io::Write, str::FromStr};
+use std::str::FromStr;
 
 sol! {
     #[derive(Debug, PartialEq, Eq)]
@@ -28,7 +27,7 @@ sol! {
     }
 }
 
-pub fn run() {
+pub fn run(criterion: &mut Criterion) {
     let (seed, iterations) = try_init_env_vars().expect("Failed to parse env vars");
 
     let run_call_data = IBURNTPIX::runCall { seed, iterations }.abi_encode();
@@ -45,46 +44,52 @@ pub fn run() {
         })
         .build_mainnet();
 
-    let started = Instant::now();
-    let tx_result = evm.replay().unwrap().result;
-    let return_data = match tx_result {
-        ExecutionResult::Success {
-            output, gas_used, ..
-        } => {
-            println!("Gas used: {:?}", gas_used);
-            println!("Time elapsed: {:?}", started.elapsed());
-            match output {
-                Output::Call(value) => value,
-                _ => unreachable!("Unexpected output type"),
-            }
-        }
-        _ => unreachable!("Execution failed: {:?}", tx_result),
-    };
-
-    // Remove returndata offset and length from output
-    let returndata_offset = 64;
-    let data = &return_data[returndata_offset..];
-
-    // Remove trailing zeros
-    let trimmed_data = data
-        .split_at(data.len() - data.iter().rev().filter(|&x| *x == 0).count())
-        .0;
-    let file_name = format!("{}_{}", seed, iterations);
-
-    svg(file_name, trimmed_data).expect("Failed to store svg");
+    criterion.bench_function("burntpix", |b| {
+        b.iter(|| {
+            let _ = evm.replay();
+        })
+    });
+
+    // Collects the data and uses it to generate the svg after running the benchmark
+    //     let return_data = match tx_result {
+    //         ExecutionResult::Success {
+    //             output, gas_used, ..
+    //         } => {
+    //             println!("Gas used: {:?}", gas_used);
+    //             match output {
+    //                 Output::Call(value) => value,
+    //                 _ => unreachable!("Unexpected output type"),
+    //             }
+    //         }
+    //         _ => unreachable!("Execution failed: {:?}", tx_result),
+    //     };
+
+    //     // Remove returndata offset and length from output
+    //     let returndata_offset = 64;
+    //     let data = &return_data[returndata_offset..];
+
+    //     // Remove trailing zeros
+    //     let trimmed_data = data
+    //         .split_at(data.len() - data.iter().rev().filter(|&x| *x == 0).count())
+    //         .0;
+    //     let file_name = format!("{}_{}", seed, iterations);
+
+    //     svg(file_name, trimmed_data).expect("Failed to store svg");
+    // });
 }
 
-fn svg(filename: String, svg_data: &[u8]) -> Result<(), Box<dyn Error>> {
-    let current_dir = std::env::current_dir()?;
-    let svg_dir = current_dir.join("burntpix").join("svgs");
-    std::fs::create_dir_all(&svg_dir)?;
+// Actually generates the svg
+// fn svg(filename: String, svg_data: &[u8]) -> Result<(), Box<dyn Error>> {
+//     let current_dir = std::env::current_dir()?;
+//     let svg_dir = current_dir.join("burntpix").join("svgs");
+//     std::fs::create_dir_all(&svg_dir)?;
 
-    let file_path = svg_dir.join(format!("{}.svg", filename));
-    let mut file = File::create(file_path)?;
-    file.write_all(svg_data)?;
+//     let file_path = svg_dir.join(format!("{}.svg", filename));
+//     let mut file = File::create(file_path)?;
+//     file.write_all(svg_data)?;
 
-    Ok(())
-}
+//     Ok(())
+// }
 
 const DEFAULT_SEED: &str = "0";
 const DEFAULT_ITERATIONS: &str = "0x7A120";
diff --git a/bins/revme/src/cmd/bench/snailtracer.rs b/bins/revme/src/cmd/bench/snailtracer.rs
@@ -1,11 +1,14 @@
+use criterion::Criterion;
 use database::{BenchmarkDB, BENCH_CALLER, BENCH_TARGET};
 use revm::{
     bytecode::Bytecode,
     primitives::{bytes, hex, Bytes, TxKind},
     Context, ExecuteEvm, MainBuilder, MainContext,
 };
 
-pub fn simple_example(bytecode: Bytecode) {
+pub fn run(criterion: &mut Criterion) {
+    let bytecode = Bytecode::new_raw(Bytes::from(hex::decode(BYTES).unwrap()));
+
     let mut evm = Context::mainnet()
         .with_db(BenchmarkDB::new_bytecode(bytecode.clone()))
         .modify_tx_chained(|tx| {
@@ -16,16 +19,11 @@ pub fn simple_example(bytecode: Bytecode) {
             tx.gas_limit = 1_000_000_000;
         })
         .build_mainnet();
-    let _ = evm.replay().unwrap();
-}
-
-pub fn run() {
-    println!("Running snailtracer example!");
-    let bytecode = Bytecode::new_raw(Bytes::from(hex::decode(BYTES).unwrap()));
-    let start = std::time::Instant::now();
-    simple_example(bytecode);
-    let elapsed = start.elapsed();
-    println!("elapsed: {:?}", elapsed);
+    criterion.bench_function("snailtracer", |b| {
+        b.iter(|| {
+            let _ = evm.replay().unwrap();
+        })
+    });
 }
 
 const BYTES: &str = include_str!("snailtracer.hex");
diff --git a/bins/revme/src/cmd/bench/transfer.rs b/bins/revme/src/cmd/bench/transfer.rs
@@ -1,13 +1,12 @@
+use criterion::Criterion;
 use database::{BenchmarkDB, BENCH_CALLER, BENCH_TARGET};
 use revm::{
     bytecode::Bytecode,
     primitives::{TxKind, U256},
     Context, ExecuteEvm, MainBuilder, MainContext,
 };
-use std::time::Instant;
 
-pub fn run() {
-    let time = Instant::now();
+pub fn run(criterion: &mut Criterion) {
     let mut evm = Context::mainnet()
         .with_db(BenchmarkDB::new_bytecode(Bytecode::new()))
         .modify_tx_chained(|tx| {
@@ -17,13 +16,9 @@ pub fn run() {
             tx.value = U256::from(10);
         })
         .build_mainnet();
-    println!("Init: {:?}", time.elapsed());
-
-    let time = Instant::now();
-    let _ = evm.replay();
-    println!("First run: {:?}", time.elapsed());
-
-    let time = Instant::now();
-    let _ = evm.replay();
-    println!("Second run: {:?}", time.elapsed());
+    criterion.bench_function("transfer", |b| {
+        b.iter(|| {
+            let _ = evm.replay();
+        })
+    });
 }
