starknet_committer,starknet_patricia: add subtree trait

Author: ArielElp

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

@@ -18,13 +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::facts_db::traversal::calculate_subtrees_roots;
-use crate::db::traversal::calculate_subtrees_roots;
+use crate::db::facts_db::types::FactsSubTree;
 
 #[cfg(test)]
 #[path = "create_facts_tree_test.rs"]
@@ -43,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>,
@@ -102,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>>(),
                         );
                     }
@@ -162,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,
@@ -193,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>(
@@ -230,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()) {

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

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

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);

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

@@ -25,15 +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 crate::db::facts_db::traversal::fetch_patricia_paths_inner;
-use super::fetch_patricia_paths_inner;
+use crate::db::facts_db::types::FactsSubTree;
 
 fn to_preimage_map(raw_preimages: HashMap<u32, Vec<u32>>) -> PreimageMap {
     raw_preimages
@@ -74,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,

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_patricia/src/patricia_merkle_tree/traversal.rs

@@ -1,11 +1,8 @@
-use starknet_api::hash::HashOutput;
 use starknet_patricia_storage::errors::{DeserializationError, StorageError};
 use starknet_patricia_storage::storage_trait::PatriciaStorageError;
 use thiserror::Error;
 
-use crate::patricia_merkle_tree::filled_tree::node_serde::PatriciaPrefix;
 use crate::patricia_merkle_tree::node_data::inner_node::PathToBottom;
-use crate::patricia_merkle_tree::node_data::leaf::Leaf;
 use crate::patricia_merkle_tree::original_skeleton_tree::utils::split_leaves;
 use crate::patricia_merkle_tree::types::{NodeIndex, SortedLeafIndices, SubTreeHeight};
 
@@ -25,88 +22,70 @@ pub enum TraversalError {
 
 pub type TraversalResult<T> = Result<T, TraversalError>;
 
-#[derive(Debug, PartialEq)]
-pub struct SubTree<'a> {
-    pub sorted_leaf_indices: SortedLeafIndices<'a>,
-    pub root_index: NodeIndex,
-    pub root_hash: HashOutput,
-}
+pub trait SubTreeTrait<'a>: Sized {
+    type ChildData: Copy;
 
-impl<'a> SubTree<'a> {
-    pub(crate) fn get_height(&self) -> SubTreeHeight {
-        SubTreeHeight::new(SubTreeHeight::ACTUAL_HEIGHT.0 - (self.root_index.bit_length() - 1))
-    }
+    fn create_child(
+        sorted_leaf_indices: SortedLeafIndices<'a>,
+        root_index: NodeIndex,
+        child_data: Self::ChildData,
+    ) -> Self;
+    fn get_root_index(&self) -> NodeIndex;
+    fn get_sorted_leaf_indices(&self) -> &SortedLeafIndices<'a>;
 
-    pub(crate) fn split_leaves(&self) -> [SortedLeafIndices<'a>; 2] {
-        split_leaves(&self.root_index, &self.sorted_leaf_indices)
+    fn get_height(&self) -> SubTreeHeight {
+        SubTreeHeight::new(
+            SubTreeHeight::ACTUAL_HEIGHT.0 - (self.get_root_index().bit_length() - 1),
+        )
     }
-
-    pub fn is_unmodified(&self) -> bool {
-        self.sorted_leaf_indices.is_empty()
+    fn split_leaves(&self) -> [SortedLeafIndices<'a>; 2] {
+        split_leaves(&self.get_root_index(), self.get_sorted_leaf_indices())
     }
-
-    pub fn get_root_prefix<L: Leaf>(&self) -> PatriciaPrefix {
-        if self.is_leaf() {
-            PatriciaPrefix::Leaf(L::get_static_prefix())
-        } else {
-            PatriciaPrefix::InnerNode
-        }
+    fn is_unmodified(&self) -> bool {
+        self.get_sorted_leaf_indices().is_empty()
     }
-
     /// Returns the bottom subtree which is referred from `self` by the given path. When creating
     /// the bottom subtree some indices that were modified under `self` are not modified under the
     /// bottom subtree (leaves that were previously empty). These indices are returned as well.
-    pub fn get_bottom_subtree(
+    fn get_bottom_subtree(
         &self,
         path_to_bottom: &PathToBottom,
-        bottom_hash: HashOutput,
-    ) -> (Self, Vec<&NodeIndex>) {
-        let bottom_index = path_to_bottom.bottom_index(self.root_index);
+        bottom_data: Self::ChildData,
+    ) -> (Self, Vec<NodeIndex>) {
+        let sorted_leaf_indices = self.get_sorted_leaf_indices();
+        let bottom_index = path_to_bottom.bottom_index(self.get_root_index());
         let bottom_height = self.get_height() - SubTreeHeight::new(path_to_bottom.length.into());
         let leftmost_in_subtree = bottom_index << bottom_height.into();
         let rightmost_in_subtree =
             leftmost_in_subtree - NodeIndex::ROOT + (NodeIndex::ROOT << bottom_height.into());
-        let leftmost_index = self.sorted_leaf_indices.bisect_left(&leftmost_in_subtree);
-        let rightmost_index = self.sorted_leaf_indices.bisect_right(&rightmost_in_subtree);
-        let bottom_leaves = self.sorted_leaf_indices.subslice(leftmost_index, rightmost_index);
-        let previously_empty_leaf_indices = self.sorted_leaf_indices.get_indices()
-            [..leftmost_index]
+        let leftmost_index = sorted_leaf_indices.bisect_left(&leftmost_in_subtree);
+        let rightmost_index = sorted_leaf_indices.bisect_right(&rightmost_in_subtree);
+        let bottom_leaves = sorted_leaf_indices.subslice(leftmost_index, rightmost_index);
+        let previously_empty_leaf_indices = sorted_leaf_indices.get_indices()[..leftmost_index]
             .iter()
-            .chain(self.sorted_leaf_indices.get_indices()[rightmost_index..].iter())
+            .chain(sorted_leaf_indices.get_indices()[rightmost_index..].iter())
+            .cloned()
             .collect();
 
         (
-            Self {
-                sorted_leaf_indices: bottom_leaves,
-                root_index: bottom_index,
-                root_hash: bottom_hash,
-            },
+            Self::create_child(bottom_leaves, bottom_index, bottom_data),
             previously_empty_leaf_indices,
         )
     }
-
-    pub fn get_children_subtrees(
+    fn get_children_subtrees(
         &self,
-        left_hash: HashOutput,
-        right_hash: HashOutput,
+        left_data: Self::ChildData,
+        right_data: Self::ChildData,
     ) -> (Self, Self) {
         let [left_leaves, right_leaves] = self.split_leaves();
-        let left_root_index = self.root_index * 2.into();
+        let left_root_index = self.get_root_index() * 2.into();
         (
-            SubTree {
-                sorted_leaf_indices: left_leaves,
-                root_index: left_root_index,
-                root_hash: left_hash,
-            },
-            SubTree {
-                sorted_leaf_indices: right_leaves,
-                root_index: left_root_index + NodeIndex::ROOT,
-                root_hash: right_hash,
-            },
+            Self::create_child(left_leaves, left_root_index, left_data),
+            Self::create_child(right_leaves, left_root_index + NodeIndex::ROOT, right_data),
         )
     }
-
-    pub fn is_leaf(&self) -> bool {
-        self.root_index.is_leaf()
+    fn is_leaf(&self) -> bool {
+        self.get_root_index().is_leaf()
     }
+    fn should_traverse_unmodified_children() -> bool;
 }