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

Author: ArielElp

.github/workflows/apollo_storage_os_input_ci.yml

@@ -12,6 +12,7 @@ on:
       - ".github/workflows/apollo_storage_os_input_ci.yml"
       - "Cargo.lock"
       - "Cargo.toml"
+      - "crates/apollo_committer/**"
       - "crates/apollo_committer_types/**"
       - "crates/apollo_storage/**"
       - "crates/apollo_config/**"
@@ -61,6 +62,7 @@ jobs:
           github_token: ${{ secrets.GITHUB_TOKEN }}
       - run: cargo test -p starknet_committer --features os_input
       - run: cargo test -p apollo_committer_types --features os_input
+      - run: cargo test -p apollo_committer --features os_input
       - run: cargo build -p apollo_batcher --features os_input
       - run: cargo test -p apollo_batcher --features os_input
       - run: cargo test -p apollo_reverts --features os_input

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,
@@ -14,11 +20,15 @@ use apollo_infra::component_definitions::{default_component_start_fn, ComponentS
 use async_trait::async_trait;
 use starknet_api::block::BlockNumber;
 use starknet_api::block_hash::state_diff_hash::calculate_state_diff_hash;
+#[cfg(feature = "os_input")]
+use starknet_api::core::ContractAddress;
 use starknet_api::core::{GlobalRoot, StateDiffCommitment};
 use starknet_api::hash::PoseidonHash;
 use starknet_api::state::ThinStateDiff;
 use starknet_committer::block_committer::commit::commit_block;
 use starknet_committer::block_committer::input::Input;
+#[cfg(feature = "os_input")]
+use starknet_committer::block_committer::input::StarknetStorageKey;
 use starknet_committer::block_committer::measurements_util::{
     Action,
     BlockDurations,
@@ -27,21 +37,35 @@ use starknet_committer::block_committer::measurements_util::{
     MeasurementsTrait,
     SingleBlockMeasurements,
 };
+#[cfg(feature = "os_input")]
+use starknet_committer::db::forest_trait::forest_trait_witnesses::{
+    ForestStorageWithWitnesses,
+    PatriciaProofsUpdate,
+    PatriciaProofsWrite,
+};
 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};
 
@@ -447,9 +471,192 @@ where
     }
 
     fn map_internal_error<E: Error>(&self, err: E) -> CommitterError {
+        self.map_internal_error_at_height(self.offset, err)
+    }
+
+    fn map_internal_error_at_height<E: Error>(
+        &self,
+        height: BlockNumber,
+        err: E,
+    ) -> CommitterError {
         let error_message = format!("{err:?}: {err}");
-        error!("Error committing block number {0}. {error_message}.", self.offset);
-        CommitterError::Internal { height: self.offset, message: error_message }
+        error!("Error committing block number {height}. {error_message}.");
+        CommitterError::Internal { height, message: error_message }
+    }
+}
+
+#[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 { storage_keys, accessed_contracts, accessed_class_hashes },
+        }: ReadPathsAndCommitBlockRequest,
+    ) -> CommitterResult<ReadPathsAndCommitBlockResponse> {
+        let class_hashes: Vec<_> = accessed_class_hashes.iter().copied().collect();
+        let contract_addresses: Vec<_> = accessed_contracts.iter().copied().collect();
+        let contract_storage_keys = storage_keys.iter().fold(
+            HashMap::<ContractAddress, Vec<StarknetStorageKey>>::new(),
+            |mut accumulator, (address, key)| {
+                accumulator.entry(*address).or_default().push(StarknetStorageKey(*key));
+                accumulator
+            },
+        );
+        let mut leaves_request = LeavesRequest::from_accessed_leaves(
+            &class_hashes,
+            &contract_addresses,
+            &contract_storage_keys,
+        );
+        info!(
+            "read_paths_and_commit_block: height {height}, accessed keys len {}, state diff len {}",
+            leaves_request.total_leaf_count(),
+            state_diff.len(),
+        );
+        let sorted_leaves: SortedLeavesRequest<'_> = (&mut leaves_request).into();
+        let digest = accessed_keys_digest(&sorted_leaves);
+
+        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(|error| self.map_internal_error_at_height(height, error))?;
+                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);
+                let witness_node_count = patricia_proofs.classes_trie_proof.len()
+                    + patricia_proofs.contracts_trie_proof.nodes.len()
+                    + patricia_proofs.contracts_trie_proof.leaves.len()
+                    + patricia_proofs
+                        .contracts_trie_storage_proofs
+                        .values()
+                        .map(|proof| proof.len())
+                        .sum::<usize>();
+                info!(
+                    "For block number {height}, writing filled forest and {witness_node_count} \
+                     witness nodes 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,
+                        PatriciaProofsUpdate::Write(PatriciaProofsWrite {
+                            block_number: 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]>> {
+        let digest_raw = self
+            .forest_storage
+            .read_metadata(ForestMetadataType::AccessedKeysDigest(DbBlockNumber(block_number)))
+            .await
+            .map_err(|error| self.map_internal_error_at_height(block_number, error))?;
+
+        digest_raw
+            .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()
     }
 }
 

crates/starknet_committer/src/patricia_merkle_tree/tree.rs

@@ -91,6 +91,16 @@ impl LeavesRequest {
             contract_storage_leaf_indices,
         }
     }
+
+    /// Total number of trie leaves requested (classes, contracts, and storage slots).
+    pub fn total_leaf_count(&self) -> usize {
+        self.class_leaf_indices.len()
+            + self.contract_leaf_indices.len()
+            + self
+                .contract_storage_leaf_indices
+                .values()
+                .fold(0, |count, leaf_indices| count + leaf_indices.len())
+    }
 }
 
 impl<'a> From<&'a mut LeavesRequest> for SortedLeavesRequest<'a> {

crates/starknet_committer/src/patricia_merkle_tree/types.rs

@@ -43,13 +43,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,