starknet_committer: underlying logic of the new read witnesses and commit endpoint

Author: ArielElp

crates/apollo_committer/Cargo.toml

@@ -8,6 +8,7 @@ description = "State root commitment computation component for the Starknet sequ
 
 [features]
 testing = []
+os_input = ["apollo_committer_types/os_input", "starknet_committer/os_input"]
 
 [dependencies]
 apollo_committer_config.workspace = true

crates/apollo_committer/src/committer.rs

@@ -3,6 +3,12 @@ use std::error::Error;
 use std::path::PathBuf;
 
 use apollo_committer_config::config::{ApolloStorage, CommitterConfig};
+#[cfg(feature = "os_input")]
+use apollo_committer_types::committer_types::{
+    AccessedKeys,
+    ReadPathsAndCommitBlockRequest,
+    ReadPathsAndCommitBlockResponse,
+};
 use apollo_committer_types::committer_types::{
     CommitBlockRequest,
     CommitBlockResponse,
@@ -27,21 +33,33 @@ use starknet_committer::block_committer::measurements_util::{
     MeasurementsTrait,
     SingleBlockMeasurements,
 };
+#[cfg(feature = "os_input")]
+use starknet_committer::db::forest_trait::ForestStorageWithWitnesses;
+#[cfg(feature = "os_input")]
+use starknet_committer::db::forest_trait::PatriciaProofsUpdates;
 use starknet_committer::db::forest_trait::{
     EmptyInitialReadContext,
     ForestMetadataType,
     ForestStorageWithEmptyReadContext,
 };
 use starknet_committer::db::index_db::IndexDb;
+#[cfg(feature = "os_input")]
+use starknet_committer::db::serde_db_utils::accessed_keys_digest;
 use starknet_committer::db::serde_db_utils::{
     deserialize_felt_no_packing,
     serialize_felt_no_packing,
     DbBlockNumber,
 };
 use starknet_committer::forest::deleted_nodes::DeletedNodes;
 use starknet_committer::forest::filled_forest::FilledForest;
+#[cfg(feature = "os_input")]
+use starknet_committer::patricia_merkle_tree::tree::{LeavesRequest, SortedLeavesRequest};
+#[cfg(feature = "os_input")]
+use starknet_patricia_storage::errors::SerializationError;
 use starknet_patricia_storage::map_storage::CachedStorage;
 use starknet_patricia_storage::rocksdb_storage::RocksDbStorage;
+#[cfg(feature = "os_input")]
+use starknet_patricia_storage::storage_trait::ImmutableReadOnlyStorage;
 use starknet_patricia_storage::storage_trait::{DbValue, Storage};
 use tracing::{debug, error, info, warn};
 
@@ -453,6 +471,162 @@ where
     }
 }
 
+#[cfg(feature = "os_input")]
+impl<S, ForestDB> Committer<S, ForestDB>
+where
+    S: StorageConstructor + ImmutableReadOnlyStorage + 'static,
+    ForestDB: ForestStorageWithWitnesses<Storage = S>,
+{
+    /// Commits the next block and returns merged Patricia witness facts for OS input, persisting
+    /// digest + payload for idempotent replay.
+    pub async fn read_paths_and_commit_block(
+        &mut self,
+        ReadPathsAndCommitBlockRequest {
+            commit: CommitBlockRequest { state_diff, state_diff_commitment, height },
+            accessed_keys: AccessedKeys { class_hashes, contract_addresses, contract_storage_keys },
+        }: ReadPathsAndCommitBlockRequest,
+    ) -> CommitterResult<ReadPathsAndCommitBlockResponse> {
+        let mut leaves_request = LeavesRequest::from_accessed_leaves(
+            &class_hashes,
+            &contract_addresses,
+            &contract_storage_keys,
+        );
+        let sorted_leaves: SortedLeavesRequest<'_> = (&mut leaves_request).into();
+        let digest = accessed_keys_digest(&sorted_leaves);
+        info!(
+            "read_paths_and_commit_block: block {height} with {} class hashes, {} contract \
+             addresses",
+            class_hashes.len(),
+            contract_addresses.len()
+        );
+
+        match self.commit_or_load(&state_diff, state_diff_commitment, height).await? {
+            CommitBlockHeightPlan::Historical { global_root } => {
+                let stored_digest = self.load_witnesses_digest(height).await?;
+                if stored_digest != Some(digest) {
+                    return Err(CommitterError::AccessedKeysDigestMismatch {
+                        height,
+                        stored: stored_digest,
+                        expected: digest,
+                    });
+                }
+                let proofs = self
+                    .forest_storage
+                    .read_witnesses(height)
+                    .await
+                    .map_err(|e| self.map_internal_error(e))?;
+                let proofs = proofs.ok_or(CommitterError::MissingPatriciaPaths { height })?;
+                Ok(ReadPathsAndCommitBlockResponse { global_root, patricia_proofs: proofs })
+            }
+            // Flow overview:
+            // 1. Fetch patricia paths for the accessed keys
+            // 2. Compute the updates from the state diff (commit) but avoid updating the underlying
+            //    DB in order to guarantee atomicity.
+            // 3. Fetch patricia paths for the post-commit tries, via running step 1 against a two
+            //    layer storage composed from the underlying storage and the modifications from 2.
+            // 4. Merge the two sets of patricia paths and write the result to the storage.
+            // 5. Update the commitment offset and return the global root and the patricia proofs.
+            CommitBlockHeightPlan::CommitTip { state_diff_commitment } => {
+                let pre_roots = self
+                    .forest_storage
+                    .read_roots(ForestDB::InitialReadContext::create_empty())
+                    .await
+                    .map_err(|e| self.map_internal_error(e))?;
+                let mut patricia_proofs = self
+                    .forest_storage
+                    .fetch_patricia_witnesses(
+                        pre_roots.classes_trie_root_hash,
+                        pre_roots.contracts_trie_root_hash,
+                        sorted_leaves.class_sorted,
+                        sorted_leaves.contract_sorted,
+                        &sorted_leaves.storage_sorted,
+                        None,
+                    )
+                    .await
+                    .map_err(|e| CommitterError::PatriciaPathsCollectionFailed {
+                        height,
+                        message: format!("pre-commit witness paths: {e:?}"),
+                    })?;
+
+                let mut block_measurements = SingleBlockMeasurements::default();
+                block_measurements.start_measurement(Action::EndToEnd);
+                let CommitStateDiffOutput { filled_forest, global_root, deleted_nodes } =
+                    self.commit_state_diff(state_diff, &mut block_measurements).await?;
+                let post_roots = filled_forest.state_roots();
+
+                let forest_updates = ForestDB::serialize_forest(&filled_forest)
+                    .map_err(|e| self.map_internal_error(e))?;
+
+                let proof_after = self
+                    .forest_storage
+                    .fetch_patricia_witnesses(
+                        post_roots.classes_trie_root_hash,
+                        post_roots.contracts_trie_root_hash,
+                        sorted_leaves.class_sorted,
+                        sorted_leaves.contract_sorted,
+                        &sorted_leaves.storage_sorted,
+                        Some(forest_updates),
+                    )
+                    .await
+                    .map_err(|e| CommitterError::PatriciaPathsCollectionFailed {
+                        height,
+                        message: format!("post-commit witness paths: {e:?}"),
+                    })?;
+
+                patricia_proofs.extend(proof_after);
+
+                let (metadata, next_offset) =
+                    commit_tip_metadata_bundle(height, global_root, state_diff_commitment);
+                info!(
+                    "For block number {height}, writing filled forest and witnesses to storage \
+                     with metadata: {metadata:?}, delete {} nodes",
+                    deleted_nodes.len()
+                );
+                block_measurements.start_measurement(Action::Write);
+                let n_write_entries = self
+                    .forest_storage
+                    .write_with_metadata_and_witnesses(
+                        &filled_forest,
+                        metadata,
+                        deleted_nodes,
+                        PatriciaProofsUpdates::Set {
+                            height,
+                            keys_digest: digest,
+                            witnesses: patricia_proofs.clone(),
+                        },
+                    )
+                    .await
+                    .map_err(|e: SerializationError| self.map_internal_error(e))?;
+                block_measurements.attempt_to_stop_measurement(Action::Write, n_write_entries).ok();
+                block_measurements.attempt_to_stop_measurement(Action::EndToEnd, 0).ok();
+                update_metrics(height, &block_measurements.block_measurement);
+                self.update_offset(next_offset);
+                Ok(ReadPathsAndCommitBlockResponse { global_root, patricia_proofs })
+            }
+        }
+    }
+
+    async fn load_witnesses_digest(
+        &mut self,
+        block_number: BlockNumber,
+    ) -> CommitterResult<Option<[u8; 32]>> {
+        self.forest_storage
+            .read_metadata(ForestMetadataType::AccessedKeysDigest(DbBlockNumber(block_number)))
+            .await
+            .map_err(|error| self.map_internal_error(error))?
+            .map(|digest_raw| {
+                digest_raw.0.as_slice().try_into().map_err(|_| CommitterError::Internal {
+                    height: block_number,
+                    message: format!(
+                        "Invalid OS witnesses digest length {} (expected 32)",
+                        digest_raw.0.len()
+                    ),
+                })
+            })
+            .transpose()
+    }
+}
+
 #[async_trait]
 impl ComponentStarter for ApolloCommitter {
     async fn start(&mut self) {

crates/apollo_committer/src/metrics.rs

@@ -1,6 +1,10 @@
 use apollo_committer_types::communication::COMMITTER_REQUEST_LABELS;
 use apollo_infra::metrics::{
-    InfraMetrics, LocalClientMetrics, LocalServerMetrics, RemoteClientMetrics, RemoteServerMetrics,
+    InfraMetrics,
+    LocalClientMetrics,
+    LocalServerMetrics,
+    RemoteClientMetrics,
+    RemoteServerMetrics,
 };
 use apollo_metrics::{define_infra_metrics, define_metrics};
 use starknet_api::block::BlockNumber;

crates/starknet_committer/src/patricia_merkle_tree/types.rs

@@ -62,13 +62,13 @@ pub type ClassesTrie = FilledTreeImpl<CompiledClassHash>;
 pub type ContractsTrie = FilledTreeImpl<ContractState>;
 pub type StorageTrieMap = HashMap<ContractAddress, StorageTrie>;
 
-#[derive(Debug, Clone, PartialEq)]
+#[derive(Clone, Debug, PartialEq)]
 pub struct ContractsTrieProof {
     pub nodes: PreimageMap,
     pub leaves: HashMap<ContractAddress, ContractState>,
 }
 
-#[derive(Debug, Clone, PartialEq)]
+#[derive(Clone, Debug, PartialEq)]
 pub struct StarknetForestProofs {
     pub classes_trie_proof: PreimageMap,
     pub contracts_trie_proof: ContractsTrieProof,