starknet_committer: add patricia paths forest reader/writer traits

Author: ArielElp

crates/starknet_committer/src/db/forest_trait.rs

@@ -30,11 +30,18 @@ use crate::forest::original_skeleton_forest::{ForestSortedIndices, OriginalSkele
 use crate::patricia_merkle_tree::leaf::leaf_impl::ContractState;
 use crate::patricia_merkle_tree::types::CompiledClassHash;
 
+#[cfg(feature = "os_input")]
+#[path = "forest_trait_witnesses.rs"]
+pub mod forest_trait_witnesses;
+
 #[derive(Clone, Debug, Deserialize, Eq, PartialEq, Hash, Serialize)]
 pub enum ForestMetadataType {
     CommitmentOffset,
     StateDiffHash(DbBlockNumber),
     StateRoot(DbBlockNumber),
+    /// BLAKE2s digest of the canonical accessed-keys set for the block.
+    #[cfg(feature = "os_input")]
+    AccessedKeysDigest(DbBlockNumber),
 }
 
 #[async_trait]

crates/starknet_committer/src/db/forest_trait_witnesses.rs

@@ -0,0 +1,85 @@
+use std::collections::HashMap;
+
+use async_trait::async_trait;
+use starknet_api::block::BlockNumber;
+use starknet_api::hash::HashOutput;
+use starknet_patricia::patricia_merkle_tree::traversal::TraversalResult;
+use starknet_patricia::patricia_merkle_tree::types::NodeIndex;
+use starknet_patricia_storage::errors::SerializationResult;
+use starknet_patricia_storage::storage_trait::{DbHashMap, DbValue};
+
+use super::{
+    EmptyInitialReadContext,
+    ForestMetadataType,
+    ForestReader,
+    ForestWriterWithMetadata,
+    StorageInitializer,
+};
+use crate::forest::deleted_nodes::DeletedNodes;
+use crate::forest::filled_forest::FilledForest;
+use crate::forest::forest_errors::ForestResult;
+use crate::patricia_merkle_tree::tree::SortedLeafIndices;
+use crate::patricia_merkle_tree::types::StarknetForestProofs;
+
+/// The information required to write Patricia proofs to the database.
+pub struct PatriciaProofsWrite {
+    pub block_number: BlockNumber,
+    pub keys_digest: [u8; 32],
+    pub witnesses: StarknetForestProofs,
+}
+
+/// Patricia proofs DB operation, which can be either delete or write.
+/// Expected by [ForestWriterWithMetadataAndWitnesses::write_with_metadata_and_witnesses], which
+/// accumulates all DB operations to guarantee atomicity.
+pub enum PatriciaProofsUpdate {
+    Write(PatriciaProofsWrite),
+    Delete(BlockNumber),
+}
+
+/// Reads committed OS-input witness payload (structured [`StarknetForestProofs`]) for a block
+/// height.
+#[async_trait]
+pub trait ForestReaderWithWitnesses:
+    ForestReader<InitialReadContext: EmptyInitialReadContext> + Send
+{
+    async fn read_witnesses(
+        &mut self,
+        height: BlockNumber,
+    ) -> ForestResult<Option<StarknetForestProofs>>;
+
+    /// Fetches Patricia witness paths for OS input, optionally staging serialized trie node KVs on
+    /// an in-memory overlay so reads match post-commit state before the forest is persisted.
+    async fn fetch_patricia_witnesses(
+        &mut self,
+        classes_trie_root_hash: HashOutput,
+        contracts_trie_root_hash: HashOutput,
+        class_sorted_leaf_indices: SortedLeafIndices<'_>,
+        contract_sorted_leaf_indices: SortedLeafIndices<'_>,
+        contract_storage_sorted_leaf_indices: &HashMap<NodeIndex, SortedLeafIndices<'_>>,
+        staged_serialized_forest: Option<DbHashMap>,
+    ) -> TraversalResult<StarknetForestProofs>;
+}
+
+/// Writes forest + metadata + deleted nodes, and applies [`PatriciaProofsUpdate`] in the same
+/// batch.
+#[async_trait]
+pub trait ForestWriterWithMetadataAndWitnesses: ForestWriterWithMetadata + Send {
+    async fn write_with_metadata_and_witnesses(
+        &mut self,
+        filled_forest: &FilledForest,
+        metadata: HashMap<ForestMetadataType, DbValue>,
+        deleted_nodes: DeletedNodes,
+        patricia_proofs_update: PatriciaProofsUpdate,
+    ) -> SerializationResult<usize>;
+}
+
+/// Forest storage with empty [`ForestReader::InitialReadContext`] plus OS-input witness read/write.
+pub trait ForestStorageWithWitnesses:
+    ForestReaderWithWitnesses + ForestWriterWithMetadataAndWitnesses + StorageInitializer
+{
+}
+
+impl<T> ForestStorageWithWitnesses for T where
+    T: ForestReaderWithWitnesses + ForestWriterWithMetadataAndWitnesses + StorageInitializer
+{
+}

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

@@ -50,6 +50,7 @@ use crate::db::index_db::types::{
     IndexLayoutSubTree,
     IndexNodeContext,
 };
+use crate::db::serde_db_utils::DbBlockNumber;
 use crate::forest::deleted_nodes::DeletedNodes;
 use crate::forest::filled_forest::FilledForest;
 use crate::forest::forest_errors::ForestResult;
@@ -84,6 +85,11 @@ static STATE_ROOT_METADATA_PREFIX: LazyLock<[u8; 32]> = LazyLock::new(|| {
     (Felt::from_bytes_be(&STATE_DIFF_HASH_METADATA_PREFIX) + Felt::ONE).to_bytes_be()
 });
 
+/// Prefix for accessed-keys digest metadata (committed per block).
+#[cfg_attr(not(feature = "os_input"), expect(dead_code))]
+pub(crate) static ACCESSED_KEYS_DIGEST_METADATA_PREFIX: LazyLock<[u8; 32]> =
+    LazyLock::new(|| (Felt::from_bytes_be(&STATE_ROOT_METADATA_PREFIX) + Felt::ONE).to_bytes_be());
+
 pub struct IndexDb<S: Storage, H = TreeHashFunctionImpl> {
     storage: S,
     phantom: PhantomData<H>,
@@ -267,28 +273,26 @@ impl<S: Storage> ForestWriter for IndexDb<S> {
 #[async_trait]
 impl<S: Storage> ForestMetadata for IndexDb<S> {
     fn metadata_key(metadata_type: ForestMetadataType) -> DbKey {
-        let mut key = Vec::with_capacity(64);
-        match metadata_type {
-            // Padding to 64byte keys to keep the 32byte prefix aligned between metadata and
-            // patricia nodes.
+        // Padding to 64byte keys to keep the 32byte prefix aligned between metadata and
+        // patricia nodes.
+        DbKey(match metadata_type {
             ForestMetadataType::CommitmentOffset => {
+                let mut key = Vec::with_capacity(64);
                 key.extend_from_slice(&*COMMITMENT_OFFSET_METADATA_PREFIX);
                 key.extend_from_slice(&[0u8; 32]);
+                key
             }
             ForestMetadataType::StateDiffHash(block_number) => {
-                key.extend_from_slice(&*STATE_DIFF_HASH_METADATA_PREFIX);
-                let block_number_bytes: [u8; 8] = block_number.serialize();
-                key.extend_from_slice(&block_number_bytes);
-                key.extend_from_slice(&[0u8; 24]);
+                metadata_block_number_key(&STATE_DIFF_HASH_METADATA_PREFIX, block_number)
             }
             ForestMetadataType::StateRoot(block_number) => {
-                key.extend_from_slice(&*STATE_ROOT_METADATA_PREFIX);
-                let block_number_bytes: [u8; 8] = block_number.serialize();
-                key.extend_from_slice(&block_number_bytes);
-                key.extend_from_slice(&[0u8; 24]);
+                metadata_block_number_key(&STATE_ROOT_METADATA_PREFIX, block_number)
+            }
+            #[cfg(feature = "os_input")]
+            ForestMetadataType::AccessedKeysDigest(block_number) => {
+                metadata_block_number_key(&ACCESSED_KEYS_DIGEST_METADATA_PREFIX, block_number)
             }
-        }
-        DbKey(key)
+        })
     }
 
     async fn get_from_storage(&mut self, db_key: DbKey) -> ForestResult<Option<DbValue>> {
@@ -320,6 +324,14 @@ impl<S: Storage> ForestWriterWithMetadata for IndexDb<S> {
     }
 }
 
+fn metadata_block_number_key(prefix: &[u8; 32], block_number: DbBlockNumber) -> Vec<u8> {
+    let mut key = Vec::with_capacity(64);
+    key.extend_from_slice(prefix);
+    key.extend_from_slice(&block_number.serialize());
+    key.extend_from_slice(&[0u8; 24]);
+    key
+}
+
 fn extract_root_hash<L: Leaf>(root: &Option<DbValue>) -> Result<HashOutput, DeserializationError>
 where
     TreeHashFunctionImpl: TreeHashFunction<L>,