refactor: make external tests generic (#140)

crates/committer/src/patricia_merkle_tree/external_test_utils.rs

@@ -5,14 +5,13 @@ use ethnum::U256;
 use rand::Rng;
 use serde_json::json;
 
-use super::filled_tree::tree::{FilledTree, StorageTrie};
+use super::filled_tree::tree::{FilledTree, FilledTreeImpl};
 use super::node_data::leaf::{Leaf, LeafModifications, SkeletonLeaf};
-use super::original_skeleton_tree::config::OriginalSkeletonStorageTrieConfig;
+use super::original_skeleton_tree::config::OriginalSkeletonTreeConfig;
 use super::original_skeleton_tree::tree::{OriginalSkeletonTree, OriginalSkeletonTreeImpl};
 use super::types::{NodeIndex, SortedLeafIndices};
-use super::updated_skeleton_tree::hash_function::TreeHashFunctionImpl;
+use super::updated_skeleton_tree::hash_function::TreeHashFunction;
 use super::updated_skeleton_tree::tree::{UpdatedSkeletonTree, UpdatedSkeletonTreeImpl};
-use crate::block_committer::input::StarknetStorageValue;
 use crate::felt::Felt;
 use crate::hash::hash_trait::HashOutput;
 use crate::patricia_merkle_tree::errors::TypesError;
@@ -63,12 +62,16 @@ pub fn get_random_u256<R: Rng>(rng: &mut R, low: U256, high: U256) -> U256 {
     result
 }
 
-pub async fn tree_computation_flow(
-    leaf_modifications: Arc<LeafModifications<StarknetStorageValue>>,
+pub async fn tree_computation_flow<L, TH>(
+    leaf_modifications: Arc<LeafModifications<L>>,
     storage: &MapStorage,
     root_hash: HashOutput,
-) -> StorageTrie {
-    let config = OriginalSkeletonStorageTrieConfig::new(&leaf_modifications, false);
+    config: impl OriginalSkeletonTreeConfig<L>,
+) -> FilledTreeImpl<L>
+where
+    TH: TreeHashFunction<L> + 'static,
+    L: Leaf + 'static,
+{
     let mut sorted_leaf_indices: Vec<NodeIndex> = leaf_modifications.keys().copied().collect();
     let sorted_leaf_indices = SortedLeafIndices::new(&mut sorted_leaf_indices);
     let mut original_skeleton =
@@ -92,24 +95,26 @@ pub async fn tree_computation_flow(
     )
     .expect("Failed to create the updated skeleton tree");
 
-    StorageTrie::create::<TreeHashFunctionImpl>(updated_skeleton.into(), leaf_modifications)
+    FilledTreeImpl::<L>::create::<TH>(updated_skeleton.into(), leaf_modifications)
         .await
         .expect("Failed to create the filled tree")
 }
 
-pub async fn single_tree_flow_test(
-    leaf_modifications: LeafModifications<StarknetStorageValue>,
+pub async fn single_tree_flow_test<L: Leaf + 'static, TH: TreeHashFunction<L> + 'static>(
+    leaf_modifications: LeafModifications<L>,
     storage: MapStorage,
     root_hash: HashOutput,
+    config: impl OriginalSkeletonTreeConfig<L>,
 ) -> String {
     // Move from leaf number to actual index.
     let leaf_modifications = leaf_modifications
         .into_iter()
         .map(|(k, v)| (NodeIndex::FIRST_LEAF + k, v))
-        .collect::<LeafModifications<StarknetStorageValue>>();
+        .collect::<LeafModifications<L>>();
 
     let filled_tree =
-        tree_computation_flow(Arc::new(leaf_modifications), &storage, root_hash).await;
+        tree_computation_flow::<L, TH>(Arc::new(leaf_modifications), &storage, root_hash, config)
+            .await;
 
     let hash_result = filled_tree.get_root_hash();
 

crates/committer/src/patricia_merkle_tree/internal_test_utils.rs

@@ -28,7 +28,7 @@ use crate::storage::db_object::{DBObject, Deserializable};
 use crate::storage::storage_trait::StorageValue;
 
 #[derive(Debug, PartialEq, Clone, Copy, Default, Eq)]
-pub(crate) struct MockLeaf(pub(crate) Felt);
+pub struct MockLeaf(pub(crate) Felt);
 
 impl DBObject for MockLeaf {
     fn serialize(&self) -> StorageValue {

crates/committer/src/patricia_merkle_tree/original_skeleton_tree/config.rs

@@ -6,7 +6,7 @@ use crate::patricia_merkle_tree::original_skeleton_tree::tree::OriginalSkeletonT
 use crate::patricia_merkle_tree::types::NodeIndex;
 
 /// Configures the creation of an original skeleton tree.
-pub(crate) trait OriginalSkeletonTreeConfig<L: Leaf> {
+pub trait OriginalSkeletonTreeConfig<L: Leaf> {
     /// Configures whether modified leaves should be compared to the previous leaves and log out a
     /// warning when encountering a trivial modification.
     fn compare_modified_leaves(&self) -> bool;
@@ -22,14 +22,14 @@ pub(crate) trait OriginalSkeletonTreeConfig<L: Leaf> {
 #[macro_export]
 macro_rules! generate_trie_config {
     ($struct_name:ident, $leaf_type:ty) => {
-        pub(crate) struct $struct_name<'a> {
+        pub struct $struct_name<'a> {
             modifications: &'a LeafModifications<$leaf_type>,
             compare_modified_leaves: bool,
         }
 
         impl<'a> $struct_name<'a> {
             #[allow(dead_code)]
-            pub(crate) fn new(
+            pub fn new(
                 modifications: &'a LeafModifications<$leaf_type>,
                 compare_modified_leaves: bool,
             ) -> Self {

crates/committer/src/patricia_merkle_tree/updated_skeleton_tree/hash_function.rs

@@ -17,7 +17,7 @@ use crate::patricia_merkle_tree::node_data::leaf::{ContractState, Leaf};
 pub mod hash_function_test;
 
 /// Trait for hash functions.
-pub(crate) trait HashFunction {
+pub trait HashFunction {
     /// Computes the hash of the given input.
     fn hash(left: &Felt, right: &Felt) -> HashOutput;
 }
@@ -38,7 +38,7 @@ impl HashFunction for PoseidonHashFunction {
     }
 }
 
-pub(crate) trait TreeHashFunction<L: Leaf> {
+pub trait TreeHashFunction<L: Leaf> {
     /// Computes the hash of the given leaf.
     fn compute_leaf_hash(leaf_data: &L) -> HashOutput;
 

crates/committer_cli/benches/committer_bench.rs

@@ -13,7 +13,9 @@ use std::sync::Arc;
 use committer::block_committer::input::StarknetStorageValue;
 use committer::patricia_merkle_tree::external_test_utils::tree_computation_flow;
 use committer::patricia_merkle_tree::node_data::leaf::LeafModifications;
+use committer::patricia_merkle_tree::original_skeleton_tree::config::OriginalSkeletonStorageTrieConfig;
 use committer::patricia_merkle_tree::types::NodeIndex;
+use committer::patricia_merkle_tree::updated_skeleton_tree::hash_function::TreeHashFunctionImpl;
 use committer_cli::commands::parse_and_commit;
 use committer_cli::tests::utils::parse_from_python::TreeFlowInput;
 use criterion::{criterion_group, criterion_main, Criterion};
@@ -36,14 +38,15 @@ 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);
+    let arc_leaf_modifications = Arc::new(leaf_modifications.clone());
 
     criterion.bench_function("tree_computation_flow", |benchmark| {
         benchmark.iter(|| {
-            runtime.block_on(tree_computation_flow(
+            runtime.block_on(tree_computation_flow::<StarknetStorageValue, TreeHashFunctionImpl>(
                 Arc::clone(&arc_leaf_modifications),
                 &storage,
                 root_hash,
+                OriginalSkeletonStorageTrieConfig::new(&leaf_modifications, false),
             ));
         })
     });

crates/committer_cli/src/tests/python_tests.rs

@@ -26,6 +26,7 @@ use committer::patricia_merkle_tree::node_data::inner_node::{
     PathToBottom,
 };
 use committer::patricia_merkle_tree::node_data::leaf::ContractState;
+use committer::patricia_merkle_tree::original_skeleton_tree::config::OriginalSkeletonStorageTrieConfig;
 use committer::patricia_merkle_tree::types::SubTreeHeight;
 use committer::patricia_merkle_tree::updated_skeleton_tree::hash_function::TreeHashFunctionImpl;
 use committer::storage::db_object::DBObject;
@@ -197,7 +198,13 @@ impl PythonTest {
                 let TreeFlowInput { leaf_modifications, storage, root_hash } =
                     serde_json::from_str(Self::non_optional_input(input)?)?;
                 // 2. Run the test.
-                let output = single_tree_flow_test(leaf_modifications, storage, root_hash).await;
+                let output = single_tree_flow_test::<StarknetStorageValue, TreeHashFunctionImpl>(
+                    leaf_modifications.clone(),
+                    storage,
+                    root_hash,
+                    OriginalSkeletonStorageTrieConfig::new(&leaf_modifications, false),
+                )
+                .await;
                 // 3. Serialize and return output.
                 Ok(output)
             }

crates/committer_cli/src/tests/regression_tests.rs

@@ -2,8 +2,10 @@ use std::collections::HashMap;
 use std::fs;
 
 use clap::Error;
-use committer::block_committer::input::{ConfigImpl, Input};
+use committer::block_committer::input::{ConfigImpl, Input, StarknetStorageValue};
 use committer::patricia_merkle_tree::external_test_utils::single_tree_flow_test;
+use committer::patricia_merkle_tree::original_skeleton_tree::config::OriginalSkeletonStorageTrieConfig;
+use committer::patricia_merkle_tree::updated_skeleton_tree::hash_function::TreeHashFunctionImpl;
 use serde::{Deserialize, Deserializer};
 use serde_json::{Map, Value};
 use tempfile::NamedTempFile;
@@ -108,7 +110,13 @@ pub async fn test_regression_single_tree() {
 
     let start = std::time::Instant::now();
     // Benchmark the single tree flow test.
-    let output = single_tree_flow_test(leaf_modifications, storage, root_hash).await;
+    let output = single_tree_flow_test::<StarknetStorageValue, TreeHashFunctionImpl>(
+        leaf_modifications.clone(),
+        storage,
+        root_hash,
+        OriginalSkeletonStorageTrieConfig::new(&leaf_modifications, false),
+    )
+    .await;
     let execution_time = std::time::Instant::now() - start;
 
     // Assert correctness of the output of the single tree flow test.