From 8d101ffee93229ab026e9f2c3dd71db4e1a51792 Mon Sep 17 00:00:00 2001
From: amos rothberg <amos@starkware.co>
Date: Tue, 30 Jul 2024 11:49:22 +0300
Subject: [PATCH] Pass benchmarks.

---
 .../external_test_utils.rs                    |  7 +++---
 .../committer_cli/benches/committer_bench.rs  | 25 +++++++++++--------
 2 files changed, 17 insertions(+), 15 deletions(-)

diff --git a/crates/committer/src/patricia_merkle_tree/external_test_utils.rs b/crates/committer/src/patricia_merkle_tree/external_test_utils.rs
index 1f4a8dcc..83590a0f 100644
--- a/crates/committer/src/patricia_merkle_tree/external_test_utils.rs
+++ b/crates/committer/src/patricia_merkle_tree/external_test_utils.rs
@@ -68,7 +68,7 @@ pub fn get_random_u256<R: Rng>(rng: &mut R, low: U256, high: U256) -> U256 {
 }
 
 pub async fn tree_computation_flow(
-    leaf_modifications: Arc<LeafModifications<StarknetStorageValue>>,
+    leaf_modifications: LeafModifications<StarknetStorageValue>,
     storage: &MapStorage,
     root_hash: HashOutput,
 ) -> StorageTrie {
@@ -98,7 +98,7 @@ pub async fn tree_computation_flow(
 
     StorageTrie::create_no_additional_output::<TreeHashFunctionImpl>(
         updated_skeleton.into(),
-        leaf_modifications,
+        Arc::new(leaf_modifications),
     )
     .await
     .expect("Failed to create the filled tree")
@@ -115,8 +115,7 @@ pub async fn single_tree_flow_test(
         .map(|(k, v)| (NodeIndex::FIRST_LEAF + k, v))
         .collect::<LeafModifications<StarknetStorageValue>>();
 
-    let filled_tree =
-        tree_computation_flow(Arc::new(leaf_modifications), &storage, root_hash).await;
+    let filled_tree = tree_computation_flow(leaf_modifications, &storage, root_hash).await;
 
     let hash_result = filled_tree.get_root_hash();
 
diff --git a/crates/committer_cli/benches/committer_bench.rs b/crates/committer_cli/benches/committer_bench.rs
index a83442c4..daa31750 100644
--- a/crates/committer_cli/benches/committer_bench.rs
+++ b/crates/committer_cli/benches/committer_bench.rs
@@ -7,7 +7,7 @@
 // Then upload the new files to GCS with this new prefix (run e.g.,
 // gcloud storage cp LOCAL_FILE gs://committer-testing-artifacts/NEW_PREFIX/tree_flow_inputs.json).
 
-use std::{collections::HashMap, sync::Arc};
+use std::collections::HashMap;
 
 use committer::{
     block_committer::input::StarknetStorageValue,
@@ -17,7 +17,7 @@ use committer::{
     },
 };
 use committer_cli::{commands::parse_and_commit, tests::utils::parse_from_python::TreeFlowInput};
-use criterion::{criterion_group, criterion_main, Criterion};
+use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
 
 const CONCURRENCY_MODE: bool = true;
 const SINGLE_TREE_FLOW_INPUT: &str = include_str!("tree_flow_inputs.json");
@@ -42,16 +42,19 @@ pub fn single_tree_flow_benchmark(criterion: &mut Criterion) {
         .into_iter()
         .map(|(k, v)| (NodeIndex::FIRST_LEAF + k, v))
         .collect::<LeafModifications<StarknetStorageValue>>();
-    let arc_leaf_modifications = Arc::new(leaf_modifications);
 
-    criterion.bench_function("tree_computation_flow", |benchmark| {
-        benchmark.iter(|| {
-            runtime.block_on(tree_computation_flow(
-                Arc::clone(&arc_leaf_modifications),
-                &storage,
-                root_hash,
-            ));
-        })
+    criterion.bench_function("tree_computation_flow", move |b| {
+        b.iter_batched(
+            || leaf_modifications.clone(),
+            |leaf_modifications_input| {
+                runtime.block_on(tree_computation_flow(
+                    leaf_modifications_input,
+                    &storage,
+                    root_hash,
+                ));
+            },
+            BatchSize::LargeInput,
+        )
     });
 }