starknet_committer,starknet_patricia: add subtree trait

Author: ArielElp

crates/starknet_committer/src/db.rs

@@ -1,6 +1,4 @@
-pub mod create_facts_tree;
 #[cfg(any(feature = "testing", test))]
 pub mod external_test_utils;
 pub mod facts_db;
 pub mod forest_trait;
-pub mod traversal;

crates/starknet_committer/src/db/external_test_utils.rs

@@ -17,7 +17,7 @@ use starknet_patricia::patricia_merkle_tree::updated_skeleton_tree::tree::{
 };
 use starknet_patricia_storage::map_storage::MapStorage;
 
-use crate::db::create_facts_tree::create_original_skeleton_tree;
+use crate::db::facts_db::create_facts_tree::create_original_skeleton_tree;
 
 pub async fn tree_computation_flow<L, TH>(
     leaf_modifications: LeafModifications<L>,

…et_committer/src/db/create_facts_tree.rs …ter/src/db/facts_db/create_facts_tree.rscrates/starknet_committer/src/db/create_facts_tree.rs renamed to crates/starknet_committer/src/db/facts_db/create_facts_tree.rs crates/starknet_committer/src/db/create_facts_tree.rs renamed to crates/starknet_committer/src/db/facts_db/create_facts_tree.rs

@@ -18,12 +18,13 @@ use starknet_patricia::patricia_merkle_tree::original_skeleton_tree::tree::{
     OriginalSkeletonTreeImpl,
     OriginalSkeletonTreeResult,
 };
-use starknet_patricia::patricia_merkle_tree::traversal::SubTree;
+use starknet_patricia::patricia_merkle_tree::traversal::SubTreeTrait;
 use starknet_patricia::patricia_merkle_tree::types::{NodeIndex, SortedLeafIndices};
 use starknet_patricia_storage::storage_trait::Storage;
 use tracing::warn;
 
-use crate::db::traversal::calculate_subtrees_roots;
+use crate::db::facts_db::traversal::calculate_subtrees_roots;
+use crate::db::facts_db::types::FactsSubTree;
 
 #[cfg(test)]
 #[path = "create_facts_tree_test.rs"]
@@ -42,7 +43,7 @@ macro_rules! log_trivial_modification {
 /// encountering a trivial modification. Fills the previous leaf values if it is not none.
 async fn fetch_nodes<'a, L: Leaf>(
     skeleton_tree: &mut OriginalSkeletonTreeImpl<'a>,
-    subtrees: Vec<SubTree<'a>>,
+    subtrees: Vec<FactsSubTree<'a>>,
     storage: &mut impl Storage,
     leaf_modifications: &LeafModifications<L>,
     config: &impl OriginalSkeletonTreeConfig<L>,
@@ -101,7 +102,7 @@ async fn fetch_nodes<'a, L: Leaf>(
                         leaves.extend(
                             previously_empty_leaves_indices
                                 .iter()
-                                .map(|idx| (**idx, L::default()))
+                                .map(|idx| (*idx, L::default()))
                                 .collect::<HashMap<NodeIndex, L>>(),
                         );
                     }
@@ -161,7 +162,7 @@ pub async fn create_original_skeleton_tree<'a, L: Leaf>(
         )?;
         return Ok(OriginalSkeletonTreeImpl::create_empty(sorted_leaf_indices));
     }
-    let main_subtree = SubTree { sorted_leaf_indices, root_index: NodeIndex::ROOT, root_hash };
+    let main_subtree = FactsSubTree { sorted_leaf_indices, root_index: NodeIndex::ROOT, root_hash };
     let mut skeleton_tree = OriginalSkeletonTreeImpl { nodes: HashMap::new(), sorted_leaf_indices };
     fetch_nodes::<L>(
         &mut skeleton_tree,
@@ -192,7 +193,7 @@ pub async fn create_original_skeleton_tree_and_get_previous_leaves<'a, L: Leaf>(
             sorted_leaf_indices.get_indices().iter().map(|idx| (*idx, L::default())).collect(),
         ));
     }
-    let main_subtree = SubTree { sorted_leaf_indices, root_index: NodeIndex::ROOT, root_hash };
+    let main_subtree = FactsSubTree { sorted_leaf_indices, root_index: NodeIndex::ROOT, root_hash };
     let mut skeleton_tree = OriginalSkeletonTreeImpl { nodes: HashMap::new(), sorted_leaf_indices };
     let mut leaves = HashMap::new();
     fetch_nodes::<L>(
@@ -229,8 +230,8 @@ pub async fn get_leaves<'a, L: Leaf>(
 /// referred subtree or not.
 fn handle_subtree<'a>(
     skeleton_tree: &mut OriginalSkeletonTreeImpl<'a>,
-    next_subtrees: &mut Vec<SubTree<'a>>,
-    subtree: SubTree<'a>,
+    next_subtrees: &mut Vec<FactsSubTree<'a>>,
+    subtree: FactsSubTree<'a>,
     should_fetch_modified_leaves: bool,
 ) {
     if !subtree.is_leaf() || (should_fetch_modified_leaves && !subtree.is_unmodified()) {

…mmitter/src/db/create_facts_tree_test.rs …rc/db/facts_db/create_facts_tree_test.rscrates/starknet_committer/src/db/create_facts_tree_test.rs renamed to crates/starknet_committer/src/db/facts_db/create_facts_tree_test.rs crates/starknet_committer/src/db/create_facts_tree_test.rs renamed to crates/starknet_committer/src/db/facts_db/create_facts_tree_test.rs

@@ -23,7 +23,7 @@ use starknet_patricia_storage::map_storage::MapStorage;
 use starknet_patricia_storage::storage_trait::{DbHashMap, DbKey, DbValue};
 use starknet_types_core::felt::Felt;
 
-use crate::db::create_facts_tree::create_original_skeleton_tree;
+use crate::db::facts_db::create_facts_tree::create_original_skeleton_tree;
 
 #[tokio::test]
 #[rstest]

…es/starknet_committer/src/db/facts_db.rs …starknet_committer/src/db/facts_db/db.rscrates/starknet_committer/src/db/facts_db.rs renamed to crates/starknet_committer/src/db/facts_db/db.rs crates/starknet_committer/src/db/facts_db.rs renamed to crates/starknet_committer/src/db/facts_db/db.rs

@@ -14,7 +14,7 @@ use crate::block_committer::input::{
     ConfigImpl,
     StarknetStorageValue,
 };
-use crate::db::create_facts_tree::{
+use crate::db::facts_db::create_facts_tree::{
     create_original_skeleton_tree,
     create_original_skeleton_tree_and_get_previous_leaves,
 };

crates/starknet_committer/src/db/facts_db/mod.rs

@@ -0,0 +1,4 @@
+pub mod create_facts_tree;
+pub mod db;
+pub mod traversal;
+pub mod types;

…s/starknet_committer/src/db/traversal.rs …t_committer/src/db/facts_db/traversal.rscrates/starknet_committer/src/db/traversal.rs renamed to crates/starknet_committer/src/db/facts_db/traversal.rs crates/starknet_committer/src/db/traversal.rs renamed to crates/starknet_committer/src/db/facts_db/traversal.rs

@@ -8,18 +8,20 @@ use starknet_patricia::patricia_merkle_tree::node_data::inner_node::{
     PreimageMap,
 };
 use starknet_patricia::patricia_merkle_tree::node_data::leaf::Leaf;
-use starknet_patricia::patricia_merkle_tree::traversal::{SubTree, TraversalResult};
+use starknet_patricia::patricia_merkle_tree::traversal::{SubTreeTrait, TraversalResult};
 use starknet_patricia::patricia_merkle_tree::types::{NodeIndex, SortedLeafIndices};
 use starknet_patricia_storage::errors::StorageError;
 use starknet_patricia_storage::storage_trait::{create_db_key, DbKey, Storage};
 
+use crate::db::facts_db::types::FactsSubTree;
+
 #[cfg(test)]
 #[path = "traversal_test.rs"]
 pub mod traversal_test;
 
 // TODO(Aviv, 17/07/2024): Split between storage prefix implementation and function logic.
 pub async fn calculate_subtrees_roots<'a, L: Leaf>(
-    subtrees: &[SubTree<'a>],
+    subtrees: &[FactsSubTree<'a>],
     storage: &mut impl Storage,
 ) -> TraversalResult<Vec<FilledNode<L>>> {
     let mut subtrees_roots = vec![];
@@ -54,7 +56,7 @@ pub async fn fetch_patricia_paths<L: Leaf>(
         return Ok(witnesses);
     }
 
-    let main_subtree = SubTree { sorted_leaf_indices, root_index: NodeIndex::ROOT, root_hash };
+    let main_subtree = FactsSubTree { sorted_leaf_indices, root_index: NodeIndex::ROOT, root_hash };
 
     fetch_patricia_paths_inner::<L>(storage, vec![main_subtree], &mut witnesses, leaves).await?;
     Ok(witnesses)
@@ -67,9 +69,9 @@ pub async 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.
-async fn fetch_patricia_paths_inner<'a, L: Leaf>(
+pub(crate) async fn fetch_patricia_paths_inner<'a, L: Leaf>(
     storage: &mut impl Storage,
-    subtrees: Vec<SubTree<'a>>,
+    subtrees: Vec<FactsSubTree<'a>>,
     witnesses: &mut PreimageMap,
     mut leaves: Option<&mut HashMap<NodeIndex, L>>,
 ) -> TraversalResult<()> {
@@ -103,7 +105,7 @@ async fn fetch_patricia_paths_inner<'a, L: Leaf>(
                         // Insert empty leaves descendent of the current subtree, that are not
                         // descendents of the bottom node.
                         for index in empty_leaves_indices {
-                            leaves_map.insert(*index, L::default());
+                            leaves_map.insert(index, L::default());
                         }
                     }
                     next_subtrees.push(bottom_subtree);

…rknet_committer/src/db/traversal_test.rs …mitter/src/db/facts_db/traversal_test.rscrates/starknet_committer/src/db/traversal_test.rs renamed to crates/starknet_committer/src/db/facts_db/traversal_test.rs crates/starknet_committer/src/db/traversal_test.rs renamed to crates/starknet_committer/src/db/facts_db/traversal_test.rs

@@ -25,14 +25,15 @@ use starknet_patricia::patricia_merkle_tree::node_data::inner_node::{
     Preimage,
     PreimageMap,
 };
-use starknet_patricia::patricia_merkle_tree::traversal::SubTree;
+use starknet_patricia::patricia_merkle_tree::traversal::SubTreeTrait;
 use starknet_patricia::patricia_merkle_tree::types::{SortedLeafIndices, SubTreeHeight};
 use starknet_patricia_storage::map_storage::MapStorage;
 use starknet_patricia_storage::storage_trait::{DbHashMap, DbKey, DbValue};
 use starknet_types_core::felt::Felt;
 use starknet_types_core::hash::Pedersen;
 
-use super::fetch_patricia_paths_inner;
+use crate::db::facts_db::traversal::fetch_patricia_paths_inner;
+use crate::db::facts_db::types::FactsSubTree;
 
 fn to_preimage_map(raw_preimages: HashMap<u32, Vec<u32>>) -> PreimageMap {
     raw_preimages
@@ -73,7 +74,7 @@ async fn test_fetch_patricia_paths_inner_impl(
         .iter()
         .map(|&idx| small_tree_index_to_full(U256::from(idx), height))
         .collect::<Vec<_>>();
-    let main_subtree = SubTree {
+    let main_subtree = FactsSubTree {
         sorted_leaf_indices: SortedLeafIndices::new(&mut leaf_indices),
         root_index: small_tree_index_to_full(U256::ONE, height),
         root_hash,
@@ -514,10 +515,10 @@ struct TestPatriciaPathsInput {
 /// The files names indicate the tree height, number of initial leaves and number of modified
 /// leaves. The hash function used in the python tests is Pedersen.
 /// The leaves values are their NodeIndices.
-#[case(include_str!("../../resources/fetch_patricia_paths_test_10_200_50.json"))]
-#[case(include_str!("../../resources/fetch_patricia_paths_test_10_5_2.json"))]
-#[case(include_str!("../../resources/fetch_patricia_paths_test_10_100_30.json"))]
-#[case(include_str!("../../resources/fetch_patricia_paths_test_8_120_70.json"))]
+#[case(include_str!("../../../resources/fetch_patricia_paths_test_10_200_50.json"))]
+#[case(include_str!("../../../resources/fetch_patricia_paths_test_10_5_2.json"))]
+#[case(include_str!("../../../resources/fetch_patricia_paths_test_10_100_30.json"))]
+#[case(include_str!("../../../resources/fetch_patricia_paths_test_8_120_70.json"))]
 async fn test_fetch_patricia_paths_inner_from_json(#[case] input_data: &str) {
     let input: TestPatriciaPathsInput = serde_json::from_str(input_data)
         .unwrap_or_else(|error| panic!("JSON was not well-formatted: {error:?}"));

crates/starknet_committer/src/db/facts_db/types.rs

@@ -0,0 +1,42 @@
+use starknet_api::hash::HashOutput;
+use starknet_patricia::patricia_merkle_tree::filled_tree::node_serde::PatriciaPrefix;
+use starknet_patricia::patricia_merkle_tree::node_data::leaf::Leaf;
+use starknet_patricia::patricia_merkle_tree::traversal::SubTreeTrait;
+use starknet_patricia::patricia_merkle_tree::types::{NodeIndex, SortedLeafIndices};
+
+#[derive(Debug, PartialEq)]
+pub struct FactsSubTree<'a> {
+    pub sorted_leaf_indices: SortedLeafIndices<'a>,
+    pub root_index: NodeIndex,
+    pub root_hash: HashOutput,
+}
+
+impl<'a> SubTreeTrait<'a> for FactsSubTree<'a> {
+    type ChildData = HashOutput;
+    fn create_child(
+        sorted_leaf_indices: SortedLeafIndices<'a>,
+        root_index: NodeIndex,
+        child_data: Self::ChildData,
+    ) -> Self {
+        Self { sorted_leaf_indices, root_index, root_hash: child_data }
+    }
+    fn get_root_index(&self) -> NodeIndex {
+        self.root_index
+    }
+    fn get_sorted_leaf_indices(&self) -> &SortedLeafIndices<'a> {
+        &self.sorted_leaf_indices
+    }
+    fn should_traverse_unmodified_children() -> bool {
+        false
+    }
+}
+
+impl<'a> FactsSubTree<'a> {
+    pub fn get_root_prefix<L: Leaf>(&self) -> PatriciaPrefix {
+        if self.is_leaf() {
+            PatriciaPrefix::Leaf(L::get_static_prefix())
+        } else {
+            PatriciaPrefix::InnerNode
+        }
+    }
+}

crates/starknet_committer/src/forest/skeleton_forest_test.rs

@@ -32,7 +32,7 @@ use crate::block_committer::input::{
     StarknetStorageValue,
     StateDiff,
 };
-use crate::db::facts_db::FactsDb;
+use crate::db::facts_db::db::FactsDb;
 use crate::db::forest_trait::ForestReader;
 use crate::forest::original_skeleton_forest::{ForestSortedIndices, OriginalSkeletonForest};
 use crate::patricia_merkle_tree::leaf::leaf_impl::ContractState;