starknet_committer: add index layout nodes

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

@@ -1 +1,2 @@
 pub mod leaves;
+pub mod types;

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

@@ -0,0 +1,145 @@
+use ethnum::U256;
+use starknet_api::hash::HashOutput;
+use starknet_patricia::patricia_merkle_tree::filled_tree::node::FilledNode;
+use starknet_patricia::patricia_merkle_tree::node_data::inner_node::{
+    BinaryData,
+    EdgeData,
+    EdgePath,
+    EdgePathLength,
+    NodeData,
+    PathToBottom,
+};
+use starknet_patricia::patricia_merkle_tree::node_data::leaf::Leaf;
+use starknet_patricia::patricia_merkle_tree::traversal::{SubTreeTrait, UnmodifiedChildTraversal};
+use starknet_patricia::patricia_merkle_tree::types::{NodeIndex, SortedLeafIndices};
+use starknet_patricia::patricia_merkle_tree::updated_skeleton_tree::hash_function::TreeHashFunction;
+use starknet_patricia_storage::db_object::{
+    DBObject,
+    EmptyDeserializationContext,
+    HasStaticPrefix,
+};
+use starknet_patricia_storage::errors::DeserializationError;
+use starknet_patricia_storage::storage_trait::{DbKeyPrefix, DbValue};
+use starknet_types_core::felt::Felt;
+
+use crate::hash_function::hash::TreeHashFunctionImpl;
+
+pub(crate) const INDEX_LAYOUT_BINARY_BYTES: usize = 32;
+pub(crate) const SERIALIZE_HASH_BYTES: usize = 32;
+
+// TODO(Ariel): Remove this wrapper once DBObject is a local trait.
+pub struct IndexFilledNode<L: Leaf>(pub FilledNode<L, ()>);
+
+pub struct IndexNodeContext {
+    pub is_leaf: bool,
+}
+
+impl<L: Leaf> HasStaticPrefix for IndexFilledNode<L> {
+    type KeyContext = <L as HasStaticPrefix>::KeyContext;
+    fn get_static_prefix(key_context: &Self::KeyContext) -> DbKeyPrefix {
+        L::get_static_prefix(key_context)
+    }
+}
+
+impl<L> DBObject for IndexFilledNode<L>
+where
+    L: Leaf,
+    TreeHashFunctionImpl: TreeHashFunction<L>,
+{
+    type DeserializeContext = IndexNodeContext;
+    fn serialize(&self) -> DbValue {
+        match &self.0.data {
+            NodeData::Binary(_) => DbValue(self.0.hash.0.to_bytes_be().to_vec()),
+            NodeData::Edge(edge_data) => {
+                let mut raw_bytes = self.0.hash.0.to_bytes_be().to_vec();
+                let bit_len: u8 = edge_data.path_to_bottom.length.into();
+                let byte_len: usize = (bit_len.saturating_add(7) / 8).into();
+                raw_bytes.push(bit_len);
+                raw_bytes
+                    .extend(edge_data.path_to_bottom.path.0.to_le_bytes()[..byte_len].to_vec());
+                DbValue(raw_bytes)
+            }
+            NodeData::Leaf(leaf) => leaf.serialize(),
+        }
+    }
+    fn deserialize(
+        value: &DbValue,
+        deserialize_context: &Self::DeserializeContext,
+    ) -> Result<Self, DeserializationError> {
+        if deserialize_context.is_leaf {
+            let leaf = L::deserialize(value, &EmptyDeserializationContext)?;
+            let hash = TreeHashFunctionImpl::compute_leaf_hash(&leaf);
+            Ok(IndexFilledNode(FilledNode { hash, data: NodeData::Leaf(leaf) }))
+        } else if value.0.len() == INDEX_LAYOUT_BINARY_BYTES {
+            Ok(IndexFilledNode(FilledNode {
+                hash: HashOutput(Felt::from_bytes_be_slice(&value.0)),
+                data: NodeData::Binary(BinaryData::<()> { left_data: (), right_data: () }),
+            }))
+        } else {
+            let node_hash = HashOutput(Felt::from_bytes_be_slice(&value.0[..SERIALIZE_HASH_BYTES]));
+            let value = &value.0[SERIALIZE_HASH_BYTES..];
+            let bit_len = value[0];
+            let byte_len = (bit_len.saturating_add(7) / 8).into();
+            let path = U256::from_le_bytes(
+                value[1..byte_len].try_into().expect("Slice with incorrect length."),
+            );
+            Ok(IndexFilledNode(FilledNode {
+                hash: node_hash,
+                data: NodeData::Edge(EdgeData::<()> {
+                    bottom_data: (),
+                    path_to_bottom: PathToBottom::new(
+                        EdgePath(path),
+                        EdgePathLength::new(bit_len)
+                            .map_err(|error| DeserializationError::ValueError(Box::new(error)))?,
+                    )
+                    .map_err(|error| DeserializationError::ValueError(Box::new(error)))?,
+                }),
+            }))
+        }
+    }
+}
+
+pub struct IndexLayoutSubTree<'a> {
+    pub sorted_leaf_indices: SortedLeafIndices<'a>,
+    pub root_index: NodeIndex,
+}
+
+impl<'a> SubTreeTrait<'a> for IndexLayoutSubTree<'a> {
+    type NodeData = ();
+    type NodeDeserializeContext = IndexNodeContext;
+
+    fn create(
+        sorted_leaf_indices: SortedLeafIndices<'a>,
+        root_index: NodeIndex,
+        _child_data: Self::NodeData,
+    ) -> Self {
+        Self { sorted_leaf_indices, root_index }
+    }
+
+    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_child(_data: Self::NodeData) -> UnmodifiedChildTraversal {
+        UnmodifiedChildTraversal::Traverse
+    }
+
+    fn get_root_context(&self) -> Self::NodeDeserializeContext {
+        Self::NodeDeserializeContext { is_leaf: self.is_leaf() }
+    }
+
+    fn get_root_prefix<L: Leaf>(
+        &self,
+        _key_context: &<L as HasStaticPrefix>::KeyContext,
+    ) -> DbKeyPrefix {
+        L::get_static_prefix(_key_context)
+    }
+
+    fn get_root_suffix(&self) -> Vec<u8> {
+        self.root_index.0.to_be_bytes().to_vec()
+    }
+}

crates/starknet_committer/src/hash_function/hash.rs

@@ -1,3 +1,5 @@
+use std::convert::AsRef;
+
 use starknet_api::hash::HashOutput;
 use starknet_patricia::patricia_merkle_tree::node_data::inner_node::NodeData;
 use starknet_patricia::patricia_merkle_tree::updated_skeleton_tree::hash_function::{
@@ -8,9 +10,13 @@ use starknet_types_core::felt::Felt;
 use starknet_types_core::hash::{Pedersen, Poseidon, StarkHash};
 
 use crate::block_committer::input::StarknetStorageValue;
+use crate::db::index_db::leaves::{
+    IndexLayoutCompiledClassHash,
+    IndexLayoutContractState,
+    IndexLayoutStarknetStorageValue,
+};
 use crate::patricia_merkle_tree::leaf::leaf_impl::ContractState;
 use crate::patricia_merkle_tree::types::CompiledClassHash;
-
 /// Implementation of HashFunction for Pedersen hash function.
 pub struct PedersenHashFunction;
 impl HashFunction for PedersenHashFunction {
@@ -43,35 +49,69 @@ impl TreeHashFunctionImpl {
 /// <https://docs.starknet.io/documentation/architecture_and_concepts/Network_Architecture/starknet-state/#trie_construction>
 impl TreeHashFunction<ContractState> for TreeHashFunctionImpl {
     fn compute_leaf_hash(contract_state: &ContractState) -> HashOutput {
-        HashOutput(Pedersen::hash(
-            &Pedersen::hash(
-                &Pedersen::hash(&contract_state.class_hash.0, &contract_state.storage_root_hash.0),
-                &contract_state.nonce.0,
-            ),
-            &Self::CONTRACT_STATE_HASH_VERSION,
-        ))
+        compute_contract_state_leaf_hash(contract_state)
     }
     fn compute_node_hash(node_data: &NodeData<ContractState, HashOutput>) -> HashOutput {
         Self::compute_node_hash_with_inner_hash_function::<PedersenHashFunction>(node_data)
     }
 }
 
+impl TreeHashFunction<IndexLayoutContractState> for TreeHashFunctionImpl {
+    fn compute_leaf_hash(contract_state: &IndexLayoutContractState) -> HashOutput {
+        compute_contract_state_leaf_hash(contract_state)
+    }
+    fn compute_node_hash(node_data: &NodeData<IndexLayoutContractState, HashOutput>) -> HashOutput {
+        Self::compute_node_hash_with_inner_hash_function::<PedersenHashFunction>(node_data)
+    }
+}
+
+fn compute_contract_state_leaf_hash<L: AsRef<ContractState>>(
+    contract_state_leaf: &L,
+) -> HashOutput {
+    let contract_state: &ContractState = contract_state_leaf.as_ref();
+    HashOutput(Pedersen::hash(
+        &Pedersen::hash(
+            &Pedersen::hash(&contract_state.class_hash.0, &contract_state.storage_root_hash.0),
+            &contract_state.nonce.0,
+        ),
+        &TreeHashFunctionImpl::CONTRACT_STATE_HASH_VERSION,
+    ))
+}
+
 /// Implementation of TreeHashFunction for the classes trie.
 /// The implementation is based on the following reference:
 /// <https://docs.starknet.io/documentation/architecture_and_concepts/Network_Architecture/starknet-state/#trie_construction>
 impl TreeHashFunction<CompiledClassHash> for TreeHashFunctionImpl {
     fn compute_leaf_hash(compiled_class_hash: &CompiledClassHash) -> HashOutput {
-        let contract_class_leaf_version: Felt = Felt::from_hex(Self::CONTRACT_CLASS_LEAF_V0)
-            .expect(
-                "could not parse hex string corresponding to b'CONTRACT_CLASS_LEAF_V0' to Felt",
-            );
-        HashOutput(Poseidon::hash(&contract_class_leaf_version, &compiled_class_hash.0))
+        compute_compiled_class_leaf_hash(compiled_class_hash)
     }
     fn compute_node_hash(node_data: &NodeData<CompiledClassHash, HashOutput>) -> HashOutput {
         Self::compute_node_hash_with_inner_hash_function::<PoseidonHashFunction>(node_data)
     }
 }
 
+impl TreeHashFunction<IndexLayoutCompiledClassHash> for TreeHashFunctionImpl {
+    fn compute_leaf_hash(compiled_class_hash: &IndexLayoutCompiledClassHash) -> HashOutput {
+        compute_compiled_class_leaf_hash(compiled_class_hash)
+    }
+    fn compute_node_hash(
+        node_data: &NodeData<IndexLayoutCompiledClassHash, HashOutput>,
+    ) -> HashOutput {
+        Self::compute_node_hash_with_inner_hash_function::<PoseidonHashFunction>(node_data)
+    }
+}
+
+fn compute_compiled_class_leaf_hash<L: AsRef<CompiledClassHash>>(
+    compiled_class_hash_leaf: &L,
+) -> HashOutput {
+    let compiled_class_hash: &CompiledClassHash = compiled_class_hash_leaf.as_ref();
+    let contract_class_leaf_version: Felt = Felt::from_hex(
+        TreeHashFunctionImpl::CONTRACT_CLASS_LEAF_V0,
+    )
+    .expect("could not parse hex string corresponding to b'CONTRACT_CLASS_LEAF_V0' to Felt");
+    HashOutput(Poseidon::hash(&contract_class_leaf_version, &compiled_class_hash.0))
+}
+
 /// Implementation of TreeHashFunction for the storage trie.
 /// The implementation is based on the following reference:
 /// <https://docs.starknet.io/documentation/architecture_and_concepts/Network_Architecture/starknet-state/#trie_construction>
@@ -84,6 +124,17 @@ impl TreeHashFunction<StarknetStorageValue> for TreeHashFunctionImpl {
     }
 }
 
+impl TreeHashFunction<IndexLayoutStarknetStorageValue> for TreeHashFunctionImpl {
+    fn compute_leaf_hash(storage_value: &IndexLayoutStarknetStorageValue) -> HashOutput {
+        HashOutput(storage_value.0.0)
+    }
+    fn compute_node_hash(
+        node_data: &NodeData<IndexLayoutStarknetStorageValue, HashOutput>,
+    ) -> HashOutput {
+        Self::compute_node_hash_with_inner_hash_function::<PedersenHashFunction>(node_data)
+    }
+}
+
 /// Combined trait for all specific implementations.
 pub(crate) trait ForestHashFunction:
     TreeHashFunction<ContractState>