chore: change bench to allow consuming leaf modifications (#477)

starkware-libs · Aug 15, 2024 · 64ce427 · 64ce427
1 parent 5606fb8
commit 64ce427
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 17 deletions.
diff --git a/crates/committer_cli/benches/committer_bench.rs b/crates/committer_cli/benches/committer_bench.rs
@@ -8,12 +8,11 @@
 // gcloud storage cp LOCAL_FILE gs://committer-testing-artifacts/NEW_PREFIX/tree_flow_inputs.json).
 
 use std::collections::HashMap;
-use std::sync::Arc;
 
 use committer_cli::commands::commit;
 use committer_cli::parse_input::read::parse_input;
 use committer_cli::tests::utils::parse_from_python::TreeFlowInput;
-use criterion::{criterion_group, criterion_main, Criterion};
+use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
 use starknet_committer::block_committer::input::StarknetStorageValue;
 use starknet_committer::hash_function::hash::TreeHashFunctionImpl;
 use starknet_committer::patricia_merkle_tree::tree::OriginalSkeletonStorageTrieConfig;
@@ -39,17 +38,22 @@ 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.clone());
 
-    criterion.bench_function("tree_computation_flow", |benchmark| {
-        benchmark.iter(|| {
-            runtime.block_on(tree_computation_flow::<StarknetStorageValue, TreeHashFunctionImpl>(
-                Arc::clone(&arc_leaf_modifications),
-                &storage,
-                root_hash,
-                OriginalSkeletonStorageTrieConfig::new(&leaf_modifications, false),
-            ));
-        })
+    criterion.bench_function("tree_computation_flow", move |benchmark| {
+        benchmark.iter_batched(
+            || leaf_modifications.clone(),
+            |leaf_modifications_input| {
+                runtime.block_on(
+                    tree_computation_flow::<StarknetStorageValue, TreeHashFunctionImpl>(
+                        leaf_modifications_input,
+                        &storage,
+                        root_hash,
+                        OriginalSkeletonStorageTrieConfig::new(&leaf_modifications, false),
+                    ),
+                );
+            },
+            BatchSize::LargeInput,
+        )
     });
 }
 

diff --git a/crates/starknet_patricia/src/patricia_merkle_tree/external_test_utils.rs b/crates/starknet_patricia/src/patricia_merkle_tree/external_test_utils.rs
@@ -1,5 +1,4 @@
 use std::collections::HashMap;
-use std::sync::Arc;
 
 use ethnum::U256;
 use rand::Rng;
@@ -66,7 +65,7 @@ pub fn get_random_u256<R: Rng>(rng: &mut R, low: U256, high: U256) -> U256 {
 }
 
 pub async fn tree_computation_flow<L, TH>(
-    leaf_modifications: Arc<LeafModifications<L>>,
+    leaf_modifications: LeafModifications<L>,
     storage: &MapStorage,
     root_hash: HashOutput,
     config: impl OriginalSkeletonTreeConfig<L>,
@@ -98,7 +97,7 @@ where
     )
     .expect("Failed to create the updated skeleton tree");
 
-    FilledTreeImpl::<L>::create::<TH>(updated_skeleton.into(), leaf_modifications)
+    FilledTreeImpl::<L>::create::<TH>(updated_skeleton.into(), leaf_modifications.into())
         .await
         .expect("Failed to create the filled tree")
 }
@@ -116,8 +115,7 @@ pub async fn single_tree_flow_test<L: Leaf + 'static, TH: TreeHashFunction<L> +
         .collect::<LeafModifications<L>>();
 
     let filled_tree =
-        tree_computation_flow::<L, TH>(Arc::new(leaf_modifications), &storage, root_hash, config)
-            .await;
+        tree_computation_flow::<L, TH>(leaf_modifications, &storage, root_hash, config).await;
 
     let hash_result = filled_tree.get_root_hash();