build: tree hash function for different leaf types (#60)

* build: tree hash function for different leaf types * test: tree hash function for different leaf types
starkware-libs · Apr 24, 2024 · 758a47a · 758a47a
1 parent 54a8717
commit 758a47a
Show file tree

Hide file tree

Showing 2 changed files with 64 additions and 11 deletions.
diff --git a/crates/committer/src/patricia_merkle_tree/types.rs b/crates/committer/src/patricia_merkle_tree/types.rs
@@ -20,9 +20,10 @@ pub(crate) struct TreeHashFunctionImpl<H: HashFunction> {
     _hash_function: PhantomData<H>,
 }
 
-/// Implementation of TreeHashFunction.
-// TODO(Aner, 11/4/25): Implement the function for LeafData::StorageValue and LeafData::StateTreeTuple
+/// Implementation of TreeHashFunction. The implementation is based on the following reference:
+/// https://docs.starknet.io/documentation/architecture_and_concepts/Network_Architecture/starknet-state/#trie_construction
 // TODO(Aner, 11/4/24): Verify the correctness of the implementation.
+const CONTRACT_STATE_HASH_VERSION: Felt = Felt::ZERO;
 impl<H: HashFunction> TreeHashFunction<LeafData, H> for TreeHashFunctionImpl<H> {
     fn compute_node_hash(node_data: &NodeData<LeafData>) -> HashOutput {
         match node_data {
@@ -36,15 +37,22 @@ impl<H: HashFunction> TreeHashFunction<LeafData, H> for TreeHashFunctionImpl<H>
             }) => HashOutput(
                 H::compute_hash(HashInputPair(hash_output.0, path.0)).0 + Felt::from(length.0),
             ),
-            NodeData::Leaf(leaf_data) => match leaf_data {
-                LeafData::StorageValue(_) => todo!(),
-                LeafData::CompiledClassHash(compiled_class_hash) => {
-                    HashOutput(compiled_class_hash.0)
-                }
-                LeafData::StateTreeTuple { .. } => {
-                    todo!()
-                }
-            },
+            NodeData::Leaf(LeafData::StorageValue(storage_value)) => HashOutput(*storage_value),
+            NodeData::Leaf(LeafData::CompiledClassHash(compiled_class_hash)) => {
+                HashOutput(compiled_class_hash.0)
+            }
+            NodeData::Leaf(LeafData::StateTreeTuple {
+                class_hash,
+                contract_state_root_hash,
+                nonce,
+            }) => H::compute_hash(HashInputPair(
+                H::compute_hash(HashInputPair(
+                    H::compute_hash(HashInputPair(class_hash.0, *contract_state_root_hash)).0,
+                    nonce.0,
+                ))
+                .0,
+                CONTRACT_STATE_HASH_VERSION,
+            )),
         }
     }
 }

diff --git a/crates/committer/src/patricia_merkle_tree/types_test.rs b/crates/committer/src/patricia_merkle_tree/types_test.rs
@@ -1,8 +1,11 @@
 use crate::hash::hash_trait::{HashFunction, HashInputPair, HashOutput};
 use crate::hash::pedersen::PedersenHashFunction;
+use crate::patricia_merkle_tree::filled_tree::node::ClassHash;
+use crate::patricia_merkle_tree::filled_tree::node::Nonce;
 use crate::patricia_merkle_tree::node_data::inner_node::{
     BinaryData, EdgeData, EdgePath, EdgePathLength, NodeData, PathToBottom,
 };
+use crate::patricia_merkle_tree::node_data::leaf::LeafData;
 use crate::patricia_merkle_tree::types::TreeHashFunction;
 use crate::patricia_merkle_tree::types::{NodeIndex, TreeHashFunctionImpl};
 use crate::types::Felt;
@@ -77,3 +80,45 @@ fn test_tree_hash_function_impl_edge_node(
     assert_eq!(hash_output, HashOutput(expected_hash));
     assert_eq!(hash_output, direct_hash_computation);
 }
+
+#[rstest]
+#[case(NodeData::Leaf(LeafData::CompiledClassHash(ClassHash(Felt::from_hex("0xACDC").unwrap()))), Felt::from_hex("0xACDC").unwrap())]
+#[case(NodeData::Leaf(LeafData::StorageValue(Felt::from_hex("0xDEAFBEEF").unwrap())), Felt::from_hex("0xDEAFBEEF").unwrap())]
+// Random StateTreeTuples and the expected hash results were generated and computed elsewhere.
+#[case(
+    NodeData::Leaf(
+        LeafData::StateTreeTuple{
+            class_hash: ClassHash(Felt::from_hex("0x150917f3bba17e3c0be685981d6cf8874098a857f354d374688e76eb69f44ab").unwrap()), 
+            contract_state_root_hash: Felt::from_hex("0x5d6b3a46ac855f3cea9af60e13cc6c81bf237b901f49f73981a521d4cb6c50c").unwrap(), 
+            nonce: Nonce(Felt::from_hex("0x38").unwrap()) 
+        }
+    ),
+    Felt::from_hex("0x3f27688d56740e5e238acde5b408154563c4f6e05514b8029e86ad51c388f8b").unwrap()
+)]
+#[case(
+    NodeData::Leaf(
+        LeafData::StateTreeTuple{
+            class_hash: ClassHash(Felt::from_hex("0x2c9982b9bd36f16e409c98616e43dbc9e4b47db8a8e8edc3b915bc4dab3c61c").unwrap()), 
+            contract_state_root_hash: Felt::from_hex("0x63d7c2a04df01e361238cf4bf07b7cee2e3c5ee58b4ec37cf279aed9ac8d013").unwrap(), 
+            nonce: Nonce(Felt::from_hex("0x1b").unwrap()) 
+        }
+    ),
+    Felt::from_hex("0x26e95377fcfa70a9882dc08653b7dc77165befe23971aa2fa5b926ff81af6cc").unwrap()
+)]
+#[case(
+    NodeData::Leaf(
+        LeafData::StateTreeTuple{
+            class_hash: ClassHash(Felt::from_hex("0x314bd50b446c4e8c3a82b0eb9326e532432fbd925ab7afea236f7f8618848a2").unwrap()), 
+            contract_state_root_hash: Felt::from_hex("0x1118159298289b372808e6b008b6bc9d68f5ebc060962ea1331fdc70cf3d047").unwrap(), 
+            nonce: Nonce(Felt::from_hex("0x47").unwrap()) 
+        }
+    ),
+    Felt::from_hex("0x1b20bbb35009bf03f86fb092b56a9c44deedbcca6addf8f7640f54a48ba5bbc").unwrap()
+)]
+fn test_tree_hash_function_impl_leaf_node(
+    #[case] node_data: NodeData<LeafData>,
+    #[case] expected_hash: Felt,
+) {
+    let hash_output = TreeHashFunctionImpl::<PedersenHashFunction>::compute_node_hash(&node_data);
+    assert_eq!(hash_output, HashOutput(expected_hash));
+}