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);