starknet_committer: define StarknetForestProofs serde

crates/starknet_committer/Cargo.toml

@@ -7,12 +7,13 @@ license.workspace = true
 description = "Computes and manages Starknet state."
 
 [features]
-os_input = ["dep:blake2", "dep:digest"]
+os_input = ["dep:bincode", "dep:blake2", "dep:digest", "dep:zstd"]
 testing = ["starknet_patricia/testing"]
 
 [dependencies]
 apollo_config.workspace = true
 async-trait.workspace = true
+bincode = { workspace = true, optional = true }
 blake2 = { workspace = true, optional = true }
 derive_more = { workspace = true, features = ["as_ref", "from", "into"] }
 digest = { workspace = true, optional = true }
@@ -31,6 +32,7 @@ strum.workspace = true
 thiserror.workspace = true
 tokio = { workspace = true, features = ["rt"] }
 tracing.workspace = true
+zstd = { workspace = true, optional = true }
 
 [dev-dependencies]
 async-recursion.workspace = true

crates/starknet_committer/src/patricia_merkle_tree.rs

@@ -1,4 +1,8 @@
 pub mod leaf;
+#[cfg(feature = "os_input")]
+mod starknet_forest_proofs_serde;
+#[cfg(all(test, feature = "os_input"))]
+mod starknet_forest_proofs_serialization_test;
 pub mod tree;
 pub mod types;
 #[cfg(test)]

crates/starknet_committer/src/patricia_merkle_tree/starknet_forest_proofs_serde.rs

@@ -0,0 +1,150 @@
+use std::collections::HashMap;
+
+use starknet_api::core::{ClassHash, ContractAddress, Nonce};
+use starknet_api::hash::HashOutput;
+use starknet_patricia::patricia_merkle_tree::node_data::inner_node::{
+    flatten_preimages,
+    Preimage,
+    PreimageMap,
+};
+use starknet_patricia_storage::errors::{
+    DeserializationError,
+    SerializationError,
+    SerializationResult,
+};
+use starknet_patricia_storage::storage_trait::DbValue;
+use starknet_types_core::felt::Felt;
+
+use crate::patricia_merkle_tree::leaf::leaf_impl::ContractState;
+use crate::patricia_merkle_tree::types::{ContractsTrieProof, StarknetForestProofs};
+
+type RawPreimages = Vec<(HashOutput, Vec<Felt>)>;
+type ContractTrieLeaves = Vec<(ContractAddress, (Nonce, HashOutput, ClassHash))>;
+type StorageTrieProofs = Vec<(ContractAddress, RawPreimages)>;
+type SerializedStarknetForestProofs =
+    (RawPreimages, RawPreimages, ContractTrieLeaves, StorageTrieProofs);
+
+impl StarknetForestProofs {
+    /// Zstd-compressed bincode payload for the OS-input witnesses.
+    ///
+    /// The inner bincode payload is a 4-tuple, in order:
+    ///
+    /// 1. Classes trie inner nodes — `RawPreimages`.
+    /// 2. Contract trie inner nodes — `RawPreimages`.
+    /// 3. Contract trie leaves — `ContractTrieLeaves`.
+    /// 4. Storage tries inner nodes — `StorageTrieProofs`.
+    ///
+    /// Each `commitment_facts` entry uses the same encoding as [`CommitmentInfo::commitment_facts`]
+    /// and OS Patricia hints:
+    ///
+    /// - Binary node — `[left: Felt, right: Felt]`.
+    /// - Edge node — `[length: Felt, path: Felt, bottom: Felt]`.
+    pub fn serialize(&self) -> SerializationResult<DbValue> {
+        let classes: RawPreimages =
+            flatten_preimages(&self.classes_trie_proof).into_iter().collect();
+        let contract_nodes: RawPreimages =
+            flatten_preimages(&self.contracts_trie_proof.nodes).into_iter().collect();
+        let contract_leaves: ContractTrieLeaves = self
+            .contracts_trie_proof
+            .leaves
+            .iter()
+            .map(|(addr, contract_state)| {
+                (
+                    *addr,
+                    (
+                        contract_state.nonce,
+                        contract_state.storage_root_hash,
+                        contract_state.class_hash,
+                    ),
+                )
+            })
+            .collect();
+
+        let storage: StorageTrieProofs = self
+            .contracts_trie_storage_proofs
+            .iter()
+            .map(|(addr, preimage_map)| {
+                (*addr, flatten_preimages(preimage_map).into_iter().collect())
+            })
+            .collect();
+
+        let bincode_payload =
+            bincode::serialize(&(classes, contract_nodes, contract_leaves, storage))
+                .map_err(bincode_ser_err)?;
+        let compressed =
+            zstd::encode_all(bincode_payload.as_slice(), zstd::DEFAULT_COMPRESSION_LEVEL)
+                .map_err(SerializationError::IOSerialize)?;
+        Ok(DbValue(compressed))
+    }
+
+    pub fn deserialize(value: &DbValue) -> Result<Self, DeserializationError> {
+        let bincode_payload = zstd::decode_all(value.0.as_slice()).map_err(|error| {
+            DeserializationError::ValueError(Box::new(std::io::Error::new(
+                std::io::ErrorKind::InvalidData,
+                error.to_string(),
+            )))
+        })?;
+        let (classes, contract_nodes, contract_leaves, storage): SerializedStarknetForestProofs =
+            bincode::deserialize(&bincode_payload).map_err(|error| {
+                DeserializationError::ValueError(Box::new(std::io::Error::new(
+                    std::io::ErrorKind::InvalidData,
+                    error.to_string(),
+                )))
+            })?;
+
+        let classes_trie_proof = preimage_map_from_commitment_facts(classes)?;
+        let contracts_trie_proof = ContractsTrieProof {
+            nodes: preimage_map_from_commitment_facts(contract_nodes)?,
+            leaves: contract_leaves.into_iter().try_fold(
+                HashMap::new(),
+                |mut leaves, (addr, (nonce, storage_root_hash, class_hash))| {
+                    if leaves
+                        .insert(addr, ContractState { nonce, storage_root_hash, class_hash })
+                        .is_some()
+                    {
+                        return Err(DeserializationError::KeyDuplicate(format!(
+                            "duplicate contracts trie leaf {addr:?}"
+                        )));
+                    }
+                    Ok(leaves)
+                },
+            )?,
+        };
+
+        let contracts_trie_storage_proofs =
+            storage.into_iter().try_fold(HashMap::new(), |mut proofs, (addr, facts)| {
+                if proofs.insert(addr, preimage_map_from_commitment_facts(facts)?).is_some() {
+                    return Err(DeserializationError::KeyDuplicate(format!(
+                        "duplicate storage trie witness address {addr:?}"
+                    )));
+                }
+                Ok(proofs)
+            })?;
+
+        Ok(Self { classes_trie_proof, contracts_trie_proof, contracts_trie_storage_proofs })
+    }
+}
+
+fn bincode_ser_err(error: bincode::Error) -> SerializationError {
+    SerializationError::IOSerialize(std::io::Error::other(error))
+}
+
+fn preimage_map_from_commitment_facts(
+    facts: RawPreimages,
+) -> Result<PreimageMap, DeserializationError> {
+    let mut preimage_map = PreimageMap::new();
+    for (hash, raw_preimage) in facts {
+        let preimage = Preimage::try_from(&raw_preimage).map_err(|error| {
+            DeserializationError::ValueError(Box::new(std::io::Error::new(
+                std::io::ErrorKind::InvalidData,
+                error.to_string(),
+            )))
+        })?;
+        if preimage_map.insert(hash, preimage).is_some() {
+            return Err(DeserializationError::KeyDuplicate(format!(
+                "duplicate preimage node hash {hash:?}"
+            )));
+        }
+    }
+    Ok(preimage_map)
+}

crates/starknet_committer/src/patricia_merkle_tree/starknet_forest_proofs_serialization_test.rs

@@ -0,0 +1,137 @@
+use std::collections::HashMap;
+
+use ethnum::U256;
+use pretty_assertions::assert_eq;
+use rstest::rstest;
+use starknet_api::core::{ClassHash, ContractAddress, Nonce};
+use starknet_api::hash::HashOutput;
+use starknet_patricia::patricia_merkle_tree::node_data::inner_node::{
+    BinaryData,
+    EdgeData,
+    EdgePath,
+    EdgePathLength,
+    PathToBottom,
+    Preimage,
+    PreimageMap,
+};
+use starknet_types_core::felt::Felt;
+
+use crate::patricia_merkle_tree::leaf::leaf_impl::ContractState;
+use crate::patricia_merkle_tree::types::{ContractsTrieProof, StarknetForestProofs};
+
+fn binary_preimage(node_hash: u128, left: u128, right: u128) -> (HashOutput, Preimage) {
+    (
+        HashOutput(Felt::from(node_hash)),
+        Preimage::Binary(BinaryData {
+            left_data: HashOutput(Felt::from(left)),
+            right_data: HashOutput(Felt::from(right)),
+        }),
+    )
+}
+
+fn edge_preimage(node_hash: u128, bottom: u128, path: u128, length: u8) -> (HashOutput, Preimage) {
+    (
+        HashOutput(Felt::from(node_hash)),
+        Preimage::Edge(EdgeData {
+            bottom_data: HashOutput(Felt::from(bottom)),
+            path_to_bottom: PathToBottom::new(
+                EdgePath(U256::from(path)),
+                EdgePathLength::new(length).unwrap(),
+            )
+            .unwrap(),
+        }),
+    )
+}
+
+fn contract_state(nonce: u128, storage_root: u128, class_hash: u128) -> ContractState {
+    ContractState {
+        nonce: Nonce(Felt::from(nonce)),
+        storage_root_hash: HashOutput(Felt::from(storage_root)),
+        class_hash: ClassHash(Felt::from(class_hash)),
+    }
+}
+
+fn empty_contracts_trie_proof() -> ContractsTrieProof {
+    ContractsTrieProof { nodes: PreimageMap::new(), leaves: HashMap::new() }
+}
+
+fn only_classes() -> StarknetForestProofs {
+    StarknetForestProofs {
+        classes_trie_proof: PreimageMap::from([
+            binary_preimage(100, 101, 102),
+            edge_preimage(200, 201, 0, 1),
+        ]),
+        contracts_trie_proof: empty_contracts_trie_proof(),
+        contracts_trie_storage_proofs: HashMap::new(),
+    }
+}
+
+fn only_contracts() -> StarknetForestProofs {
+    let address = ContractAddress::try_from(Felt::from(7)).unwrap();
+    StarknetForestProofs {
+        classes_trie_proof: PreimageMap::new(),
+        contracts_trie_proof: ContractsTrieProof {
+            nodes: PreimageMap::from([binary_preimage(300, 301, 302)]),
+            leaves: HashMap::from([(address, contract_state(1, 400, 500))]),
+        },
+        contracts_trie_storage_proofs: HashMap::new(),
+    }
+}
+
+fn classes_and_contracts() -> StarknetForestProofs {
+    let address = ContractAddress::try_from(Felt::from(9)).unwrap();
+    StarknetForestProofs {
+        classes_trie_proof: PreimageMap::from([binary_preimage(110, 111, 112)]),
+        contracts_trie_proof: ContractsTrieProof {
+            nodes: PreimageMap::from([edge_preimage(310, 311, 1, 1)]),
+            leaves: HashMap::from([(address, contract_state(2, 401, 501))]),
+        },
+        contracts_trie_storage_proofs: HashMap::new(),
+    }
+}
+
+fn all_leaf_types_single_storage() -> StarknetForestProofs {
+    let address = ContractAddress::try_from(Felt::from(11)).unwrap();
+    StarknetForestProofs {
+        classes_trie_proof: PreimageMap::from([binary_preimage(120, 121, 122)]),
+        contracts_trie_proof: ContractsTrieProof {
+            nodes: PreimageMap::from([binary_preimage(320, 321, 322)]),
+            leaves: HashMap::from([(address, contract_state(3, 402, 502))]),
+        },
+        contracts_trie_storage_proofs: HashMap::from([(
+            address,
+            PreimageMap::from([edge_preimage(410, 411, 0, 1)]),
+        )]),
+    }
+}
+
+fn all_leaf_types_multiple_storages() -> StarknetForestProofs {
+    let address_a = ContractAddress::try_from(Felt::from(13)).unwrap();
+    let address_b = ContractAddress::try_from(Felt::from(14)).unwrap();
+    StarknetForestProofs {
+        classes_trie_proof: PreimageMap::from([edge_preimage(130, 131, 1, 1)]),
+        contracts_trie_proof: ContractsTrieProof {
+            nodes: PreimageMap::from([binary_preimage(330, 331, 332)]),
+            leaves: HashMap::from([
+                (address_a, contract_state(4, 403, 503)),
+                (address_b, contract_state(5, 404, 504)),
+            ]),
+        },
+        contracts_trie_storage_proofs: HashMap::from([
+            (address_a, PreimageMap::from([binary_preimage(420, 421, 422)])),
+            (address_b, PreimageMap::from([edge_preimage(430, 431, 1, 1)])),
+        ]),
+    }
+}
+
+#[rstest]
+#[case::only_classes(only_classes())]
+#[case::only_contracts(only_contracts())]
+#[case::classes_and_contracts(classes_and_contracts())]
+#[case::all_leaf_types_single_storage(all_leaf_types_single_storage())]
+#[case::all_leaf_types_multiple_storages(all_leaf_types_multiple_storages())]
+fn test_starknet_forest_proofs_serialization_round_trip(#[case] proofs: StarknetForestProofs) {
+    let encoded = proofs.serialize().unwrap();
+    let decoded = StarknetForestProofs::deserialize(&encoded).unwrap();
+    assert_eq!(proofs, decoded);
+}

crates/starknet_committer/src/patricia_merkle_tree/types.rs

@@ -36,19 +36,21 @@ pub type ClassesTrie = FilledTreeImpl<CompiledClassHash>;
 pub type ContractsTrie = FilledTreeImpl<ContractState>;
 pub type StorageTrieMap = HashMap<ContractAddress, StorageTrie>;
 
+#[derive(Debug, PartialEq)]
 pub struct ContractsTrieProof {
     pub nodes: PreimageMap,
     pub leaves: HashMap<ContractAddress, ContractState>,
 }
 
+#[derive(Debug, PartialEq)]
 pub struct StarknetForestProofs {
     pub classes_trie_proof: PreimageMap,
     pub contracts_trie_proof: ContractsTrieProof,
     pub contracts_trie_storage_proofs: HashMap<ContractAddress, PreimageMap>,
 }
 
 impl StarknetForestProofs {
-    pub(crate) fn extend(&mut self, other: Self) {
+    pub fn extend(&mut self, other: Self) {
         self.classes_trie_proof.extend(other.classes_trie_proof);
         self.contracts_trie_proof.nodes.extend(other.contracts_trie_proof.nodes);
         self.contracts_trie_proof.leaves.extend(other.contracts_trie_proof.leaves);