starknet_committer,starknet_patricia: add layout for base trait

Author: ArielElp

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

@@ -3,7 +3,7 @@ use std::collections::HashMap;
 use async_trait::async_trait;
 use starknet_api::core::ContractAddress;
 use starknet_api::hash::HashOutput;
-use starknet_patricia::db_layout::{NodeLayout, TrieType};
+use starknet_patricia::db_layout::{NodeLayout, NodeLayoutFor, TrieType};
 use starknet_patricia::patricia_merkle_tree::filled_tree::node::{FactDbFilledNode, FilledNode};
 use starknet_patricia::patricia_merkle_tree::filled_tree::node_serde::FactNodeDeserializationContext;
 use starknet_patricia::patricia_merkle_tree::filled_tree::tree::FilledTree;
@@ -61,6 +61,18 @@ impl<'a, L: LeafWithEmptyKeyContext> NodeLayout<'a, L> for FactsNodeLayout {
     }
 }
 
+impl NodeLayoutFor<StarknetStorageValue> for FactsNodeLayout {
+    type DbLeaf = StarknetStorageValue;
+}
+
+impl NodeLayoutFor<ContractState> for FactsNodeLayout {
+    type DbLeaf = ContractState;
+}
+
+impl NodeLayoutFor<CompiledClassHash> for FactsNodeLayout {
+    type DbLeaf = CompiledClassHash;
+}
+
 pub struct FactsDb<S: Storage> {
     // TODO(Yoav): Define StorageStats trait and impl it here. Then, make the storage field
     // private.
@@ -92,7 +104,7 @@ impl<S: Storage> ForestReader<FactsDbInitialRead> for FactsDb<S> {
         forest_sorted_indices: &'a ForestSortedIndices<'a>,
         config: ReaderConfig,
     ) -> ForestResult<(OriginalSkeletonForest<'a>, HashMap<NodeIndex, ContractState>)> {
-        read_forest::<S, StarknetStorageValue, ContractState, CompiledClassHash, FactsNodeLayout>(
+        read_forest::<S, FactsNodeLayout>(
             &mut self.storage,
             context,
             storage_updates,

crates/starknet_committer/src/db/forest_trait.rs

@@ -3,8 +3,8 @@ use std::collections::HashMap;
 use async_trait::async_trait;
 use serde::{Deserialize, Serialize};
 use starknet_api::core::ContractAddress;
-use starknet_patricia::db_layout::NodeLayout;
-use starknet_patricia::patricia_merkle_tree::node_data::leaf::{Leaf, LeafModifications};
+use starknet_patricia::db_layout::NodeLayoutFor;
+use starknet_patricia::patricia_merkle_tree::node_data::leaf::LeafModifications;
 use starknet_patricia::patricia_merkle_tree::types::NodeIndex;
 use starknet_patricia_storage::errors::SerializationResult;
 use starknet_patricia_storage::storage_trait::{DbHashMap, DbKey, DbValue, Storage};
@@ -69,7 +69,7 @@ pub trait ForestReader<I: InputContext> {
 }
 
 /// Helper function containing layout-common read logic.
-pub(crate) async fn read_forest<'a, S, StorageLeaf, ContractStateLeaf, ClassesLeaf, Layout>(
+pub(crate) async fn read_forest<'a, S, Layout>(
     storage: &mut S,
     context: FactsDbInitialRead,
     storage_updates: &'a HashMap<ContractAddress, LeafModifications<StarknetStorageValue>>,
@@ -79,29 +79,25 @@ pub(crate) async fn read_forest<'a, S, StorageLeaf, ContractStateLeaf, ClassesLe
 ) -> ForestResult<(OriginalSkeletonForest<'a>, HashMap<NodeIndex, ContractState>)>
 where
     S: Storage,
-    ContractStateLeaf: Leaf + Into<ContractState>,
-    StorageLeaf: Leaf + From<StarknetStorageValue>,
-    ClassesLeaf: Leaf + From<CompiledClassHash>,
-    Layout: NodeLayout<'a, StorageLeaf>
-        + NodeLayout<'a, ContractStateLeaf>
-        + NodeLayout<'a, ClassesLeaf>,
+    Layout: NodeLayoutFor<StarknetStorageValue>
+        + NodeLayoutFor<ContractState>
+        + NodeLayoutFor<CompiledClassHash>,
 {
-    let (contracts_trie, original_contracts_trie_leaves) =
-        create_contracts_trie::<ContractStateLeaf, Layout>(
-            storage,
-            context.0.contracts_trie_root_hash,
-            forest_sorted_indices.contracts_trie_sorted_indices,
-        )
-        .await?;
-    let storage_tries = create_storage_tries::<StorageLeaf, Layout>(
+    let (contracts_trie, original_contracts_trie_leaves) = create_contracts_trie::<Layout>(
+        storage,
+        context.0.contracts_trie_root_hash,
+        forest_sorted_indices.contracts_trie_sorted_indices,
+    )
+    .await?;
+    let storage_tries = create_storage_tries::<Layout>(
         storage,
         storage_updates,
         &original_contracts_trie_leaves,
         &config,
         &forest_sorted_indices.storage_tries_sorted_indices,
     )
     .await?;
-    let classes_trie = create_classes_trie::<ClassesLeaf, Layout>(
+    let classes_trie = create_classes_trie::<Layout>(
         storage,
         classes_updates,
         context.0.classes_trie_root_hash,

crates/starknet_committer/src/db/index_db/db.rs

@@ -3,7 +3,7 @@ use std::collections::HashMap;
 use async_trait::async_trait;
 use starknet_api::core::ContractAddress;
 use starknet_api::hash::HashOutput;
-use starknet_patricia::db_layout::{NodeLayout, TrieType};
+use starknet_patricia::db_layout::{NodeLayout, NodeLayoutFor, TrieType};
 use starknet_patricia::patricia_merkle_tree::filled_tree::node::FilledNode;
 use starknet_patricia::patricia_merkle_tree::filled_tree::tree::FilledTree;
 use starknet_patricia::patricia_merkle_tree::node_data::leaf::{Leaf, LeafModifications};
@@ -76,6 +76,18 @@ where
     }
 }
 
+impl NodeLayoutFor<StarknetStorageValue> for IndexNodeLayout {
+    type DbLeaf = IndexLayoutStarknetStorageValue;
+}
+
+impl NodeLayoutFor<ContractState> for IndexNodeLayout {
+    type DbLeaf = IndexLayoutContractState;
+}
+
+impl NodeLayoutFor<CompiledClassHash> for IndexNodeLayout {
+    type DbLeaf = IndexLayoutCompiledClassHash;
+}
+
 // TODO(Ariel): define an IndexDbInitialRead empty type, and check whether each tree is empty inside
 // create_xxx_trie.
 #[async_trait]
@@ -91,13 +103,7 @@ impl<S: Storage> ForestReader<FactsDbInitialRead> for IndexDb<S> {
         forest_sorted_indices: &'a ForestSortedIndices<'a>,
         config: ReaderConfig,
     ) -> ForestResult<(OriginalSkeletonForest<'a>, HashMap<NodeIndex, ContractState>)> {
-        read_forest::<
-            S,
-            IndexLayoutStarknetStorageValue,
-            IndexLayoutContractState,
-            IndexLayoutCompiledClassHash,
-            IndexNodeLayout,
-        >(
+        read_forest::<S, IndexNodeLayout>(
             &mut self.storage,
             context,
             storage_updates,

crates/starknet_committer/src/db/index_db/leaves.rs

@@ -24,10 +24,12 @@ use crate::patricia_merkle_tree::types::CompiledClassHash;
 )]
 pub struct IndexLayoutContractState(pub ContractState);
 
-#[derive(Clone, Debug, Default, Eq, PartialEq, derive_more::AsRef, derive_more::From)]
+#[derive(
+    Clone, Debug, Default, Eq, PartialEq, derive_more::AsRef, derive_more::From, derive_more::Into,
+)]
 pub struct IndexLayoutCompiledClassHash(pub CompiledClassHash);
 
-#[derive(Clone, Debug, Default, Eq, PartialEq, derive_more::From)]
+#[derive(Clone, Debug, Default, Eq, PartialEq, derive_more::From, derive_more::Into)]
 pub struct IndexLayoutStarknetStorageValue(pub StarknetStorageValue);
 
 /// Set to 2^251 + 1 to avoid collisions with contract addresses prefixes.

crates/starknet_committer/src/db/trie_traversal.rs

@@ -4,7 +4,7 @@ use std::fmt::Debug;
 
 use starknet_api::core::ContractAddress;
 use starknet_api::hash::HashOutput;
-use starknet_patricia::db_layout::{NodeLayout, TrieType};
+use starknet_patricia::db_layout::{NodeLayout, NodeLayoutFor, TrieType};
 use starknet_patricia::patricia_merkle_tree::filled_tree::node::FilledNode;
 use starknet_patricia::patricia_merkle_tree::node_data::inner_node::{
     BinaryData,
@@ -341,11 +341,7 @@ pub async fn create_original_skeleton_tree<'a, L: Leaf, Layout: NodeLayout<'a, L
     Ok(skeleton_tree)
 }
 
-pub async fn create_storage_tries<
-    'a,
-    L: Leaf + From<StarknetStorageValue>,
-    Layout: NodeLayout<'a, L>,
->(
+pub async fn create_storage_tries<'a, Layout: NodeLayoutFor<StarknetStorageValue>>(
     storage: &mut impl Storage,
     actual_storage_updates: &HashMap<ContractAddress, LeafModifications<StarknetStorageValue>>,
     original_contracts_trie_leaves: &HashMap<NodeIndex, ContractState>,
@@ -364,14 +360,14 @@ pub async fn create_storage_tries<
             config.warn_on_trivial_modifications(),
         );
 
-        let original_skeleton = create_original_skeleton_tree::<L, Layout>(
+        let original_skeleton = create_original_skeleton_tree::<Layout::DbLeaf, Layout>(
             storage,
             contract_state.storage_root_hash,
             *sorted_leaf_indices,
             &config,
             // TODO(Ariel): Change `LeafModifications` in `actual_storage_updates` to be an
             // iterator over borrowed data so that the conversion below is costless.
-            &updates.iter().map(|(idx, value)| (*idx, L::from(*value))).collect(),
+            &updates.iter().map(|(idx, value)| (*idx, Layout::DbLeaf::from(*value))).collect(),
             None,
             &Layout::generate_key_context(TrieType::StorageTrie(*address)),
         )
@@ -383,15 +379,15 @@ pub async fn create_storage_tries<
 
 /// Creates the contracts trie original skeleton.
 /// Also returns the previous contracts state of the modified contracts.
-pub async fn create_contracts_trie<'a, L: Leaf + Into<ContractState>, Layout: NodeLayout<'a, L>>(
+pub async fn create_contracts_trie<'a, Layout: NodeLayoutFor<ContractState>>(
     storage: &mut impl Storage,
     contracts_trie_root_hash: HashOutput,
     contracts_trie_sorted_indices: SortedLeafIndices<'a>,
 ) -> ForestResult<(OriginalSkeletonTreeImpl<'a>, HashMap<NodeIndex, ContractState>)> {
     let config = OriginalSkeletonTrieConfig::new_for_contracts_trie();
 
     let mut leaves = HashMap::new();
-    let skeleton_tree = create_original_skeleton_tree::<L, Layout>(
+    let skeleton_tree = create_original_skeleton_tree::<Layout::DbLeaf, Layout>(
         storage,
         contracts_trie_root_hash,
         contracts_trie_sorted_indices,
@@ -408,11 +404,7 @@ pub async fn create_contracts_trie<'a, L: Leaf + Into<ContractState>, Layout: No
     Ok((skeleton_tree, leaves))
 }
 
-pub async fn create_classes_trie<
-    'a,
-    L: Leaf + From<CompiledClassHash>,
-    Layout: NodeLayout<'a, L>,
->(
+pub async fn create_classes_trie<'a, Layout: NodeLayoutFor<CompiledClassHash>>(
     storage: &mut impl Storage,
     actual_classes_updates: &LeafModifications<CompiledClassHash>,
     classes_trie_root_hash: HashOutput,
@@ -423,14 +415,17 @@ pub async fn create_classes_trie<
         config.warn_on_trivial_modifications(),
     );
 
-    Ok(create_original_skeleton_tree::<L, Layout>(
+    Ok(create_original_skeleton_tree::<Layout::DbLeaf, Layout>(
         storage,
         classes_trie_root_hash,
         contracts_trie_sorted_indices,
         &config,
         // TODO(Ariel): Change `actual_classes_updates` to be an iterator over borrowed data so
         // that the conversion below is costless.
-        &actual_classes_updates.iter().map(|(idx, value)| (*idx, L::from(*value))).collect(),
+        &actual_classes_updates
+            .iter()
+            .map(|(idx, value)| (*idx, Layout::DbLeaf::from(*value)))
+            .collect(),
         None,
         &Layout::generate_key_context(TrieType::ClassesTrie),
     )

crates/starknet_patricia/src/db_layout.rs

@@ -83,3 +83,17 @@ pub trait NodeLayout<'a, L: Leaf> {
         FilledNode { hash: filled_node.hash, data: node_data }
     }
 }
+
+/// A layout trait for one of the base leaf types: StarknetStorageValue, ContractState, or
+/// CompiledClassHash. Use `NodeLayout<DbLeaf>` when you only deal with the DB representation, use
+/// `NodeLayoutFor<BaseLeaf>` when you need to refer to both the logical type and the DB type.
+///
+/// We require that the BaseLeaf type and its DB representation `NodeLayout::DbLeaf` are isomorphic.
+/// The `From<BaseLeaf>` conversions are needed to convert incoming modifications, and the
+/// `Into<BaseLeaf>` conversions are used in the `ContractState` case where we need to return the
+/// leaves. For more details see `create_contracts_trie`, `create_classes_trie`, and
+/// `create_storage_tries` in `starknet_committer`.
+pub trait NodeLayoutFor<BaseLeaf>: for<'a> NodeLayout<'a, Self::DbLeaf> {
+    /// This layout's DB representation of BaseLeaf.
+    type DbLeaf: Leaf + From<BaseLeaf> + Into<BaseLeaf>;
+}