From 4781ea6561afa0b6420f4840195ec2763ec44f60 Mon Sep 17 00:00:00 2001
From: "coderabbitai[bot]"
 <136622811+coderabbitai[bot]@users.noreply.github.com>
Date: Sat, 12 Sep 2026 12:52:51 +0000
Subject: [PATCH] =?UTF-8?q?=F0=9F=93=9D=20CodeRabbit=20Chat:=20Update=20ad?=
 =?UTF-8?q?d=20benchmark=20to=20use=20built-in=20Itme=20module?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 rust/benches/optimize/add_bench.rs | 42 +++++++++++++-----------------
 1 file changed, 18 insertions(+), 24 deletions(-)

diff --git a/rust/benches/optimize/add_bench.rs b/rust/benches/optimize/add_bench.rs
index 1c2bcb0c..9546fccf 100644
--- a/rust/benches/optimize/add_bench.rs
+++ b/rust/benches/optimize/add_bench.rs
@@ -8,12 +8,14 @@
  * http://www.apache.org/licenses/LICENSE-2.0
  */
 
-// TODO: fix
-use criterion::{Bencher, Criterion, criterion_group, criterion_main};
 use lightvm::{LightVM, types::capability::Capability};
-use std::time::Duration;
-fn bench_vm_execution(c: &mut Criterion) {
-  let mut vm = LightVM::new(vec![Capability::Control, Capability::Observe]);
+fn main() {
+  let capabilities = vec![
+    Capability::Control,
+    Capability::Debug,
+    Capability::Observe,
+  ];
+  let mut vm = LightVM::new(capabilities.clone());
   let raw = r#"[
     ["val", "x"],
     ["push", 5],
@@ -21,23 +23,15 @@ fn bench_vm_execution(c: &mut Criterion) {
     ["add", "i16"],
     ["set", "x"]
   ]"#;
-  let optimized_json = LightVM::tools().optimize_bytecode(raw);
-  vm.load(optimized_json.clone());
-  let mut group = c.benchmark_group("LightVM Execution");
-  group.bench_function("add_bench", |b: &mut Bencher| {
-    b.iter(|| vm.run(None));
-  });
-  group.finish();
+  let tools = vm.tools();
+  let optimized_json = tools.optimize_bytecode(raw);
+  let benchmark = tools.bench("add_bench").expect("benchmark requires debug capability");
+  benchmark.run(
+    || {
+      let mut vm = LightVM::new(capabilities.clone());
+      vm.load(optimized_json.clone());
+      vm
+    },
+    |vm| vm.run(None),
+  );
 }
-fn custom_config() -> Criterion {
-  Criterion::default()
-    .sample_size(300)
-    .measurement_time(Duration::from_secs(15))
-    .warm_up_time(Duration::from_secs(3))
-}
-criterion_group! {
-  name = benches;
-  config = custom_config();
-  targets = bench_vm_execution
-}
-criterion_main!(benches);