starknet_committer,starknet_patricia: don't use impl Storage, use generics

Author: dorimedini-starkware

crates/starknet_committer/src/block_committer/commit.rs

@@ -23,7 +23,7 @@ use crate::patricia_merkle_tree::types::class_hash_into_node_index;
 
 type BlockCommitmentResult<T> = Result<T, BlockCommitmentError>;
 
-pub async fn commit_block<S: Storage>(
+pub async fn commit_block<S: Storage + ?Sized>(
     input: Input<ConfigImpl>,
     storage: &mut S,
     mut time_measurement: Option<&mut TimeMeasurement>,

crates/starknet_committer/src/forest/filled_forest.rs

@@ -34,7 +34,7 @@ pub struct FilledForest {
 impl FilledForest {
     /// Writes the node serialization of the filled trees to storage. Returns the number of new
     /// objects written to storage.
-    pub fn write_to_storage(&self, storage: &mut impl Storage) -> usize {
+    pub fn write_to_storage<S: Storage + ?Sized>(&self, storage: &mut S) -> usize {
         // Serialize all trees to one hash map.
         let new_db_objects: DbHashMap = self
             .storage_tries

crates/starknet_committer/src/forest/original_skeleton_forest.rs

@@ -35,8 +35,8 @@ impl<'a> OriginalSkeletonForest<'a> {
     /// Creates an original skeleton forest that includes the storage tries of the modified
     /// contracts, the classes trie and the contracts trie. Additionally, returns the original
     /// contract states that are needed to compute the contract state tree.
-    pub(crate) fn create(
-        storage: &mut impl Storage,
+    pub(crate) fn create<S: Storage + ?Sized>(
+        storage: &mut S,
         contracts_trie_root_hash: HashOutput,
         classes_trie_root_hash: HashOutput,
         storage_updates: &HashMap<ContractAddress, LeafModifications<StarknetStorageValue>>,
@@ -72,9 +72,9 @@ impl<'a> OriginalSkeletonForest<'a> {
 
     /// Creates the contracts trie original skeleton.
     /// Also returns the previous contracts state of the modified contracts.
-    fn create_contracts_trie(
+    fn create_contracts_trie<S: Storage + ?Sized>(
         contracts_trie_root_hash: HashOutput,
-        storage: &mut impl Storage,
+        storage: &mut S,
         contracts_trie_sorted_indices: SortedLeafIndices<'a>,
     ) -> ForestResult<(OriginalSkeletonTreeImpl<'a>, HashMap<NodeIndex, ContractState>)> {
         Ok(OriginalSkeletonTreeImpl::create_and_get_previous_leaves(
@@ -86,10 +86,10 @@ impl<'a> OriginalSkeletonForest<'a> {
         )?)
     }
 
-    fn create_storage_tries(
+    fn create_storage_tries<S: Storage + ?Sized>(
         actual_storage_updates: &HashMap<ContractAddress, LeafModifications<StarknetStorageValue>>,
         original_contracts_trie_leaves: &HashMap<NodeIndex, ContractState>,
-        storage: &mut impl Storage,
+        storage: &mut S,
         config: &impl Config,
         storage_tries_sorted_indices: &HashMap<ContractAddress, SortedLeafIndices<'a>>,
     ) -> ForestResult<HashMap<ContractAddress, OriginalSkeletonTreeImpl<'a>>> {
@@ -116,10 +116,10 @@ impl<'a> OriginalSkeletonForest<'a> {
         Ok(storage_tries)
     }
 
-    fn create_classes_trie(
+    fn create_classes_trie<S: Storage + ?Sized>(
         actual_classes_updates: &LeafModifications<CompiledClassHash>,
         classes_trie_root_hash: HashOutput,
-        storage: &mut impl Storage,
+        storage: &mut S,
         config: &impl Config,
         contracts_trie_sorted_indices: SortedLeafIndices<'a>,
     ) -> ForestResult<OriginalSkeletonTreeImpl<'a>> {

crates/starknet_committer/src/patricia_merkle_tree/tree.rs

@@ -44,8 +44,8 @@ impl OriginalSkeletonContractsTrieConfig {
 /// Fetch the leaves in the contracts trie only, to be able to get the storage root hashes.
 /// Assumption: `contract_sorted_leaf_indices` contains all `contract_storage_sorted_leaf_indices`
 /// keys.
-fn fetch_all_patricia_paths(
-    storage: &mut impl Storage,
+fn fetch_all_patricia_paths<S: Storage + ?Sized>(
+    storage: &mut S,
     classes_trie_root_hash: HashOutput,
     contracts_trie_root_hash: HashOutput,
     class_sorted_leaf_indices: SortedLeafIndices<'_>,
@@ -70,7 +70,7 @@ fn fetch_all_patricia_paths(
 
     // Classes trie - no need to fetch the leaves.
     let leaves = None;
-    let classes_trie_proof = fetch_patricia_paths::<CompiledClassHash>(
+    let classes_trie_proof = fetch_patricia_paths::<S, CompiledClassHash>(
         storage,
         classes_trie_root_hash,
         class_sorted_leaf_indices,
@@ -79,7 +79,7 @@ fn fetch_all_patricia_paths(
 
     // Contracts trie - the leaves are required.
     let mut leaves = HashMap::new();
-    let contracts_proof_nodes = fetch_patricia_paths::<ContractState>(
+    let contracts_proof_nodes = fetch_patricia_paths::<S, ContractState>(
         storage,
         contracts_trie_root_hash,
         contract_sorted_leaf_indices,
@@ -97,7 +97,7 @@ fn fetch_all_patricia_paths(
             .storage_root_hash;
         // No need to fetch the leaves.
         let leaves = None;
-        let proof = fetch_patricia_paths::<StarknetStorageValue>(
+        let proof = fetch_patricia_paths::<S, StarknetStorageValue>(
             storage,
             storage_root_hash,
             *sorted_leaf_indices,

crates/starknet_patricia/src/patricia_merkle_tree/original_skeleton_tree/create_tree.rs

@@ -33,10 +33,10 @@ impl<'a> OriginalSkeletonTreeImpl<'a> {
     /// Given a list of subtrees, traverses towards their leaves and fetches all non-empty,
     /// unmodified nodes. If `compare_modified_leaves` is set, function logs out a warning when
     /// encountering a trivial modification. Fills the previous leaf values if it is not none.
-    fn fetch_nodes<L: Leaf>(
+    fn fetch_nodes<S: Storage + ?Sized, L: Leaf>(
         &mut self,
         subtrees: Vec<SubTree<'a>>,
-        storage: &mut impl Storage,
+        storage: &mut S,
         leaf_modifications: &LeafModifications<L>,
         config: &impl OriginalSkeletonTreeConfig<L>,
         mut previous_leaves: Option<&mut HashMap<NodeIndex, L>>,
@@ -47,7 +47,7 @@ impl<'a> OriginalSkeletonTreeImpl<'a> {
         let should_fetch_modified_leaves =
             config.compare_modified_leaves() || previous_leaves.is_some();
         let mut next_subtrees = Vec::new();
-        let filled_roots = calculate_subtrees_roots::<L>(&subtrees, storage)?;
+        let filled_roots = calculate_subtrees_roots::<S, L>(&subtrees, storage)?;
         for (filled_root, subtree) in filled_roots.into_iter().zip(subtrees.iter()) {
             match filled_root.data {
                 // Binary node.
@@ -127,11 +127,17 @@ impl<'a> OriginalSkeletonTreeImpl<'a> {
                 }
             }
         }
-        self.fetch_nodes::<L>(next_subtrees, storage, leaf_modifications, config, previous_leaves)
+        self.fetch_nodes::<S, L>(
+            next_subtrees,
+            storage,
+            leaf_modifications,
+            config,
+            previous_leaves,
+        )
     }
 
-    pub(crate) fn create_impl<L: Leaf>(
-        storage: &mut impl Storage,
+    pub(crate) fn create_impl<S: Storage + ?Sized, L: Leaf>(
+        storage: &mut S,
         root_hash: HashOutput,
         sorted_leaf_indices: SortedLeafIndices<'a>,
         config: &impl OriginalSkeletonTreeConfig<L>,
@@ -150,7 +156,7 @@ impl<'a> OriginalSkeletonTreeImpl<'a> {
         }
         let main_subtree = SubTree { sorted_leaf_indices, root_index: NodeIndex::ROOT, root_hash };
         let mut skeleton_tree = Self { nodes: HashMap::new(), sorted_leaf_indices };
-        skeleton_tree.fetch_nodes::<L>(
+        skeleton_tree.fetch_nodes::<S, L>(
             vec![main_subtree],
             storage,
             leaf_modifications,
@@ -160,8 +166,8 @@ impl<'a> OriginalSkeletonTreeImpl<'a> {
         Ok(skeleton_tree)
     }
 
-    pub(crate) fn create_and_get_previous_leaves_impl<L: Leaf>(
-        storage: &mut impl Storage,
+    pub(crate) fn create_and_get_previous_leaves_impl<S: Storage + ?Sized, L: Leaf>(
+        storage: &mut S,
         root_hash: HashOutput,
         sorted_leaf_indices: SortedLeafIndices<'a>,
         leaf_modifications: &LeafModifications<L>,
@@ -180,7 +186,7 @@ impl<'a> OriginalSkeletonTreeImpl<'a> {
         let main_subtree = SubTree { sorted_leaf_indices, root_index: NodeIndex::ROOT, root_hash };
         let mut skeleton_tree = Self { nodes: HashMap::new(), sorted_leaf_indices };
         let mut leaves = HashMap::new();
-        skeleton_tree.fetch_nodes::<L>(
+        skeleton_tree.fetch_nodes::<S, L>(
             vec![main_subtree],
             storage,
             leaf_modifications,

crates/starknet_patricia/src/patricia_merkle_tree/original_skeleton_tree/create_tree_test.rs

@@ -210,7 +210,7 @@ fn test_create_tree(
     let config = OriginalSkeletonMockTrieConfig::new(compare_modified_leaves);
     let mut sorted_leaf_indices: Vec<NodeIndex> = leaf_modifications.keys().copied().collect();
     let sorted_leaf_indices = SortedLeafIndices::new(&mut sorted_leaf_indices);
-    let skeleton_tree = OriginalSkeletonTreeImpl::create::<MockLeaf>(
+    let skeleton_tree = OriginalSkeletonTreeImpl::create::<MapStorage, MockLeaf>(
         &mut storage,
         root_hash,
         sorted_leaf_indices,

crates/starknet_patricia/src/patricia_merkle_tree/original_skeleton_tree/tree.rs

@@ -20,8 +20,8 @@ pub type OriginalSkeletonTreeResult<T> = Result<T, OriginalSkeletonTreeError>;
 /// update. It also contains the hashes (for edge siblings - also the edge data) of the unmodified
 /// nodes on the Merkle paths from the updated leaves to the root.
 pub trait OriginalSkeletonTree<'a>: Sized {
-    fn create<L: Leaf>(
-        storage: &mut impl Storage,
+    fn create<S: Storage + ?Sized, L: Leaf>(
+        storage: &mut S,
         root_hash: HashOutput,
         sorted_leaf_indices: SortedLeafIndices<'a>,
         config: &impl OriginalSkeletonTreeConfig<L>,
@@ -32,8 +32,8 @@ pub trait OriginalSkeletonTree<'a>: Sized {
 
     fn get_nodes_mut(&mut self) -> &mut OriginalSkeletonNodeMap;
 
-    fn create_and_get_previous_leaves<L: Leaf>(
-        storage: &mut impl Storage,
+    fn create_and_get_previous_leaves<S: Storage + ?Sized, L: Leaf>(
+        storage: &mut S,
         root_hash: HashOutput,
         sorted_leaf_indices: SortedLeafIndices<'a>,
         config: &impl OriginalSkeletonTreeConfig<L>,
@@ -52,8 +52,8 @@ pub struct OriginalSkeletonTreeImpl<'a> {
 }
 
 impl<'a> OriginalSkeletonTree<'a> for OriginalSkeletonTreeImpl<'a> {
-    fn create<L: Leaf>(
-        storage: &mut impl Storage,
+    fn create<S: Storage + ?Sized, L: Leaf>(
+        storage: &mut S,
         root_hash: HashOutput,
         sorted_leaf_indices: SortedLeafIndices<'a>,
         config: &impl OriginalSkeletonTreeConfig<L>,
@@ -70,8 +70,8 @@ impl<'a> OriginalSkeletonTree<'a> for OriginalSkeletonTreeImpl<'a> {
         &mut self.nodes
     }
 
-    fn create_and_get_previous_leaves<L: Leaf>(
-        storage: &mut impl Storage,
+    fn create_and_get_previous_leaves<S: Storage + ?Sized, L: Leaf>(
+        storage: &mut S,
         root_hash: HashOutput,
         sorted_leaf_indices: SortedLeafIndices<'a>,
         config: &impl OriginalSkeletonTreeConfig<L>,
@@ -92,8 +92,8 @@ impl<'a> OriginalSkeletonTree<'a> for OriginalSkeletonTreeImpl<'a> {
 }
 
 impl<'a> OriginalSkeletonTreeImpl<'a> {
-    pub fn get_leaves<L: Leaf>(
-        storage: &mut impl Storage,
+    pub fn get_leaves<S: Storage + ?Sized, L: Leaf>(
+        storage: &mut S,
         root_hash: HashOutput,
         sorted_leaf_indices: SortedLeafIndices<'a>,
     ) -> OriginalSkeletonTreeResult<HashMap<NodeIndex, L>> {

crates/starknet_patricia/src/patricia_merkle_tree/traversal.rs

@@ -125,9 +125,9 @@ impl<'a> SubTree<'a> {
 }
 
 // TODO(Aviv, 17/07/2024): Split between storage prefix implementation and function logic.
-pub(crate) fn calculate_subtrees_roots<'a, L: Leaf>(
+pub(crate) fn calculate_subtrees_roots<'a, S: Storage + ?Sized, L: Leaf>(
     subtrees: &[SubTree<'a>],
-    storage: &mut impl Storage,
+    storage: &mut S,
 ) -> TraversalResult<Vec<FilledNode<L>>> {
     let mut subtrees_roots = vec![];
     let db_keys: Vec<DbKey> = subtrees
@@ -149,8 +149,8 @@ pub(crate) fn calculate_subtrees_roots<'a, L: Leaf>(
 /// given tree according to the `root_hash`.
 /// If `leaves` is not `None`, it also fetches the modified leaves and inserts them into the
 /// provided map.
-pub fn fetch_patricia_paths<L: Leaf>(
-    storage: &mut impl Storage,
+pub fn fetch_patricia_paths<S: Storage + ?Sized, L: Leaf>(
+    storage: &mut S,
     root_hash: HashOutput,
     sorted_leaf_indices: SortedLeafIndices<'_>,
     leaves: Option<&mut HashMap<NodeIndex, L>>,
@@ -163,7 +163,7 @@ pub fn fetch_patricia_paths<L: Leaf>(
 
     let main_subtree = SubTree { sorted_leaf_indices, root_index: NodeIndex::ROOT, root_hash };
 
-    fetch_patricia_paths_inner::<L>(storage, vec![main_subtree], &mut witnesses, leaves)?;
+    fetch_patricia_paths_inner::<S, L>(storage, vec![main_subtree], &mut witnesses, leaves)?;
     Ok(witnesses)
 }
 
@@ -174,8 +174,8 @@ pub fn fetch_patricia_paths<L: Leaf>(
 /// inner nodes in their paths.
 /// If `leaves` is not `None`, it also fetches the modified leaves and inserts them into the
 /// provided map.
-fn fetch_patricia_paths_inner<'a, L: Leaf>(
-    storage: &mut impl Storage,
+fn fetch_patricia_paths_inner<'a, S: Storage + ?Sized, L: Leaf>(
+    storage: &mut S,
     subtrees: Vec<SubTree<'a>>,
     witnesses: &mut PreimageMap,
     mut leaves: Option<&mut HashMap<NodeIndex, L>>,
@@ -184,7 +184,7 @@ fn fetch_patricia_paths_inner<'a, L: Leaf>(
         return Ok(());
     }
     let mut next_subtrees = Vec::new();
-    let filled_roots = calculate_subtrees_roots::<L>(&subtrees, storage)?;
+    let filled_roots = calculate_subtrees_roots::<S, L>(&subtrees, storage)?;
     for (filled_root, subtree) in filled_roots.into_iter().zip(subtrees.iter()) {
         // Always insert root.
         // No need to insert an unmodified node (which is not the root), because its parent is
@@ -225,5 +225,5 @@ fn fetch_patricia_paths_inner<'a, L: Leaf>(
             }
         }
     }
-    fetch_patricia_paths_inner::<L>(storage, next_subtrees, witnesses, leaves)
+    fetch_patricia_paths_inner::<S, L>(storage, next_subtrees, witnesses, leaves)
 }

crates/starknet_patricia/src/patricia_merkle_tree/traversal_test.rs

@@ -648,7 +648,7 @@ fn test_fetch_patricia_paths_inner(
     let mut nodes = HashMap::new();
     let mut fetched_leaves = HashMap::new();
 
-    fetch_patricia_paths_inner::<MockLeaf>(
+    fetch_patricia_paths_inner::<MapStorage, MockLeaf>(
         &mut storage,
         vec![main_subtree],
         &mut nodes,

crates/starknet_patricia/src/patricia_merkle_tree/updated_skeleton_tree/create_tree_helper_test.rs

@@ -498,7 +498,7 @@ async fn test_update_non_modified_storage_tree(#[case] root_hash: HashOutput) {
     let empty_map = HashMap::new();
     let mut empty_storage = MapStorage::default();
     let config = OriginalSkeletonMockTrieConfig::new(false);
-    let mut original_skeleton_tree = OriginalSkeletonTreeImpl::create_impl::<MockLeaf>(
+    let mut original_skeleton_tree = OriginalSkeletonTreeImpl::create_impl::<MapStorage, MockLeaf>(
         &mut empty_storage,
         root_hash,
         SortedLeafIndices::new(&mut []),