starknet_committer,starknet_patricia: layout-based original tree creation tests

Author: ArielElp

Cargo.lock

@@ -31,6 +31,7 @@ tokio = { workspace = true, features = ["rt"] }
 tracing.workspace = true
 
 [dev-dependencies]
+async-recursion.workspace = true
 rstest.workspace = true
 rstest_reuse.workspace = true
 starknet_api = { workspace = true, features = ["testing"] }

crates/starknet_committer/Cargo.toml

@@ -0,0 +1,50 @@
+use rstest::rstest;
+use rstest_reuse::apply;
+use starknet_api::hash::HashOutput;
+use starknet_patricia::patricia_merkle_tree::external_test_utils::MockLeaf;
+use starknet_patricia::patricia_merkle_tree::node_data::inner_node::NodeData;
+use starknet_patricia::patricia_merkle_tree::updated_skeleton_tree::hash_function::TreeHashFunction;
+use starknet_patricia_storage::db_object::EmptyKeyContext;
+use starknet_types_core::felt::Felt;
+
+use crate::db::create_original_skeleton_tests::case_helpers::CreateTreeCase;
+use crate::db::create_original_skeleton_tests::{create_tree_cases, test_create_original_skeleton};
+use crate::db::index_db::db::IndexNodeLayout;
+use crate::db::index_db::test_utils::convert_facts_db_to_index_db;
+use crate::hash_function::hash::TreeHashFunctionImpl;
+
+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>) -> HashOutput {
+        HashOutput(Felt::ZERO)
+    }
+}
+
+#[apply(create_tree_cases)]
+#[rstest]
+#[tokio::test]
+async fn test_create_tree_index_layout(
+    #[case] mut case: CreateTreeCase,
+    #[values(true, false)] compare_modified_leaves: bool,
+) {
+    let mut storage = convert_facts_db_to_index_db::<MockLeaf, MockLeaf, EmptyKeyContext>(
+        &mut case.storage,
+        case.root_hash,
+        &EmptyKeyContext,
+        &mut None,
+    )
+    .await;
+
+    test_create_original_skeleton::<MockLeaf, IndexNodeLayout>(
+        &mut storage,
+        &case.leaf_modifications,
+        case.root_hash,
+        &case.expected_skeleton_nodes,
+        case.subtree_height,
+        compare_modified_leaves,
+    )
+    .await;
+}

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

@@ -1,5 +1,9 @@
+#[cfg(test)]
+pub mod create_index_tree_test;
 pub mod db;
 pub mod leaves;
 #[cfg(test)]
 pub mod serde_tests;
+#[cfg(any(feature = "testing", test))]
+pub mod test_utils;
 pub mod types;

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

@@ -0,0 +1,267 @@
+use async_recursion::async_recursion;
+use starknet_api::core::ContractAddress;
+use starknet_api::hash::{HashOutput, StateRoots};
+use starknet_patricia::patricia_merkle_tree::filled_tree::node::{FactDbFilledNode, FilledNode};
+use starknet_patricia::patricia_merkle_tree::filled_tree::node_serde::{
+    FactNodeDeserializationContext,
+    PatriciaPrefix,
+    FACT_LAYOUT_DB_KEY_SEPARATOR,
+};
+use starknet_patricia::patricia_merkle_tree::node_data::inner_node::{
+    BinaryData,
+    EdgeData,
+    NodeData,
+};
+use starknet_patricia::patricia_merkle_tree::node_data::leaf::Leaf;
+use starknet_patricia::patricia_merkle_tree::types::NodeIndex;
+use starknet_patricia::patricia_merkle_tree::updated_skeleton_tree::hash_function::TreeHashFunction;
+use starknet_patricia_storage::db_object::{DBObject, EmptyKeyContext, HasStaticPrefix};
+use starknet_patricia_storage::map_storage::MapStorage;
+use starknet_patricia_storage::storage_trait::{create_db_key, DbHashMap, DbValue, Storage};
+
+use crate::block_committer::input::{try_node_index_into_contract_address, StarknetStorageValue};
+use crate::db::index_db::leaves::{
+    IndexLayoutCompiledClassHash,
+    IndexLayoutContractState,
+    IndexLayoutStarknetStorageValue,
+    INDEX_LAYOUT_DB_KEY_SEPARATOR,
+};
+use crate::db::index_db::types::{EmptyNodeData, IndexFilledNode};
+use crate::hash_function::hash::TreeHashFunctionImpl;
+use crate::patricia_merkle_tree::leaf::leaf_impl::ContractState;
+use crate::patricia_merkle_tree::types::CompiledClassHash;
+
+/// Converts a complete Facts-layout forest DB to Index-layout.
+/// Expects all nodes to exist (panics if a node is missing).
+/// Use this for converting input storage before running tests.
+pub async fn convert_facts_forest_db_to_index_db(
+    storage: &mut MapStorage,
+    roots: StateRoots,
+) -> MapStorage {
+    convert_forest(storage, roots, false).await
+}
+
+/// Converts a Facts-layout filled forest (delta) to Index-layout.
+/// Skips missing branches (nodes that weren't updated).
+/// Use this for converting expected_facts which only contains updated nodes.
+pub async fn convert_facts_filled_forest_to_index(
+    storage: &mut MapStorage,
+    roots: StateRoots,
+) -> MapStorage {
+    convert_forest(storage, roots, true).await
+}
+
+/// Converts a single Facts-layout trie to Index-layout.
+/// Expects all nodes to exist (panics if a node is missing).
+pub async fn convert_facts_db_to_index_db<FactsLeaf, IndexLeaf, KeyContext>(
+    storage: &mut MapStorage,
+    root_hash: HashOutput,
+    key_context: &KeyContext,
+    current_leaves: &mut Option<Vec<(NodeIndex, FactsLeaf)>>,
+) -> MapStorage
+where
+    FactsLeaf: Leaf + Into<IndexLeaf> + HasStaticPrefix<KeyContext = KeyContext>,
+    IndexLeaf: Leaf + HasStaticPrefix<KeyContext = KeyContext>,
+    TreeHashFunctionImpl: TreeHashFunction<IndexLeaf>,
+    KeyContext: Sync,
+{
+    convert_single_trie(storage, root_hash, key_context, current_leaves, false).await
+}
+
+/// Converts all tries in a forest from Facts-layout to Index-layout.
+///
+/// If skip_missing is true, missing nodes are silently skipped. Used when the original facts DB
+/// only contains the tree subset that was necessary for tests. Use false when the input is known to
+/// contain a full tree.
+async fn convert_forest(
+    storage: &mut MapStorage,
+    roots: StateRoots,
+    skip_missing: bool,
+) -> MapStorage {
+    let mut contract_leaves: Option<Vec<(NodeIndex, ContractState)>> = Some(Vec::new());
+    let mut index_storage =
+        convert_single_trie::<ContractState, IndexLayoutContractState, EmptyKeyContext>(
+            storage,
+            roots.contracts_trie_root_hash,
+            &EmptyKeyContext,
+            &mut contract_leaves,
+            skip_missing,
+        )
+        .await
+        .0;
+
+    for (index, contract_leaf) in contract_leaves.unwrap() {
+        let contract_address = try_node_index_into_contract_address(&index).unwrap();
+        let storage_root = contract_leaf.storage_root_hash;
+        let contract_storage_entries =
+            convert_single_trie::<
+                StarknetStorageValue,
+                IndexLayoutStarknetStorageValue,
+                ContractAddress,
+            >(storage, storage_root, &contract_address, &mut None, skip_missing)
+            .await;
+        index_storage.extend(contract_storage_entries.0);
+    }
+
+    let classes_storage =
+        convert_single_trie::<CompiledClassHash, IndexLayoutCompiledClassHash, EmptyKeyContext>(
+            storage,
+            roots.classes_trie_root_hash,
+            &EmptyKeyContext,
+            &mut None,
+            skip_missing,
+        )
+        .await;
+    index_storage.extend(classes_storage.0);
+
+    MapStorage(index_storage)
+}
+
+/// Converts a single trie from Facts-layout to Index-layout.
+async fn convert_single_trie<FactsLeaf, IndexLeaf, KeyContext>(
+    storage: &mut MapStorage,
+    root_hash: HashOutput,
+    key_context: &KeyContext,
+    current_leaves: &mut Option<Vec<(NodeIndex, FactsLeaf)>>,
+    skip_missing: bool,
+) -> MapStorage
+where
+    FactsLeaf: Leaf + Into<IndexLeaf> + HasStaticPrefix<KeyContext = KeyContext>,
+    IndexLeaf: Leaf + HasStaticPrefix<KeyContext = KeyContext>,
+    TreeHashFunctionImpl: TreeHashFunction<IndexLeaf>,
+    KeyContext: Sync,
+{
+    let mut index_layout_storage = MapStorage(DbHashMap::new());
+    traverse_and_convert::<FactsLeaf, IndexLeaf, KeyContext>(
+        storage,
+        &mut index_layout_storage,
+        NodeIndex::ROOT,
+        root_hash,
+        key_context,
+        current_leaves,
+        skip_missing,
+    )
+    .await;
+    index_layout_storage
+}
+
+/// Recursively traverses a Facts-layout trie and converts each node to Index-layout.
+#[async_recursion]
+async fn traverse_and_convert<FactsLeaf, IndexLeaf, KeyContext>(
+    facts_storage: &mut MapStorage,
+    index_layout_storage: &mut MapStorage,
+    current_index: NodeIndex,
+    current_hash: HashOutput,
+    key_context: &KeyContext,
+    current_leaves: &mut Option<Vec<(NodeIndex, FactsLeaf)>>,
+    skip_missing: bool,
+) where
+    FactsLeaf: Leaf + Into<IndexLeaf> + HasStaticPrefix<KeyContext = KeyContext>,
+    IndexLeaf: Leaf + HasStaticPrefix<KeyContext = KeyContext>,
+    TreeHashFunctionImpl: TreeHashFunction<IndexLeaf>,
+    KeyContext: Sync,
+{
+    if current_hash == HashOutput::ROOT_OF_EMPTY_TREE {
+        return;
+    }
+    let db_prefix = if current_index.is_leaf() {
+        FactsLeaf::get_static_prefix(key_context)
+    } else {
+        PatriciaPrefix::InnerNode.into()
+    };
+    let facts_db_key =
+        create_db_key(db_prefix, FACT_LAYOUT_DB_KEY_SEPARATOR, &current_hash.0.to_bytes_be());
+
+    // Try to get the node from storage.
+    let filled_node_raw: Option<DbValue> = facts_storage.get(&facts_db_key).await.unwrap();
+
+    // Handle missing nodes based on the skip_missing flag.
+    let Some(filled_node_raw) = filled_node_raw else {
+        if skip_missing {
+            return;
+        } else {
+            panic!(
+                "Node not found in storage: index={:?}, hash={:?}. If converting a filled forest, \
+                 use convert_facts_filled_forest_to_index.",
+                current_index, current_hash
+            );
+        }
+    };
+
+    let facts_filled_node = FactDbFilledNode::<FactsLeaf>::deserialize(
+        &filled_node_raw,
+        &FactNodeDeserializationContext {
+            node_hash: current_hash,
+            is_leaf: current_index.is_leaf(),
+        },
+    )
+    .unwrap();
+
+    let index_filled_node: IndexFilledNode<IndexLeaf> = match facts_filled_node.data {
+        NodeData::Binary(binary_data) => {
+            let children_indices = current_index.get_children_indices();
+            traverse_and_convert::<FactsLeaf, IndexLeaf, KeyContext>(
+                facts_storage,
+                index_layout_storage,
+                children_indices[0],
+                binary_data.left_data,
+                key_context,
+                current_leaves,
+                skip_missing,
+            )
+            .await;
+            traverse_and_convert::<FactsLeaf, IndexLeaf, KeyContext>(
+                facts_storage,
+                index_layout_storage,
+                children_indices[1],
+                binary_data.right_data,
+                key_context,
+                current_leaves,
+                skip_missing,
+            )
+            .await;
+            IndexFilledNode(FilledNode {
+                hash: current_hash,
+                data: NodeData::Binary(BinaryData {
+                    left_data: EmptyNodeData,
+                    right_data: EmptyNodeData,
+                }),
+            })
+        }
+        NodeData::Edge(edge_data) => {
+            let bottom_index =
+                NodeIndex::compute_bottom_index(current_index, &edge_data.path_to_bottom);
+            traverse_and_convert::<FactsLeaf, IndexLeaf, KeyContext>(
+                facts_storage,
+                index_layout_storage,
+                bottom_index,
+                edge_data.bottom_data,
+                key_context,
+                current_leaves,
+                skip_missing,
+            )
+            .await;
+            IndexFilledNode(FilledNode {
+                hash: current_hash,
+                data: NodeData::Edge(EdgeData {
+                    bottom_data: EmptyNodeData,
+                    path_to_bottom: edge_data.path_to_bottom,
+                }),
+            })
+        }
+        NodeData::Leaf(leaf) => {
+            if let Some(leaves) = current_leaves {
+                leaves.push((current_index, leaf.clone()));
+            }
+
+            IndexFilledNode(FilledNode { hash: current_hash, data: NodeData::Leaf(leaf.into()) })
+        }
+    };
+
+    let index_db_key = create_db_key(
+        IndexLeaf::get_static_prefix(key_context),
+        INDEX_LAYOUT_DB_KEY_SEPARATOR,
+        &current_index.0.to_be_bytes(),
+    );
+    index_layout_storage.set(index_db_key, index_filled_node.serialize().unwrap()).await.unwrap();
+}

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

@@ -68,15 +68,12 @@ impl NodeIndex {
     }
 
     // TODO(Amos, 1/5/2024): Move to EdgePath.
-    pub(crate) fn compute_bottom_index(
-        index: NodeIndex,
-        path_to_bottom: &PathToBottom,
-    ) -> NodeIndex {
+    pub fn compute_bottom_index(index: NodeIndex, path_to_bottom: &PathToBottom) -> NodeIndex {
         let PathToBottom { path, length, .. } = path_to_bottom;
         (index << u8::from(*length)) + Self::new(path.into())
     }
 
-    pub(crate) fn get_children_indices(&self) -> [Self; 2] {
+    pub fn get_children_indices(&self) -> [Self; 2] {
         let left_child = *self << 1;
         [left_child, left_child + 1]
     }