apollo_committer: add new request handler

Author: ArielElp

crates/apollo_committer/Cargo.toml

@@ -7,8 +7,8 @@ license.workspace = true
 description = "State root commitment computation component for the Starknet sequencer."
 
 [features]
-testing = []
 os_input = ["apollo_committer_types/os_input", "starknet_committer/os_input"]
+testing = []
 
 [dependencies]
 apollo_committer_config.workspace = true
@@ -17,7 +17,7 @@ apollo_infra.workspace = true
 apollo_metrics.workspace = true
 async-trait.workspace = true
 starknet_api.workspace = true
-starknet_committer.workspace = true
+starknet_committer = { workspace = true, features = ["os_input"] }
 starknet_patricia_storage = { workspace = true, features = ["rocksdb_storage"] }
 tracing.workspace = true
 

crates/apollo_committer/src/communication.rs

@@ -2,17 +2,27 @@ use apollo_committer_types::communication::{CommitterRequest, CommitterResponse}
 use apollo_infra::component_definitions::ComponentRequestHandler;
 use apollo_infra::component_server::{LocalComponentServer, RemoteComponentServer};
 use async_trait::async_trait;
+#[cfg(feature = "os_input")]
+use starknet_committer::db::forest_trait::forest_trait_witnesses::ForestStorageWithWitnesses;
+#[cfg(not(feature = "os_input"))]
+use starknet_committer::db::forest_trait::forest_trait_witnesses::ForestWriterWithMetadataAndWitnesses;
+#[cfg(not(feature = "os_input"))]
 use starknet_committer::db::forest_trait::ForestStorageWithEmptyReadContext;
+#[cfg(feature = "os_input")]
+use starknet_patricia_storage::storage_trait::ImmutableReadOnlyStorage;
 
 use crate::committer::{ApolloCommitter, Committer, StorageConstructor};
 
 pub type LocalCommitterServer =
     LocalComponentServer<ApolloCommitter, CommitterRequest, CommitterResponse>;
 pub type RemoteCommitterServer = RemoteComponentServer<CommitterRequest, CommitterResponse>;
 
+#[cfg(not(feature = "os_input"))]
 #[async_trait]
-impl<S: StorageConstructor, ForestDB: ForestStorageWithEmptyReadContext<Storage = S>>
-    ComponentRequestHandler<CommitterRequest, CommitterResponse> for Committer<S, ForestDB>
+impl<S: StorageConstructor, ForestDB> ComponentRequestHandler<CommitterRequest, CommitterResponse>
+    for Committer<S, ForestDB>
+where
+    ForestDB: ForestStorageWithEmptyReadContext<Storage = S> + ForestWriterWithMetadataAndWitnesses,
 {
     async fn handle_request(&mut self, request: CommitterRequest) -> CommitterResponse {
         match request {
@@ -25,3 +35,28 @@ impl<S: StorageConstructor, ForestDB: ForestStorageWithEmptyReadContext<Storage
         }
     }
 }
+
+#[cfg(feature = "os_input")]
+#[async_trait]
+impl<S, ForestDB> ComponentRequestHandler<CommitterRequest, CommitterResponse>
+    for Committer<S, ForestDB>
+where
+    S: StorageConstructor + ImmutableReadOnlyStorage + 'static,
+    ForestDB: ForestStorageWithWitnesses<Storage = S>,
+{
+    async fn handle_request(&mut self, request: CommitterRequest) -> CommitterResponse {
+        match request {
+            CommitterRequest::CommitBlock(commit_block_request) => {
+                CommitterResponse::CommitBlock(self.commit_block(commit_block_request).await)
+            }
+            CommitterRequest::RevertBlock(revert_block_request) => {
+                CommitterResponse::RevertBlock(self.revert_block(revert_block_request).await)
+            }
+            CommitterRequest::ReadPathsAndCommitBlock(req) => {
+                CommitterResponse::ReadPathsAndCommitBlock(
+                    self.read_paths_and_commit_block(req).await,
+                )
+            }
+        }
+    }
+}

crates/apollo_committer/src/request_paths_and_commit_block_tests.rs

@@ -1,6 +1,7 @@
-use std::collections::HashMap;
+use std::collections::{BTreeSet, HashMap};
 
 use apollo_committer_types::committer_types::{
+    AccessedKeys,
     CommitBlockRequest,
     ReadPathsAndCommitBlockRequest,
     RevertBlockRequest,
@@ -13,16 +14,12 @@ use starknet_api::hash::HashOutput;
 use starknet_api::state::ThinStateDiff;
 use starknet_committer::block_committer::input::{
     contract_address_into_node_index,
-    AccessedKeys,
     StarknetStorageKey,
     StarknetStorageValue,
 };
 use starknet_committer::db::facts_db::db::{FactDbFilledNode, FactsNodeLayout};
-use starknet_committer::db::forest_trait::{
-    EmptyInitialReadContext,
-    ForestReader,
-    ForestReaderWithWitnesses,
-};
+use starknet_committer::db::forest_trait::forest_trait_witnesses::ForestReaderWithWitnesses;
+use starknet_committer::db::forest_trait::{EmptyInitialReadContext, ForestReader};
 use starknet_committer::db::index_db::IndexDbReadContext;
 use starknet_committer::db::trie_traversal::fetch_patricia_paths;
 use starknet_committer::hash_function::hash::TreeHashFunctionImpl;
@@ -80,12 +77,12 @@ fn accessed_storage_key_2() -> StarknetStorageKey {
 
 fn accessed_keys() -> AccessedKeys {
     AccessedKeys {
-        class_hashes: vec![accessed_class_hash()],
-        contract_addresses: vec![accessed_contract_1(), accessed_contract_2()],
-        contract_storage_keys: HashMap::from([
-            (accessed_contract_1(), vec![accessed_storage_key_1()]),
-            (accessed_contract_2(), vec![accessed_storage_key_2()]),
+        storage_keys: BTreeSet::from([
+            (accessed_contract_1(), accessed_storage_key_1().0),
+            (accessed_contract_2(), accessed_storage_key_2().0),
         ]),
+        accessed_contracts: BTreeSet::from([accessed_contract_1(), accessed_contract_2()]),
+        accessed_class_hashes: BTreeSet::from([accessed_class_hash()]),
     }
 }
 
@@ -447,7 +444,7 @@ async fn verify_witness_patricia_paths(
     let contract_root_hashes = [HashOutput::ROOT_OF_EMPTY_TREE, contracts_trie_root];
 
     let mut class_leaf_indices: Vec<NodeIndex> =
-        accessed_keys.class_hashes.iter().map(class_hash_into_node_index).collect();
+        accessed_keys.accessed_class_hashes.iter().map(class_hash_into_node_index).collect();
     verify_preimage_map_paths_exist::<CommitterCompiledClassHash>(
         &patricia_proofs.classes_trie_proof,
         &class_root_hashes,
@@ -458,7 +455,7 @@ async fn verify_witness_patricia_paths(
     .await;
 
     let mut contract_leaf_indices: Vec<NodeIndex> =
-        accessed_keys.contract_addresses.iter().map(contract_address_into_node_index).collect();
+        accessed_keys.accessed_contracts.iter().map(contract_address_into_node_index).collect();
     verify_preimage_map_paths_exist::<ContractState>(
         &patricia_proofs.contracts_trie_proof.nodes,
         &contract_root_hashes,
@@ -471,7 +468,7 @@ async fn verify_witness_patricia_paths(
     )
     .await;
 
-    for contract_address in &accessed_keys.contract_addresses {
+    for contract_address in &accessed_keys.accessed_contracts {
         let storage_proof = patricia_proofs
             .contracts_trie_storage_proofs
             .get(contract_address)
@@ -481,10 +478,13 @@ async fn verify_witness_patricia_paths(
             .leaves
             .get(contract_address)
             .unwrap_or_else(|| panic!("missing contracts trie leaf for {contract_address:?}"));
-        let storage_keys = accessed_keys
-            .contract_storage_keys
-            .get(contract_address)
-            .unwrap_or_else(|| panic!("missing accessed storage keys for {contract_address:?}"));
+        let storage_keys: Vec<_> = accessed_keys
+            .storage_keys
+            .iter()
+            .filter(|(address, _)| address == contract_address)
+            .map(|(_, key)| StarknetStorageKey(*key))
+            .collect();
+        assert!(!storage_keys.is_empty(), "missing accessed storage keys for {contract_address:?}");
         let mut storage_leaf_indices: Vec<NodeIndex> =
             storage_keys.iter().map(NodeIndex::from).collect();
         let storage_root_hashes =

crates/apollo_committer_types/src/communication.rs

@@ -21,6 +21,8 @@ use crate::committer_types::{
     RevertBlockRequest,
     RevertBlockResponse,
 };
+#[cfg(feature = "os_input")]
+use crate::committer_types::{ReadPathsAndCommitBlockRequest, ReadPathsAndCommitBlockResponse};
 use crate::errors::{CommitterClientError, CommitterClientResult, CommitterResult};
 
 pub type LocalCommitterClient = LocalComponentClient<CommitterRequest, CommitterResponse>;
@@ -43,6 +45,14 @@ pub trait CommitterClient: Send + Sync {
         &self,
         input: RevertBlockRequest,
     ) -> CommitterClientResult<RevertBlockResponse>;
+
+    #[cfg(feature = "os_input")]
+    /// Applies the state diff, collects merged Patricia witnesses for OS input, and persists replay
+    /// data (digest + payload).
+    async fn read_paths_and_commit_block(
+        &self,
+        input: ReadPathsAndCommitBlockRequest,
+    ) -> CommitterClientResult<ReadPathsAndCommitBlockResponse>;
 }
 
 #[derive(Serialize, Deserialize, Clone, AsRefStr, EnumDiscriminants)]
@@ -54,6 +64,8 @@ pub trait CommitterClient: Send + Sync {
 pub enum CommitterRequest {
     CommitBlock(CommitBlockRequest),
     RevertBlock(RevertBlockRequest),
+    #[cfg(feature = "os_input")]
+    ReadPathsAndCommitBlock(ReadPathsAndCommitBlockRequest),
 }
 
 impl_debug_for_infra_requests_and_responses!(CommitterRequest);
@@ -64,6 +76,8 @@ impl PrioritizedRequest for CommitterRequest {}
 pub enum CommitterResponse {
     CommitBlock(CommitterResult<CommitBlockResponse>),
     RevertBlock(CommitterResult<RevertBlockResponse>),
+    #[cfg(feature = "os_input")]
+    ReadPathsAndCommitBlock(CommitterResult<ReadPathsAndCommitBlockResponse>),
 }
 
 impl_debug_for_infra_requests_and_responses!(CommitterResponse);
@@ -109,4 +123,21 @@ where
             Direct
         )
     }
+
+    #[cfg(feature = "os_input")]
+    async fn read_paths_and_commit_block(
+        &self,
+        input: ReadPathsAndCommitBlockRequest,
+    ) -> CommitterClientResult<ReadPathsAndCommitBlockResponse> {
+        let request = CommitterRequest::ReadPathsAndCommitBlock(input);
+        handle_all_response_variants!(
+            self,
+            request,
+            CommitterResponse,
+            ReadPathsAndCommitBlock,
+            CommitterClientError,
+            CommitterError,
+            Direct
+        )
+    }
 }