Skip to content

Commit

Permalink
refactor: use mock test utils in filled tree test
Browse files Browse the repository at this point in the history
  • Loading branch information
AvivYossef-starkware committed Jul 22, 2024
1 parent f8fbe2f commit 031e6ae
Show file tree
Hide file tree
Showing 2 changed files with 50 additions and 103 deletions.
134 changes: 34 additions & 100 deletions crates/committer/src/patricia_merkle_tree/filled_tree/tree_test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,11 +6,11 @@ use crate::felt::Felt;
use crate::hash::hash_trait::HashOutput;
use crate::patricia_merkle_tree::filled_tree::node::FilledNode;
use crate::patricia_merkle_tree::filled_tree::tree::{FilledTree, FilledTreeImpl};
use crate::patricia_merkle_tree::internal_test_utils::{MockLeaf, OriginalSkeletonMockTrieConfig};
use crate::patricia_merkle_tree::node_data::inner_node::{
BinaryData, EdgeData, EdgePathLength, NodeData, PathToBottom,
};
use crate::patricia_merkle_tree::node_data::leaf::SkeletonLeaf;
use crate::patricia_merkle_tree::original_skeleton_tree::config::OriginalSkeletonStorageTrieConfig;
use crate::patricia_merkle_tree::original_skeleton_tree::tree::OriginalSkeletonTreeImpl;
use crate::patricia_merkle_tree::types::{NodeIndex, SortedLeafIndices};
use crate::patricia_merkle_tree::updated_skeleton_tree::hash_function::TreeHashFunctionImpl;
Expand All @@ -24,7 +24,7 @@ use crate::storage::map_storage::MapStorage;
/// This test is a sanity test for computing the root hash of the patricia merkle tree with a single node that is a leaf with hash==1.
async fn test_filled_tree_sanity() {
let mut skeleton_tree: UpdatedSkeletonNodeMap = HashMap::new();
let new_filled_leaf = StarknetStorageValue(Felt::ONE);
let new_filled_leaf = MockLeaf(Felt::ONE);
let new_leaf_index = NodeIndex::ROOT;
skeleton_tree.insert(new_leaf_index, UpdatedSkeletonNode::Leaf);
let modifications = HashMap::from([(new_leaf_index, new_filled_leaf)]);
Expand Down Expand Up @@ -84,7 +84,7 @@ async fn test_small_filled_tree() {
.map(|(index, value)| {
(
NodeIndex::from(*index),
StarknetStorageValue(Felt::from_hex(value).unwrap()),
MockLeaf(Felt::from_hex(value).unwrap()),
)
})
.collect();
Expand All @@ -100,53 +100,17 @@ async fn test_small_filled_tree() {
let root_hash = filled_tree.get_root_hash();

// The expected hash values were computed separately.
let expected_root_hash = HashOutput(
Felt::from_hex("0xe8899e8c731a35f5e9ce4c4bc32aabadcc81c5cdcc1aeba74fa7509046c338").unwrap(),
);
let expected_root_hash = HashOutput(Felt::from_hex("0x21").unwrap());
let expected_filled_tree_map = HashMap::from([
create_binary_entry_for_testing(
1,
"0xe8899e8c731a35f5e9ce4c4bc32aabadcc81c5cdcc1aeba74fa7509046c338",
"0x4e970ad06a06486b44fff5606c4f65486d31e05e323d65a618d4ef8cdf6d3a0",
"0x2955a96b09495fb2ce4ed65cf679c54e54aefc2c6972d7f3042590000bb7543",
),
create_edge_entry_for_testing(
2,
"0x4e970ad06a06486b44fff5606c4f65486d31e05e323d65a618d4ef8cdf6d3a0",
0,
1,
"0x5d36a1ae900ef417a5696417dde9a0244b873522f40b552e4a60acde0991bc9",
),
create_edge_entry_for_testing(
3,
"0x2955a96b09495fb2ce4ed65cf679c54e54aefc2c6972d7f3042590000bb7543",
15,
4,
"0x3",
),
create_binary_entry_for_testing(
4,
"0x5d36a1ae900ef417a5696417dde9a0244b873522f40b552e4a60acde0991bc9",
"0x582d984e4005c27b9c886cd00ec9a82ed5323aa629f6ea6b3ed7c0386ae6256",
"0x39eb7b85bcc9deac314406d6b73154b09b008f8af05e2f58ab623f4201d0b88",
),
create_edge_entry_for_testing(
8,
"0x582d984e4005c27b9c886cd00ec9a82ed5323aa629f6ea6b3ed7c0386ae6256",
3,
2,
"0x1",
),
create_edge_entry_for_testing(
9,
"0x39eb7b85bcc9deac314406d6b73154b09b008f8af05e2f58ab623f4201d0b88",
0,
2,
"0x2",
),
create_leaf_entry_for_testing(35, "0x1"),
create_leaf_entry_for_testing(36, "0x2"),
create_leaf_entry_for_testing(63, "0x3"),
create_mock_binary_entry_for_testing(1, "0x21", "0xb", "0x16"),
create_mock_edge_entry_for_testing(2, "0xb", 0, 1, "0xa"),
create_mock_edge_entry_for_testing(3, "0x16", 15, 4, "0x3"),
create_mock_binary_entry_for_testing(4, "0xa", "0x6", "0x4"),
create_mock_edge_entry_for_testing(8, "0x6", 3, 2, "0x1"),
create_mock_edge_entry_for_testing(9, "0x4", 0, 2, "0x2"),
create_mock_leaf_entry_for_testing(35, "0x1"),
create_mock_leaf_entry_for_testing(36, "0x2"),
create_mock_leaf_entry_for_testing(63, "0x3"),
]);
assert_eq!(filled_tree_map, &expected_filled_tree_map);
assert_eq!(root_hash, expected_root_hash, "Root hash mismatch");
Expand All @@ -158,12 +122,12 @@ async fn test_small_filled_tree() {
/// i=1: binary
/// / \
/// i=2: edge i=3: unmodified
/// l=1, p=0 hash=0x2955a96b09495fb2ce4ed65cf679c54e54aefc2c6972d7f3042590000bb7543
/// l=1, p=0 hash=0x3
/// /
/// i=4: binary
/// / \
/// i=8: edge i=9: unmodified
/// l=2, p=3 hash=0x39eb7b85bcc9deac314406d6b73154b09b008f8af05e2f58ab623f4201d0b88
/// l=2, p=3 hash=0x4
/// \
/// \
/// i=35: leaf
Expand All @@ -174,24 +138,18 @@ async fn test_small_tree_with_unmodified_nodes() {
let nodes_in_skeleton_tree = [
create_binary_updated_skeleton_node_for_testing(1),
create_path_to_bottom_edge_updated_skeleton_node_for_testing(2, 0, 1),
create_unmodified_updated_skeleton_node_for_testing(
3,
"0x2955a96b09495fb2ce4ed65cf679c54e54aefc2c6972d7f3042590000bb7543",
),
create_unmodified_updated_skeleton_node_for_testing(3, "0x3"),
create_binary_updated_skeleton_node_for_testing(4),
create_path_to_bottom_edge_updated_skeleton_node_for_testing(8, 3, 2),
create_unmodified_updated_skeleton_node_for_testing(
9,
"0x39eb7b85bcc9deac314406d6b73154b09b008f8af05e2f58ab623f4201d0b88",
),
create_unmodified_updated_skeleton_node_for_testing(9, "0x4"),
create_leaf_updated_skeleton_node_for_testing(new_leaf_index),
];
let skeleton_tree: UpdatedSkeletonNodeMap = nodes_in_skeleton_tree.into_iter().collect();

let updated_skeleton_tree = UpdatedSkeletonTreeImpl { skeleton_tree };
let modifications = HashMap::from([(
NodeIndex::from(new_leaf_index),
StarknetStorageValue(Felt::from_hex(new_leaf).unwrap()),
MockLeaf(Felt::from_hex(new_leaf).unwrap()),
)]);

// Compute the hash values.
Expand All @@ -207,37 +165,13 @@ async fn test_small_tree_with_unmodified_nodes() {
// The expected hash values were computed separately. Note that the unmodified nodes are not
// computed in the filled tree, but the hash values are directly used. The hashes of unmodified
// nodes should not appear in the filled tree.
let expected_root_hash = HashOutput(
Felt::from_hex("0xe8899e8c731a35f5e9ce4c4bc32aabadcc81c5cdcc1aeba74fa7509046c338").unwrap(),
);
let expected_root_hash = HashOutput(Felt::from_hex("0xe").unwrap());
let expected_filled_tree_map = HashMap::from([
create_binary_entry_for_testing(
1,
"0xe8899e8c731a35f5e9ce4c4bc32aabadcc81c5cdcc1aeba74fa7509046c338",
"0x4e970ad06a06486b44fff5606c4f65486d31e05e323d65a618d4ef8cdf6d3a0",
"0x2955a96b09495fb2ce4ed65cf679c54e54aefc2c6972d7f3042590000bb7543",
),
create_edge_entry_for_testing(
2,
"0x4e970ad06a06486b44fff5606c4f65486d31e05e323d65a618d4ef8cdf6d3a0",
0,
1,
"0x5d36a1ae900ef417a5696417dde9a0244b873522f40b552e4a60acde0991bc9",
),
create_binary_entry_for_testing(
4,
"0x5d36a1ae900ef417a5696417dde9a0244b873522f40b552e4a60acde0991bc9",
"0x582d984e4005c27b9c886cd00ec9a82ed5323aa629f6ea6b3ed7c0386ae6256",
"0x39eb7b85bcc9deac314406d6b73154b09b008f8af05e2f58ab623f4201d0b88",
),
create_edge_entry_for_testing(
8,
"0x582d984e4005c27b9c886cd00ec9a82ed5323aa629f6ea6b3ed7c0386ae6256",
3,
2,
"0x1",
),
create_leaf_entry_for_testing(35, "0x1"),
create_mock_binary_entry_for_testing(1, "0xe", "0xb", "0x3"),
create_mock_edge_entry_for_testing(2, "0b", 0, 1, "0xa"),
create_mock_binary_entry_for_testing(4, "0xa", "0x6", "0x4"),
create_mock_edge_entry_for_testing(8, "0x6", 3, 2, "0x1"),
create_mock_leaf_entry_for_testing(35, "0x1"),
]);
assert_eq!(filled_tree_map, &expected_filled_tree_map);
assert_eq!(root_hash, expected_root_hash, "Root hash mismatch");
Expand All @@ -246,8 +180,8 @@ async fn test_small_tree_with_unmodified_nodes() {
#[tokio::test(flavor = "multi_thread")]
/// Test that deleting a leaf that does not exist in the tree succeeds.
async fn test_delete_leaf_from_empty_tree() {
let storage_modifications: HashMap<NodeIndex, StarknetStorageValue> =
HashMap::from([(NodeIndex::FIRST_LEAF, StarknetStorageValue(Felt::ZERO))]);
let storage_modifications: HashMap<NodeIndex, MockLeaf> =
HashMap::from([(NodeIndex::FIRST_LEAF, MockLeaf(Felt::ZERO))]);

let mut indices = [NodeIndex::FIRST_LEAF];
// Create an empty original skeleton tree with a single leaf modified.
Expand All @@ -257,7 +191,7 @@ async fn test_delete_leaf_from_empty_tree() {
},
HashOutput::ROOT_OF_EMPTY_TREE,
SortedLeafIndices::new(&mut indices),
&OriginalSkeletonStorageTrieConfig::new(&storage_modifications, false),
&OriginalSkeletonMockTrieConfig::new(&storage_modifications, false),
)
.unwrap();

Expand Down Expand Up @@ -318,12 +252,12 @@ fn create_leaf_updated_skeleton_node_for_testing(index: u128) -> (NodeIndex, Upd
(NodeIndex::from(index), UpdatedSkeletonNode::Leaf)
}

fn create_binary_entry_for_testing(
fn create_mock_binary_entry_for_testing(
index: u128,
hash: &str,
left_hash: &str,
right_hash: &str,
) -> (NodeIndex, FilledNode<StarknetStorageValue>) {
) -> (NodeIndex, FilledNode<MockLeaf>) {
(
NodeIndex::from(index),
FilledNode {
Expand All @@ -336,13 +270,13 @@ fn create_binary_entry_for_testing(
)
}

fn create_edge_entry_for_testing(
fn create_mock_edge_entry_for_testing(
index: u128,
hash: &str,
path: u128,
length: u8,
bottom_hash: &str,
) -> (NodeIndex, FilledNode<StarknetStorageValue>) {
) -> (NodeIndex, FilledNode<MockLeaf>) {
(
NodeIndex::from(index),
FilledNode {
Expand All @@ -359,15 +293,15 @@ fn create_edge_entry_for_testing(
)
}

fn create_leaf_entry_for_testing(
fn create_mock_leaf_entry_for_testing(
index: u128,
hash: &str,
) -> (NodeIndex, FilledNode<StarknetStorageValue>) {
) -> (NodeIndex, FilledNode<MockLeaf>) {
(
NodeIndex::from(index),
FilledNode {
hash: HashOutput(Felt::from_hex(hash).unwrap()),
data: NodeData::Leaf(StarknetStorageValue(Felt::from_hex(hash).unwrap())),
data: NodeData::Leaf(MockLeaf(Felt::from_hex(hash).unwrap())),
},
)
}
19 changes: 16 additions & 3 deletions crates/committer/src/patricia_merkle_tree/internal_test_utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,17 +4,20 @@ use crate::felt::Felt;
use crate::hash::hash_trait::HashOutput;
use crate::patricia_merkle_tree::external_test_utils::get_random_u256;

use super::node_data::errors::LeafResult;
use super::node_data::leaf::{Leaf, LeafModifications};
use super::updated_skeleton_tree::hash_function::HashFunction;
use crate::generate_trie_config;
use crate::patricia_merkle_tree::node_data::errors::LeafResult;
use crate::patricia_merkle_tree::node_data::inner_node::NodeData;
use crate::patricia_merkle_tree::node_data::inner_node::{EdgePathLength, PathToBottom};
use crate::patricia_merkle_tree::node_data::leaf::SkeletonLeaf;
use crate::patricia_merkle_tree::node_data::leaf::{Leaf, LeafModifications};
use crate::patricia_merkle_tree::original_skeleton_tree::config::OriginalSkeletonTreeConfig;
use crate::patricia_merkle_tree::original_skeleton_tree::errors::OriginalSkeletonTreeError;
use crate::patricia_merkle_tree::original_skeleton_tree::node::OriginalSkeletonNode;
use crate::patricia_merkle_tree::original_skeleton_tree::tree::OriginalSkeletonTreeResult;
use crate::patricia_merkle_tree::types::{NodeIndex, SubTreeHeight};
use crate::patricia_merkle_tree::updated_skeleton_tree::hash_function::{
HashFunction, TreeHashFunction, TreeHashFunctionImpl,
};
use crate::patricia_merkle_tree::updated_skeleton_tree::node::UpdatedSkeletonNode;
use crate::patricia_merkle_tree::updated_skeleton_tree::tree::UpdatedSkeletonTreeImpl;
use crate::storage::db_object::{DBObject, Deserializable};
Expand Down Expand Up @@ -61,6 +64,16 @@ impl Leaf for MockLeaf {
}
}

impl TreeHashFunction<MockLeaf> for TreeHashFunctionImpl {
fn compute_leaf_hash(leaf_data: &MockLeaf) -> HashOutput {
HashOutput(leaf_data.0)
}

fn compute_node_hash(node_data: &NodeData<MockLeaf>) -> HashOutput {
Self::compute_node_hash_with_inner_hash_function::<MockHash>(node_data)
}
}

generate_trie_config!(OriginalSkeletonMockTrieConfig, MockLeaf);

struct MockHashFunction;
Expand Down

0 comments on commit 031e6ae

Please sign in to comment.