starknet_committer,starknet_patricia: add serialization utilities to node layout

Author: ArielElp

crates/starknet_committer/src/db/facts_db/db.rs

@@ -4,18 +4,19 @@ use async_trait::async_trait;
 use starknet_api::core::ContractAddress;
 use starknet_api::hash::HashOutput;
 use starknet_patricia::db_layout::{NodeLayout, TrieType};
-use starknet_patricia::patricia_merkle_tree::filled_tree::node::FactDbFilledNode;
+use starknet_patricia::patricia_merkle_tree::filled_tree::node::{FactDbFilledNode, FilledNode};
 use starknet_patricia::patricia_merkle_tree::filled_tree::node_serde::FactNodeDeserializationContext;
 use starknet_patricia::patricia_merkle_tree::filled_tree::tree::FilledTree;
 use starknet_patricia::patricia_merkle_tree::node_data::leaf::{
+    Leaf,
     LeafModifications,
     LeafWithEmptyKeyContext,
 };
 use starknet_patricia::patricia_merkle_tree::types::NodeIndex;
-use starknet_patricia_storage::db_object::{EmptyKeyContext, HasStaticPrefix};
+use starknet_patricia_storage::db_object::{DBObject, EmptyKeyContext, HasStaticPrefix};
 use starknet_patricia_storage::errors::SerializationResult;
 use starknet_patricia_storage::map_storage::MapStorage;
-use starknet_patricia_storage::storage_trait::{DbHashMap, Storage};
+use starknet_patricia_storage::storage_trait::{DbHashMap, DbKey, Storage};
 
 use crate::block_committer::input::{ReaderConfig, StarknetStorageValue};
 use crate::db::facts_db::types::{FactsDbInitialRead, FactsSubTree};
@@ -46,6 +47,18 @@ impl<'a, L: LeafWithEmptyKeyContext> NodeLayout<'a, L> for FactsNodeLayout {
     fn generate_key_context(_trie_type: TrieType) -> <L as HasStaticPrefix>::KeyContext {
         EmptyKeyContext
     }
+
+    fn get_db_object<LeafBase: Leaf + Into<L>>(
+        _node_index: NodeIndex,
+        key_context: &<L as HasStaticPrefix>::KeyContext,
+        filled_node: FilledNode<LeafBase, HashOutput>,
+    ) -> (DbKey, Self::NodeDbObject) {
+        let db_filled_node = Self::convert_node_data_and_leaf(filled_node);
+
+        let key = db_filled_node.get_db_key(key_context, &db_filled_node.hash.0.to_bytes_be());
+
+        (key, db_filled_node)
+    }
 }
 
 pub struct FactsDb<S: Storage> {

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

@@ -2,14 +2,16 @@ use std::collections::HashMap;
 
 use async_trait::async_trait;
 use starknet_api::core::ContractAddress;
+use starknet_api::hash::HashOutput;
 use starknet_patricia::db_layout::{NodeLayout, TrieType};
+use starknet_patricia::patricia_merkle_tree::filled_tree::node::FilledNode;
 use starknet_patricia::patricia_merkle_tree::filled_tree::tree::FilledTree;
 use starknet_patricia::patricia_merkle_tree::node_data::leaf::{Leaf, LeafModifications};
 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::{EmptyKeyContext, HasStaticPrefix};
+use starknet_patricia_storage::db_object::{DBObject, EmptyKeyContext, HasStaticPrefix};
 use starknet_patricia_storage::errors::SerializationResult;
-use starknet_patricia_storage::storage_trait::{DbHashMap, Storage};
+use starknet_patricia_storage::storage_trait::{DbHashMap, DbKey, Storage};
 
 use crate::block_committer::input::{ReaderConfig, StarknetStorageValue};
 use crate::db::facts_db::types::FactsDbInitialRead;
@@ -58,6 +60,20 @@ where
     fn generate_key_context(trie_type: TrieType) -> <L as HasStaticPrefix>::KeyContext {
         trie_type
     }
+
+    fn get_db_object<LeafBase: Leaf + Into<L>>(
+        node_index: NodeIndex,
+        key_context: &<L as HasStaticPrefix>::KeyContext,
+        filled_node: FilledNode<LeafBase, HashOutput>,
+    ) -> (DbKey, Self::NodeDbObject) {
+        let filled_node = Self::convert_node_data_and_leaf(filled_node);
+
+        let db_filled_node = IndexFilledNode(filled_node);
+
+        let key = db_filled_node.get_db_key(key_context, &node_index.0.to_be_bytes());
+
+        (key, db_filled_node)
+    }
 }
 
 // TODO(Ariel): define an IndexDbInitialRead empty type, and check whether each tree is empty inside

crates/starknet_patricia/src/db_layout.rs

@@ -1,10 +1,13 @@
 use starknet_api::core::ContractAddress;
 use starknet_api::hash::HashOutput;
 use starknet_patricia_storage::db_object::{DBObject, HasStaticPrefix};
+use starknet_patricia_storage::storage_trait::DbKey;
 
 use crate::patricia_merkle_tree::filled_tree::node::FilledNode;
+use crate::patricia_merkle_tree::node_data::inner_node::{BinaryData, EdgeData, NodeData};
 use crate::patricia_merkle_tree::node_data::leaf::Leaf;
 use crate::patricia_merkle_tree::traversal::SubTreeTrait;
+use crate::patricia_merkle_tree::types::NodeIndex;
 
 // TODO(Ariel): Delete this enum and use `CommitmentType` instead.
 #[derive(Debug, PartialEq)]
@@ -31,8 +34,10 @@ pub trait NodeLayout<'a, L: Leaf> {
     type DeserializationContext;
 
     /// The storage representation of the node.
-    type NodeDbObject: DBObject<DeserializeContext = Self::DeserializationContext>
-        + Into<FilledNode<L, Self::NodeData>>;
+    type NodeDbObject: DBObject<
+            DeserializeContext = Self::DeserializationContext,
+            KeyContext = <L as HasStaticPrefix>::KeyContext,
+        > + Into<FilledNode<L, Self::NodeData>>;
 
     /// The type of the subtree that is used to traverse the trie.
     type SubTree: SubTreeTrait<
@@ -44,4 +49,37 @@ pub trait NodeLayout<'a, L: Leaf> {
     /// Generates the key context for the given trie type. Used for reading nodes of a specific
     /// tree (contracts, classes, or storage), to construct a skeleton tree.
     fn generate_key_context(trie_type: TrieType) -> <L as HasStaticPrefix>::KeyContext;
+
+    /// Converts `FilledTree` nodes to db objects.
+    ///
+    /// During the construction of a `FilledTree` we computee the hashes and carry `FilledNode<L,
+    /// HashOutput>`, hence, the `FilledNode` type  is not necessarily what we want to store.
+    fn get_db_object<LeafBase: Leaf + Into<L>>(
+        node_index: NodeIndex,
+        key_context: &<L as HasStaticPrefix>::KeyContext,
+        filled_node: FilledNode<LeafBase, HashOutput>,
+    ) -> (DbKey, Self::NodeDbObject);
+
+    /// A utility function to convert a `FilledNode<LeafBase, HashOutput>` to a `FilledNode<L,
+    /// Self::NodeData>`. Used during the serialization of a `FilledTree`.
+    ///
+    /// LeafBase is one of StarknetStorageValue, CompiledClassHash, or ContractState, while L can be
+    /// layout-dependent wrappers.
+    fn convert_node_data_and_leaf<LeafBase: Leaf + Into<L>>(
+        filled_node: FilledNode<LeafBase, HashOutput>,
+    ) -> FilledNode<L, Self::NodeData> {
+        let data: NodeData<L, Self::NodeData> = match filled_node.data {
+            NodeData::Binary(data) => NodeData::Binary(BinaryData {
+                left_data: data.left_data.into(),
+                right_data: data.right_data.into(),
+            }),
+            NodeData::Edge(data) => NodeData::Edge(EdgeData {
+                bottom_data: data.bottom_data.into(),
+                path_to_bottom: data.path_to_bottom.into(),
+            }),
+            NodeData::Leaf(leaf) => NodeData::Leaf(leaf.into()),
+        };
+
+        FilledNode { hash: filled_node.hash, data }
+    }
 }