index layout forest creation tests

ArielElp · ArielElp · commit 437648a1f19f · 2025-12-31T16:31:21.000+02:00
diff --git a/crates/starknet_committer/src/db/index_db/create_index_tree_test.rs b/crates/starknet_committer/src/db/index_db/create_index_tree_test.rs
@@ -35,6 +35,7 @@ pub async fn convert_facts_db_to_index_db<FactsLeaf, IndexLeaf>(
     storage: &mut MapStorage,
     root_hash: HashOutput,
     trie_type: Option<TrieType>,
+    current_leaves: &mut Option<Vec<(NodeIndex, FactsLeaf)>>,
 ) -> MapStorage
 where
     FactsLeaf: Leaf + Into<IndexLeaf> + HasStaticPrefix<KeyContext = EmptyKeyContext>,
@@ -49,6 +50,7 @@ where
         NodeIndex::ROOT,
         root_hash,
         &trie_type,
+        current_leaves,
     )
     .await;
     index_layout_storage
@@ -61,11 +63,16 @@ async fn convert_facts_to_index_layout_inner<FactsLeaf, IndexLeaf>(
     current_root_index: NodeIndex,
     current_root_hash: HashOutput,
     trie_type: &TrieType,
+    current_leaves: &mut Option<Vec<(NodeIndex, FactsLeaf)>>,
 ) where
     FactsLeaf: Leaf + Into<IndexLeaf> + HasStaticPrefix<KeyContext = EmptyKeyContext>,
     IndexLeaf: Leaf + HasStaticPrefix<KeyContext = TrieType>,
     TreeHashFunctionImpl: TreeHashFunction<IndexLeaf>,
 {
+    println!("index: {:?}, hash: {:?}", current_root_index, current_root_hash);
+    if current_root_hash == HashOutput::ROOT_OF_EMPTY_TREE {
+        return;
+    }
     let db_prefix = if current_root_index.is_leaf() {
         FactsLeaf::get_static_prefix(&EmptyKeyContext)
     } else {
@@ -92,6 +99,7 @@ async fn convert_facts_to_index_layout_inner<FactsLeaf, IndexLeaf>(
                 children_indices[0],
                 binary_data.left_data,
                 trie_type,
+                current_leaves,
             )
             .await;
             convert_facts_to_index_layout_inner::<FactsLeaf, IndexLeaf>(
@@ -100,6 +108,7 @@ async fn convert_facts_to_index_layout_inner<FactsLeaf, IndexLeaf>(
                 children_indices[1],
                 binary_data.right_data,
                 trie_type,
+                current_leaves,
             )
             .await;
             IndexFilledNode(FilledNode {
@@ -116,6 +125,7 @@ async fn convert_facts_to_index_layout_inner<FactsLeaf, IndexLeaf>(
                 bottom_index,
                 edge_data.bottom_data,
                 trie_type,
+                current_leaves,
             )
             .await;
             IndexFilledNode(FilledNode {
@@ -126,10 +136,16 @@ async fn convert_facts_to_index_layout_inner<FactsLeaf, IndexLeaf>(
                 }),
             })
         }
-        NodeData::Leaf(leaf) => IndexFilledNode(FilledNode {
-            hash: current_root_hash,
-            data: NodeData::Leaf(leaf.into()),
-        }),
+        NodeData::Leaf(leaf) => {
+            if let Some(leaves) = current_leaves {
+                leaves.push((current_root_index, leaf.clone()));
+            }
+
+            IndexFilledNode(FilledNode {
+                hash: current_root_hash,
+                data: NodeData::Leaf(leaf.into()),
+            })
+        }
     };
 
     let index_db_key =
@@ -158,6 +174,7 @@ async fn test_create_tree_index_layout(
         &mut case.storage,
         case.root_hash,
         Some(TrieType::StorageTrie(ContractAddress::from(0_u128))),
+        &mut None,
     )
     .await;
 
diff --git a/crates/starknet_committer/src/forest/skeleton_forest_test.rs b/crates/starknet_committer/src/forest/skeleton_forest_test.rs
@@ -2,8 +2,10 @@ use std::collections::HashMap;
 
 use pretty_assertions::assert_eq;
 use rstest::rstest;
+use rstest_reuse::{apply, template};
 use starknet_api::core::{ClassHash, ContractAddress, Nonce};
 use starknet_api::hash::{HashOutput, StateRoots};
+use starknet_patricia::db_layout::TrieType;
 use starknet_patricia::patricia_merkle_tree::external_test_utils::{
     create_32_bytes_entry,
     create_binary_entry_from_u128,
@@ -21,11 +23,14 @@ use starknet_patricia_storage::db_object::{DBObject, EmptyKeyContext};
 use starknet_patricia_storage::map_storage::MapStorage;
 use starknet_patricia_storage::storage_trait::{DbHashMap, DbKey, DbValue};
 use starknet_types_core::felt::Felt;
+use starknet_types_core::hash::StarkHash;
 
 use crate::block_committer::commit::get_all_modified_indices;
 use crate::block_committer::input::{
     contract_address_into_node_index,
+    try_node_index_into_contract_address,
     Input,
+    InputContext,
     ReaderConfig,
     StarknetStorageKey,
     StarknetStorageValue,
@@ -34,6 +39,13 @@ use crate::block_committer::input::{
 use crate::db::facts_db::db::FactsDb;
 use crate::db::facts_db::types::FactsDbInitialRead;
 use crate::db::forest_trait::ForestReader;
+use crate::db::index_db::create_index_tree_test::convert_facts_db_to_index_db;
+use crate::db::index_db::db::IndexDb;
+use crate::db::index_db::leaves::{
+    IndexLayoutCompiledClassHash,
+    IndexLayoutContractState,
+    IndexLayoutStarknetStorageValue,
+};
 use crate::forest::original_skeleton_forest::{ForestSortedIndices, OriginalSkeletonForest};
 use crate::patricia_merkle_tree::leaf::leaf_impl::ContractState;
 use crate::patricia_merkle_tree::types::CompiledClassHash;
@@ -58,14 +70,24 @@ pub(crate) fn create_compiled_class_leaf_entry(val: u128) -> (DbKey, DbValue) {
     (leaf.get_db_key(&EmptyKeyContext, &leaf.0.to_bytes_be()), leaf.serialize().unwrap())
 }
 
-pub(crate) fn create_contract_state_leaf_entry(val: u128) -> (DbKey, DbValue) {
-    let felt = Felt::from(val);
+pub(crate) fn create_contract_state_leaf_entry(
+    nonce: u128,
+    storage_root: u128,
+    class_hash: u128,
+    hash: Option<u128>,
+) -> (DbKey, DbValue) {
+    let (nonce, storage_root, class_hash) =
+        (Felt::from(nonce), Felt::from(storage_root), Felt::from(class_hash));
     let leaf = ContractState {
-        nonce: Nonce(felt),
-        storage_root_hash: HashOutput(felt),
-        class_hash: ClassHash(felt),
+        nonce: Nonce(nonce),
+        storage_root_hash: HashOutput(storage_root),
+        class_hash: ClassHash(class_hash),
     };
-    (leaf.get_db_key(&EmptyKeyContext, &felt.to_bytes_be()), leaf.serialize().unwrap())
+    let mut hash = hash.map(Felt::from);
+    if hash.is_none() {
+        hash = Some(AdditionHash::hash_array(&[class_hash, storage_root, nonce]));
+    }
+    (leaf.get_db_key(&EmptyKeyContext, &hash.unwrap().to_bytes_be()), leaf.serialize().unwrap())
 }
 
 // This test uses addition hash for simplicity (i.e hash(a,b) = a + b).
@@ -130,7 +152,7 @@ pub(crate) fn create_contract_state_leaf_entry(val: u128) -> (DbKey, DbValue) {
 ///       /  \     \     \                             / \   \     \
 ///      NZ   2     NZ    NZ                          NZ  9  16    15
 
-#[tokio::test]
+#[template]
 #[rstest]
 #[case(
     Input {
@@ -162,9 +184,11 @@ pub(crate) fn create_contract_state_leaf_entry(val: u128) -> (DbKey, DbValue) {
         create_edge_entry_from_u128::<AdditionHash>(554, 1, 1),
         create_binary_entry_from_u128::<AdditionHash>(305, 556),
         // Contracts trie leaves.
-        create_contract_state_leaf_entry(277),
-        create_contract_state_leaf_entry(303),
-        create_contract_state_leaf_entry(1),
+        // TODO(Ariel): Why both layouts pass when we lie about the hash?
+        // probably becaues the original skeleton doesn't care about these hashes (there were modifications).
+        create_contract_state_leaf_entry(277, 277, 277, Some(277)),
+        create_contract_state_leaf_entry(303, 303, 303, Some(303)),
+        create_contract_state_leaf_entry(0, 0, 1, None),
         // Classes trie inner nodes.
         create_binary_entry_from_u128::<AdditionHash>(33, 47),
         create_edge_entry_from_u128::<AdditionHash>(72, 1, 1),
@@ -295,17 +319,85 @@ pub(crate) fn create_contract_state_leaf_entry(val: u128) -> (DbKey, DbValue) {
         vec![6, 7, 0],
         vec![7, 6, 0],
 )]
-async fn test_create_original_skeleton_forest(
+fn create_original_skeleton_forest_cases(
+    #[case] input: Input<FactsDbInitialRead>,
+    #[case] storage: MapStorage,
+    #[case] expected_forest: OriginalSkeletonForest<'_>,
+    #[case] expected_original_contracts_trie_leaves: HashMap<ContractAddress, ContractState>,
+    #[case] expected_storage_tries_sorted_indices: HashMap<u128, Vec<u128>>,
+    #[case] expected_contracts_trie_sorted_indices: Vec<u128>,
+    #[case] expected_classes_trie_sorted_indices: Vec<u128>,
+) {
+}
+
+#[apply(create_original_skeleton_forest_cases)]
+#[rstest]
+#[tokio::test]
+async fn test_create_original_skeleton_forest_facts_layout(
     #[case] input: Input<FactsDbInitialRead>,
     #[case] storage: MapStorage,
     #[case] expected_forest: OriginalSkeletonForest<'_>,
     #[case] expected_original_contracts_trie_leaves: HashMap<ContractAddress, ContractState>,
     #[case] expected_storage_tries_sorted_indices: HashMap<u128, Vec<u128>>,
     #[case] expected_contracts_trie_sorted_indices: Vec<u128>,
     #[case] expected_classes_trie_sorted_indices: Vec<u128>,
+) {
+    test_create_original_skeleton_forest::<FactsDbInitialRead, FactsDb<MapStorage>>(
+        input,
+        storage,
+        expected_forest,
+        expected_original_contracts_trie_leaves,
+        expected_storage_tries_sorted_indices,
+        expected_contracts_trie_sorted_indices,
+        expected_classes_trie_sorted_indices,
+        |storage| FactsDb::new(storage),
+    )
+    .await;
+}
+
+#[apply(create_original_skeleton_forest_cases)]
+#[rstest]
+#[tokio::test]
+async fn test_create_original_skeleton_forest_index_layout(
+    #[case] input: Input<FactsDbInitialRead>,
+    #[case] mut storage: MapStorage,
+    #[case] expected_forest: OriginalSkeletonForest<'_>,
+    #[case] expected_original_contracts_trie_leaves: HashMap<ContractAddress, ContractState>,
+    #[case] expected_storage_tries_sorted_indices: HashMap<u128, Vec<u128>>,
+    #[case] expected_contracts_trie_sorted_indices: Vec<u128>,
+    #[case] expected_classes_trie_sorted_indices: Vec<u128>,
+) {
+    let index_storage =
+        convert_facts_forest_db_to_index_db(&mut storage, input.initial_read_context.0).await;
+    test_create_original_skeleton_forest::<FactsDbInitialRead, IndexDb<MapStorage>>(
+        input,
+        index_storage,
+        expected_forest,
+        expected_original_contracts_trie_leaves,
+        expected_storage_tries_sorted_indices,
+        expected_contracts_trie_sorted_indices,
+        expected_classes_trie_sorted_indices,
+        |storage| IndexDb::new(storage),
+    )
+    .await;
+}
+
+async fn test_create_original_skeleton_forest<
+    ReaderInputContext: InputContext,
+    Reader: ForestReader<ReaderInputContext>,
+>(
+    input: Input<ReaderInputContext>,
+    storage: MapStorage,
+    expected_forest: OriginalSkeletonForest<'_>,
+    expected_original_contracts_trie_leaves: HashMap<ContractAddress, ContractState>,
+    expected_storage_tries_sorted_indices: HashMap<u128, Vec<u128>>,
+    expected_contracts_trie_sorted_indices: Vec<u128>,
+    expected_classes_trie_sorted_indices: Vec<u128>,
+    reader_generator: impl FnOnce(MapStorage) -> Reader,
 ) {
     let (mut storage_tries_indices, mut contracts_trie_indices, mut classes_trie_indices) =
         get_all_modified_indices(&input.state_diff);
+
     let forest_sorted_indices = ForestSortedIndices {
         storage_tries_sorted_indices: storage_tries_indices
             .iter_mut()
@@ -317,7 +409,8 @@ async fn test_create_original_skeleton_forest(
 
     let actual_storage_updates = input.state_diff.actual_storage_updates();
     let actual_classes_updates = input.state_diff.actual_classes_updates();
-    let (actual_forest, original_contracts_trie_leaves) = FactsDb::new(storage)
+    let mut reader = reader_generator(storage);
+    let (actual_forest, original_contracts_trie_leaves) = reader
         .read(
             input.initial_read_context,
             &actual_storage_updates,
@@ -355,6 +448,49 @@ async fn test_create_original_skeleton_forest(
     }
 }
 
+async fn convert_facts_forest_db_to_index_db(
+    storage: &mut MapStorage,
+    roots: StateRoots,
+) -> MapStorage {
+    let mut contract_leaves: Option<Vec<(NodeIndex, ContractState)>> = Some(Vec::new());
+    let mut index_storage =
+        convert_facts_db_to_index_db::<ContractState, IndexLayoutContractState>(
+            storage,
+            roots.contracts_trie_root_hash,
+            Some(TrieType::ContractsTrie),
+            &mut contract_leaves,
+        )
+        .await
+        .0;
+
+    for (index, contract_leaf) in contract_leaves.unwrap() {
+        println!("storage tries conversion");
+        let contract_address = try_node_index_into_contract_address(&index).unwrap();
+        let storage_root = contract_leaf.storage_root_hash;
+        let contract_storage_entries =
+            convert_facts_db_to_index_db::<StarknetStorageValue, IndexLayoutStarknetStorageValue>(
+                storage,
+                storage_root,
+                Some(TrieType::StorageTrie(contract_address)),
+                &mut None,
+            )
+            .await;
+        index_storage.extend(contract_storage_entries.0);
+    }
+
+    let classes_storage =
+        convert_facts_db_to_index_db::<CompiledClassHash, IndexLayoutCompiledClassHash>(
+            storage,
+            roots.classes_trie_root_hash,
+            Some(TrieType::ClassesTrie),
+            &mut None,
+        )
+        .await;
+    index_storage.extend(classes_storage.0);
+
+    MapStorage(index_storage)
+}
+
 fn create_contract_leaves(leaves: &[(u128, u128)]) -> HashMap<ContractAddress, ContractState> {
     leaves
         .iter()
diff --git a/crates/starknet_committer/src/hash_function/hash.rs b/crates/starknet_committer/src/hash_function/hash.rs
@@ -56,7 +56,22 @@ impl TreeHashFunction<ContractState> for TreeHashFunctionImpl {
 
 impl TreeHashFunction<IndexLayoutContractState> for TreeHashFunctionImpl {
     fn compute_leaf_hash(contract_state: &IndexLayoutContractState) -> HashOutput {
-        compute_contract_state_leaf_hash(contract_state)
+        #[cfg(not(test))]
+        {
+            compute_contract_state_leaf_hash(contract_state)
+        }
+        #[cfg(test)]
+        {
+            use starknet_patricia::patricia_merkle_tree::external_test_utils::AdditionHash;
+            HashOutput(AdditionHash::hash_array(
+                vec![
+                    contract_state.0.class_hash.0,
+                    contract_state.0.storage_root_hash.0,
+                    contract_state.0.nonce.0,
+                ]
+                .as_slice(),
+            ))
+        }
     }
     fn compute_node_hash(node_data: &NodeData<IndexLayoutContractState, HashOutput>) -> HashOutput {
         Self::compute_node_hash_with_inner_hash_function::<PedersenHashFunction>(node_data)
@@ -90,7 +105,15 @@ impl TreeHashFunction<CompiledClassHash> for TreeHashFunctionImpl {
 
 impl TreeHashFunction<IndexLayoutCompiledClassHash> for TreeHashFunctionImpl {
     fn compute_leaf_hash(compiled_class_hash: &IndexLayoutCompiledClassHash) -> HashOutput {
-        compute_compiled_class_leaf_hash(compiled_class_hash)
+        #[cfg(not(test))]
+        {
+            compute_compiled_class_leaf_hash(compiled_class_hash)
+        }
+        #[cfg(test)]
+        {
+            use starknet_patricia::patricia_merkle_tree::external_test_utils::AdditionHash;
+            HashOutput(AdditionHash::hash_array(vec![compiled_class_hash.0.0].as_slice()))
+        }
     }
     fn compute_node_hash(
         node_data: &NodeData<IndexLayoutCompiledClassHash, HashOutput>,
