starknet_committer: add original tree creation test templates

crates/starknet_committer/Cargo.toml

@@ -19,7 +19,6 @@ hex.workspace = true
 pretty_assertions.workspace = true
 rand.workspace = true
 rand_distr.workspace = true
-rstest.workspace = true
 serde = { workspace = true, features = ["derive"] }
 serde_json.workspace = true
 starknet-types-core = { workspace = true, features = ["hash", "papyrus-serialization"] }
@@ -32,6 +31,8 @@ tokio = { workspace = true, features = ["rt"] }
 tracing.workspace = true
 
 [dev-dependencies]
+rstest.workspace = true
+rstest_reuse.workspace = true
 starknet_api = { workspace = true, features = ["testing"] }
 starknet_patricia = { workspace = true, features = ["testing"] }
 

crates/starknet_committer/src/db.rs

@@ -1,3 +1,5 @@
+#[cfg(test)]
+pub mod create_original_skeleton_tests;
 pub mod db_layout;
 #[cfg(any(feature = "testing", test))]
 pub mod external_test_utils;

crates/starknet_committer/src/db/create_original_skeleton_tests.rs

@@ -0,0 +1,283 @@
+use std::collections::HashMap;
+
+use rstest_reuse::template;
+use starknet_api::hash::HashOutput;
+use starknet_patricia::db_layout::NodeLayout;
+use starknet_patricia::patricia_merkle_tree::external_test_utils::{
+    create_binary_entry_from_u128,
+    create_binary_skeleton_node,
+    create_edge_entry_from_u128,
+    create_edge_skeleton_node,
+    create_expected_skeleton_nodes,
+    create_root_edge_entry,
+    create_unmodified_subtree_skeleton_node,
+    AdditionHash,
+    MockLeaf,
+};
+use starknet_patricia::patricia_merkle_tree::node_data::leaf::{Leaf, LeafModifications};
+use starknet_patricia::patricia_merkle_tree::original_skeleton_tree::node::OriginalSkeletonNode;
+use starknet_patricia::patricia_merkle_tree::types::{NodeIndex, SortedLeafIndices, SubTreeHeight};
+use starknet_patricia_storage::db_object::{DBObject, EmptyKeyContext, HasStaticPrefix};
+use starknet_patricia_storage::map_storage::MapStorage;
+use starknet_patricia_storage::storage_trait::{DbHashMap, DbKey, DbValue};
+use starknet_types_core::felt::Felt;
+
+use crate::db::trie_traversal::create_original_skeleton_tree;
+use crate::patricia_merkle_tree::tree::OriginalSkeletonTrieConfig;
+
+pub(crate) fn create_mock_leaf_entry(val: u128) -> (DbKey, DbValue) {
+    let leaf = MockLeaf(Felt::from(val));
+    (leaf.get_db_key(&EmptyKeyContext, &leaf.0.to_bytes_be()), leaf.serialize().unwrap())
+}
+
+fn create_mock_leaf_modifications(
+    leaf_modifications: Vec<(u128, u128)>,
+) -> LeafModifications<MockLeaf> {
+    leaf_modifications
+        .into_iter()
+        .map(|(idx, val)| (NodeIndex::from(idx), MockLeaf(Felt::from(val))))
+        .collect()
+}
+
+pub(crate) struct CreateTreeCase {
+    pub storage: MapStorage,
+    pub leaf_modifications: LeafModifications<MockLeaf>,
+    pub root_hash: HashOutput,
+    pub expected_skeleton_nodes: HashMap<NodeIndex, OriginalSkeletonNode>,
+    pub subtree_height: SubTreeHeight,
+}
+
+pub(crate) fn simple_tree_of_height_3() -> CreateTreeCase {
+    CreateTreeCase {
+        storage: MapStorage(DbHashMap::from([
+            create_root_edge_entry(50, SubTreeHeight::new(3)),
+            create_binary_entry_from_u128::<AdditionHash>(8, 9),
+            create_edge_entry_from_u128::<AdditionHash>(11, 1, 1),
+            create_binary_entry_from_u128::<AdditionHash>(17, 13),
+            create_edge_entry_from_u128::<AdditionHash>(15, 3, 2),
+            create_binary_entry_from_u128::<AdditionHash>(30, 20),
+            create_mock_leaf_entry(8),
+            create_mock_leaf_entry(9),
+            create_mock_leaf_entry(11),
+            create_mock_leaf_entry(15),
+        ])),
+        leaf_modifications: create_mock_leaf_modifications(vec![(8, 8), (10, 3), (13, 2)]),
+        root_hash: HashOutput(Felt::from(50_u128 + 248_u128)),
+        expected_skeleton_nodes: create_expected_skeleton_nodes(
+            vec![
+                create_binary_skeleton_node(1),
+                create_binary_skeleton_node(2),
+                create_edge_skeleton_node(3, 3, 2),
+                create_binary_skeleton_node(4),
+                create_edge_skeleton_node(5, 1, 1),
+                create_unmodified_subtree_skeleton_node(9, 9),
+                create_unmodified_subtree_skeleton_node(15, 15),
+                create_unmodified_subtree_skeleton_node(11, 11),
+            ],
+            3,
+        ),
+        subtree_height: SubTreeHeight::new(3),
+    }
+}
+
+pub(crate) fn another_simple_tree_of_height_3() -> CreateTreeCase {
+    CreateTreeCase {
+        storage: MapStorage(DbHashMap::from([
+            create_root_edge_entry(29, SubTreeHeight::new(3)),
+            create_binary_entry_from_u128::<AdditionHash>(10, 2),
+            create_edge_entry_from_u128::<AdditionHash>(3, 1, 1),
+            create_binary_entry_from_u128::<AdditionHash>(4, 7),
+            create_edge_entry_from_u128::<AdditionHash>(12, 0, 1),
+            create_binary_entry_from_u128::<AdditionHash>(5, 11),
+            create_binary_entry_from_u128::<AdditionHash>(13, 16),
+            create_mock_leaf_entry(10),
+            create_mock_leaf_entry(2),
+            create_mock_leaf_entry(3),
+            create_mock_leaf_entry(4),
+            create_mock_leaf_entry(7),
+        ])),
+        leaf_modifications: create_mock_leaf_modifications(vec![(8, 5), (11, 1), (13, 3)]),
+        root_hash: HashOutput(Felt::from(29_u128 + 248_u128)),
+        expected_skeleton_nodes: create_expected_skeleton_nodes(
+            vec![
+                create_binary_skeleton_node(1),
+                create_edge_skeleton_node(2, 0, 1),
+                create_binary_skeleton_node(3),
+                create_binary_skeleton_node(4),
+                create_edge_skeleton_node(6, 1, 1),
+                create_unmodified_subtree_skeleton_node(7, 11),
+                create_unmodified_subtree_skeleton_node(9, 2),
+            ],
+            3,
+        ),
+        subtree_height: SubTreeHeight::new(3),
+    }
+}
+
+pub(crate) fn tree_of_height_4_with_long_edge() -> CreateTreeCase {
+    CreateTreeCase {
+        storage: MapStorage(DbHashMap::from([
+            create_root_edge_entry(116, SubTreeHeight::new(4)),
+            create_binary_entry_from_u128::<AdditionHash>(11, 13),
+            create_edge_entry_from_u128::<AdditionHash>(5, 0, 1),
+            create_binary_entry_from_u128::<AdditionHash>(19, 40),
+            create_edge_entry_from_u128::<AdditionHash>(20, 3, 2),
+            create_binary_entry_from_u128::<AdditionHash>(6, 59),
+            create_edge_entry_from_u128::<AdditionHash>(24, 0, 2),
+            create_binary_entry_from_u128::<AdditionHash>(25, 65),
+            create_binary_entry_from_u128::<AdditionHash>(26, 90),
+            create_mock_leaf_entry(11),
+            create_mock_leaf_entry(13),
+            create_mock_leaf_entry(20),
+            create_mock_leaf_entry(5),
+            create_mock_leaf_entry(19),
+            create_mock_leaf_entry(40),
+        ])),
+        leaf_modifications: create_mock_leaf_modifications(vec![
+            (18, 5),
+            (25, 1),
+            (29, 15),
+            (30, 19),
+        ]),
+        root_hash: HashOutput(Felt::from(116_u128 + 247_u128)),
+        expected_skeleton_nodes: create_expected_skeleton_nodes(
+            vec![
+                create_binary_skeleton_node(1),
+                create_edge_skeleton_node(2, 0, 2),
+                create_binary_skeleton_node(3),
+                create_edge_skeleton_node(6, 3, 2),
+                create_binary_skeleton_node(7),
+                create_unmodified_subtree_skeleton_node(8, 24),
+                create_edge_skeleton_node(14, 0, 1),
+                create_binary_skeleton_node(15),
+                create_unmodified_subtree_skeleton_node(27, 20),
+                create_unmodified_subtree_skeleton_node(28, 5),
+                create_unmodified_subtree_skeleton_node(31, 40),
+            ],
+            4,
+        ),
+        subtree_height: SubTreeHeight::new(4),
+    }
+}
+
+pub(crate) async fn test_create_original_skeleton<L, Layout>(
+    storage: &mut MapStorage,
+    leaf_modifications: &LeafModifications<L>,
+    root_hash: HashOutput,
+    expected_skeleton_nodes: &HashMap<NodeIndex, OriginalSkeletonNode>,
+    subtree_height: SubTreeHeight,
+    compare_modified_leaves: bool,
+) where
+    L: Leaf + HasStaticPrefix<KeyContext = EmptyKeyContext>,
+    Layout: for<'a> NodeLayout<'a, L>,
+{
+    let leaf_modifications: LeafModifications<L> = leaf_modifications
+        .iter()
+        .map(|(idx, leaf)| (NodeIndex::from_subtree_index(*idx, subtree_height), leaf.clone()))
+        .collect();
+
+    let config = OriginalSkeletonTrieConfig::new_for_tests(compare_modified_leaves);
+
+    let mut sorted_leaf_indices: Vec<NodeIndex> = leaf_modifications.keys().copied().collect();
+    let sorted_leaf_indices = SortedLeafIndices::new(&mut sorted_leaf_indices);
+
+    let skeleton_tree = create_original_skeleton_tree::<L, Layout>(
+        storage,
+        root_hash,
+        sorted_leaf_indices,
+        &config,
+        &leaf_modifications,
+        None,
+        &EmptyKeyContext,
+    )
+    .await
+    .unwrap();
+
+    assert_eq!(&skeleton_tree.nodes, expected_skeleton_nodes);
+}
+
+#[template]
+#[rstest]
+// This test uses addition hash for simplicity (i.e hash(a,b) = a + b).
+///                 Old tree structure:
+///
+///                             50
+///                           /   \
+///                         30     20
+///                        /  \     \
+///                       17  13     *
+///                      /  \   \     \
+///                     8    9  11     15
+///
+///                   Modified leaves indices: [8, 10, 13]
+///
+///                   Expected skeleton:
+///
+///                             B
+///                           /   \
+///                          B     E
+///                         / \     \
+///                        B   E     *
+///                       / \   \     \
+///                      NZ  9  11    15
+
+#[case::simple_tree_of_height_3(
+    crate::db::create_original_skeleton_tests::simple_tree_of_height_3()
+)]
+///                 Old tree structure:
+///
+///                             29
+///                           /    \
+///                         13      16
+///                        /      /    \
+///                       12      5     11
+///                      /  \      \    /  \
+///                     10   2      3   4   7
+///
+///                   Modified leaves indices: [8, 11, 13]
+///
+///                   Expected skeleton:
+///
+///                             B
+///                           /   \
+///                         E      B
+///                        /     /    \
+///                       B      E     E
+///                      /  \     \     \
+///                     NZ   2     NZ    NZ
+#[case::another_simple_tree_of_height_3(
+    crate::db::create_original_skeleton_tests::another_simple_tree_of_height_3()
+)]
+///                  Old tree structure:
+///
+///                             116
+///                           /     \
+///                         26       90
+///                        /      /     \
+///                       /      25      65
+///                      /        \     /  \
+///                     24         \   6   59
+///                    /  \         \  /  /  \
+///                   11  13       20  5  19 40
+///
+///                   Modified leaves indices: [18, 25, 29, 30]
+///
+///                   Expected skeleton:
+///
+///                              B
+///                           /     \
+///                          E       B
+///                         /     /     \
+///                        /     E       B
+///                       /       \     /  \
+///                      24        \   E    B
+///                                 \  /     \
+///                                 20 5     40
+#[case::tree_of_height_4_with_long_edge(
+    crate::db::create_original_skeleton_tests::tree_of_height_4_with_long_edge()
+)]
+pub fn create_tree_cases(
+    #[case] mut case: CreateTreeCase,
+    #[values(true, false)] compare_modified_leaves: bool,
+) {
+}