starknet_committer: add index layout leaves

crates/starknet_committer/Cargo.toml

@@ -13,6 +13,7 @@ testing = ["starknet_patricia/testing"]
 apollo_config.workspace = true
 async-trait.workspace = true
 csv.workspace = true
+derive_more.workspace = true
 ethnum.workspace = true
 hex.workspace = true
 pretty_assertions.workspace = true
@@ -21,7 +22,7 @@ rand_distr.workspace = true
 rstest.workspace = true
 serde = { workspace = true, features = ["derive"] }
 serde_json.workspace = true
-starknet-types-core = { workspace = true, features = ["hash"] }
+starknet-types-core = { workspace = true, features = ["hash", "papyrus-serialization"] }
 starknet_api.workspace = true
 starknet_patricia.workspace = true
 starknet_patricia_storage.workspace = true

crates/starknet_committer/src/db.rs

@@ -3,3 +3,4 @@ mod db_layout;
 pub mod external_test_utils;
 pub mod facts_db;
 pub mod forest_trait;
+pub mod index_db;

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

@@ -0,0 +1,188 @@
+use std::sync::LazyLock;
+
+use starknet_api::core::{ClassHash, ContractAddress, Nonce, PATRICIA_KEY_UPPER_BOUND};
+use starknet_api::hash::HashOutput;
+use starknet_patricia::patricia_merkle_tree::node_data::errors::LeafResult;
+use starknet_patricia::patricia_merkle_tree::node_data::leaf::Leaf;
+use starknet_patricia_storage::db_object::{
+    DBObject,
+    EmptyDeserializationContext,
+    HasStaticPrefix,
+};
+use starknet_patricia_storage::errors::{DeserializationError, SerializationError};
+use starknet_patricia_storage::storage_trait::{DbKeyPrefix, DbValue};
+use starknet_types_core::felt::Felt;
+
+use crate::block_committer::input::StarknetStorageValue;
+use crate::patricia_merkle_tree::leaf::leaf_impl::ContractState;
+use crate::patricia_merkle_tree::types::CompiledClassHash;
+
+// Wrap the leaves types so that we can implement the [DBObject] trait differently in index
+// layout.
+#[derive(Clone, Debug, Default, Eq, PartialEq, derive_more::AsRef, derive_more::From)]
+pub struct IndexLayoutContractState(pub ContractState);
+
+#[derive(Clone, Debug, Default, Eq, PartialEq, derive_more::AsRef, derive_more::From)]
+pub struct IndexLayoutCompiledClassHash(pub CompiledClassHash);
+
+#[derive(Clone, Debug, Default, Eq, PartialEq, derive_more::From)]
+pub struct IndexLayoutStarknetStorageValue(pub StarknetStorageValue);
+
+/// Set to 2^251 + 1 to avoid collisions with contract addresses prefixes.
+static FIRST_AVAILABLE_PREFIX_FELT: LazyLock<Felt> =
+    LazyLock::new(|| Felt::from_hex_unchecked(PATRICIA_KEY_UPPER_BOUND) + Felt::ONE);
+
+/// The db key prefix of nodes in the contracts trie.
+///
+/// Set to FIRST_AVAILABLE_PREFIX_FELT
+static CONTRACTS_TREE_PREFIX: LazyLock<[u8; 32]> =
+    LazyLock::new(|| FIRST_AVAILABLE_PREFIX_FELT.to_bytes_be());
+
+/// The db key prefix of nodes in the contracts trie.
+///
+/// Set to FIRST_AVAILABLE_PREFIX_FELT + 1
+static CLASSES_TREE_PREFIX: LazyLock<[u8; 32]> =
+    LazyLock::new(|| (*FIRST_AVAILABLE_PREFIX_FELT + Felt::ONE).to_bytes_be());
+
+// TODO(Ariel): Delete this enum and use `CommitmentType` instead.
+#[derive(Debug, PartialEq)]
+pub enum TrieType {
+    ContractsTrie,
+    ClassesTrie,
+    StorageTrie(ContractAddress),
+}
+
+impl TrieType {
+    fn db_prefix(&self) -> DbKeyPrefix {
+        match self {
+            Self::ContractsTrie => DbKeyPrefix::new((&CONTRACTS_TREE_PREFIX[..]).into()),
+            Self::ClassesTrie => DbKeyPrefix::new((&CLASSES_TREE_PREFIX[..]).into()),
+            Self::StorageTrie(contract_address) => {
+                let prefix = contract_address.to_bytes_be().to_vec();
+                DbKeyPrefix::new(prefix.into())
+            }
+        }
+    }
+}
+
+macro_rules! impl_has_static_prefix_for_index_layouts {
+    ($($ty:ty),* $(,)?) => {
+        $(
+            impl HasStaticPrefix for $ty {
+                type KeyContext = TrieType;
+                fn get_static_prefix(key_context: &Self::KeyContext) -> DbKeyPrefix {
+                    key_context.db_prefix()
+                }
+            }
+        )*
+    };
+}
+
+impl_has_static_prefix_for_index_layouts! {
+    IndexLayoutContractState,
+    IndexLayoutCompiledClassHash,
+    IndexLayoutStarknetStorageValue,
+}
+
+macro_rules! impl_leaf_for_wrappers {
+    ($($wrapper:ty => $inner:ty),+ $(,)?) => {
+        $(
+            impl Leaf for $wrapper {
+                type Input = <$inner as Leaf>::Input;
+                type Output = <$inner as Leaf>::Output;
+
+                fn is_empty(&self) -> bool {
+                    // assumes `pub struct Wrapper(pub Inner);`
+                    self.0.is_empty()
+                }
+
+                async fn create(
+                    input: Self::Input,
+                ) -> LeafResult<(Self, Self::Output)> {
+                    let (created_leaf, output) = <$inner as Leaf>::create(input).await?;
+                    Ok((Self(created_leaf), output))
+                }
+            }
+        )+
+    };
+}
+
+impl_leaf_for_wrappers!(
+    IndexLayoutContractState => ContractState,
+    IndexLayoutStarknetStorageValue => StarknetStorageValue,
+    IndexLayoutCompiledClassHash => CompiledClassHash,
+);
+
+impl DBObject for IndexLayoutContractState {
+    type DeserializeContext = EmptyDeserializationContext;
+    fn serialize(&self) -> Result<DbValue, SerializationError> {
+        serialize_felts(&[self.0.class_hash.0, self.0.storage_root_hash.0, self.0.nonce.0])
+    }
+
+    fn deserialize(
+        value: &DbValue,
+        _deserialize_context: &Self::DeserializeContext,
+    ) -> Result<Self, DeserializationError> {
+        let mut cursor: &[u8] = &value.0;
+        let err = || DeserializationError::FeltDeserialization(value.clone());
+
+        let class_hash = deserialize_felt(&mut cursor, err)?;
+        let storage_root_hash = deserialize_felt(&mut cursor, err)?;
+        let nonce = deserialize_felt(&mut cursor, err)?;
+
+        Ok(Self(ContractState {
+            class_hash: ClassHash(class_hash),
+            storage_root_hash: HashOutput(storage_root_hash),
+            nonce: Nonce(nonce),
+        }))
+    }
+}
+
+impl DBObject for IndexLayoutCompiledClassHash {
+    type DeserializeContext = EmptyDeserializationContext;
+
+    fn serialize(&self) -> Result<DbValue, SerializationError> {
+        serialize_felts(&[self.0.0])
+    }
+
+    fn deserialize(
+        value: &DbValue,
+        _deserialize_context: &Self::DeserializeContext,
+    ) -> Result<Self, DeserializationError> {
+        Ok(Self(CompiledClassHash(deserialize_felt(&mut &value.0[..], || {
+            DeserializationError::FeltDeserialization(value.clone())
+        })?)))
+    }
+}
+
+impl DBObject for IndexLayoutStarknetStorageValue {
+    type DeserializeContext = EmptyDeserializationContext;
+
+    fn serialize(&self) -> Result<DbValue, SerializationError> {
+        serialize_felts(&[self.0.0])
+    }
+
+    fn deserialize(
+        value: &DbValue,
+        _deserialize_context: &Self::DeserializeContext,
+    ) -> Result<Self, DeserializationError> {
+        Ok(Self(StarknetStorageValue(deserialize_felt(&mut &value.0[..], || {
+            DeserializationError::FeltDeserialization(value.clone())
+        })?)))
+    }
+}
+
+fn deserialize_felt(
+    cursor: &mut &[u8],
+    mk_err: impl Fn() -> DeserializationError,
+) -> Result<Felt, DeserializationError> {
+    Felt::deserialize(cursor).ok_or_else(mk_err)
+}
+
+fn serialize_felts(felts: &[Felt]) -> Result<DbValue, SerializationError> {
+    let mut buffer = Vec::new();
+    for felt in felts {
+        felt.serialize(&mut buffer)?;
+    }
+    Ok(DbValue(buffer))
+}

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

@@ -0,0 +1,76 @@
+use rstest::rstest;
+use starknet_api::core::{ClassHash, Nonce};
+use starknet_api::hash::HashOutput;
+use starknet_patricia::patricia_merkle_tree::node_data::leaf::Leaf;
+use starknet_patricia_storage::db_object::EmptyDeserializationContext;
+use starknet_patricia_storage::storage_trait::DbValue;
+use starknet_types_core::felt::Felt;
+
+use crate::block_committer::input::StarknetStorageValue;
+use crate::db::index_db::leaves::{
+    IndexLayoutCompiledClassHash,
+    IndexLayoutContractState,
+    IndexLayoutStarknetStorageValue,
+};
+use crate::patricia_merkle_tree::leaf::leaf_impl::ContractState;
+use crate::patricia_merkle_tree::types::CompiledClassHash;
+
+fn contract_state_leaf() -> IndexLayoutContractState {
+    IndexLayoutContractState(ContractState {
+        class_hash: ClassHash(Felt::from(1)),
+        storage_root_hash: HashOutput(Felt::from(2)),
+        nonce: Nonce(Felt::from(3)),
+    })
+}
+
+fn compiled_class_hash_leaf() -> IndexLayoutCompiledClassHash {
+    IndexLayoutCompiledClassHash(CompiledClassHash(Felt::ONE))
+}
+
+fn starknet_storage_value_leaf() -> IndexLayoutStarknetStorageValue {
+    IndexLayoutStarknetStorageValue(StarknetStorageValue(Felt::ONE))
+}
+
+fn starknet_storage_value_leaf_96_bits() -> IndexLayoutStarknetStorageValue {
+    // 2^96 (12 bytes, under the 27 nibbles threshold)
+    IndexLayoutStarknetStorageValue(StarknetStorageValue(Felt::from(1_u128 << 95)))
+}
+
+fn starknet_storage_value_leaf_136_bits() -> IndexLayoutStarknetStorageValue {
+    // 2^136 (reaching the 34 nibbles / 17 bytes serialization threshold)
+    let mut bytes = [0u8; 32];
+    bytes[15] = 128;
+    IndexLayoutStarknetStorageValue(StarknetStorageValue(Felt::from_bytes_be(&bytes)))
+}
+
+#[rstest]
+#[case::index_layout_contract_state(contract_state_leaf())]
+#[case::index_layout_compiled_class_hash(compiled_class_hash_leaf())]
+#[case::index_layout_starknet_storage_value(starknet_storage_value_leaf())]
+fn test_index_layout_leaf_serde<L: Leaf>(#[case] leaf: L) {
+    let serialized = leaf.serialize().unwrap();
+    let deserialized = L::deserialize(&serialized, &EmptyDeserializationContext).unwrap();
+    assert_eq!(leaf, deserialized);
+}
+
+#[rstest]
+#[case(contract_state_leaf(), DbValue(vec![1, 2, 3]))]
+#[case(compiled_class_hash_leaf(), DbValue(vec![1]))]
+#[case(starknet_storage_value_leaf(), DbValue(vec![1]))]
+// We are serializing 2^96. The 4 MSB of the first byte are the chooser. For values >= 16 but under
+// 27 nibbles, the chooser is the number of bytes. In this case, the first byte will be 11000000
+// (chooser 12, i.e. we need 12 bytes) followed by the value.
+#[case(starknet_storage_value_leaf_96_bits(), DbValue([vec![192, 128], vec![0; 11]].concat()))]
+// We are serializing 2^136, which exceeds the 34 nibbles threshold where the encoding utilizes the
+// full 32 bytes. This case is marked by chooser = 15, followed by the value, starting immediately
+// after the chooser (hence the first 116 bits after the chooser are 0).
+#[case(starknet_storage_value_leaf_136_bits(), DbValue([
+    vec![240], vec![0; 14], vec![128], vec![0; 16]].concat()
+))]
+fn test_leaf_serialization_regression<L: Leaf>(
+    #[case] leaf: L,
+    #[case] expected_serialize: DbValue,
+) {
+    let actual_serialize = leaf.serialize().unwrap();
+    assert_eq!(actual_serialize, expected_serialize);
+}

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

@@ -0,0 +1,3 @@
+pub mod leaves;
+#[cfg(test)]
+pub mod leaves_test;