apollo_committer,starknet_committer: use blockifier's AccessedKeys

Author: ArielElp

Cargo.lock

@@ -3,32 +3,27 @@ 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,
     RevertBlockRequest,
     RevertBlockResponse,
 };
+#[cfg(feature = "os_input")]
+use apollo_committer_types::committer_types::{
+    ReadPathsAndCommitBlockRequest,
+    ReadPathsAndCommitBlockResponse,
+};
 use apollo_committer_types::errors::{CommitterError, CommitterResult};
 use apollo_infra::component_definitions::{default_component_start_fn, ComponentStarter};
 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,
@@ -59,7 +54,7 @@ use starknet_committer::db::serde_db_utils::{
 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};
+use starknet_committer::patricia_merkle_tree::tree::LeavesRequest;
 #[cfg(feature = "os_input")]
 use starknet_patricia_storage::errors::SerializationError;
 use starknet_patricia_storage::map_storage::CachedStorage;
@@ -512,29 +507,16 @@ where
         &mut self,
         ReadPathsAndCommitBlockRequest {
             commit: CommitBlockRequest { state_diff, state_diff_commitment, height },
-            accessed_keys: AccessedKeys { storage_keys, accessed_contracts, accessed_class_hashes },
+            accessed_keys,
         }: 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,
-        );
+        let mut leaves_request = LeavesRequest::from(&accessed_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 sorted_leaves = leaves_request.sorted();
         let digest = accessed_keys_digest(&sorted_leaves);
 
         match self.commit_or_load(&state_diff, state_diff_commitment, height).await? {

crates/apollo_committer/src/committer.rs

@@ -7,9 +7,7 @@ use starknet_api::block::BlockNumber;
 use starknet_api::core::{ClassHash, ContractAddress, BLOCK_HASH_TABLE_ADDRESS};
 use starknet_api::state::StorageKey;
 
-#[cfg(any(feature = "testing", test))]
-use super::cached_state::StateChangesKeys;
-use super::cached_state::{CommitmentStateDiff, StorageEntry};
+use super::cached_state::{CommitmentStateDiff, StateChangesKeys, StorageEntry};
 use super::stateful_compression::predicted_alias_storage_entries;
 use crate::blockifier_versioned_constants::VersionedConstants;
 use crate::transaction::objects::TransactionExecutionInfo;
@@ -47,6 +45,19 @@ impl From<AccessedKeys> for StateChangesKeys {
     }
 }
 
+impl From<StateChangesKeys> for AccessedKeys {
+    fn from(state_changes_keys: StateChangesKeys) -> Self {
+        Self {
+            storage_keys: state_changes_keys.storage_keys.into_iter().collect(),
+            accessed_contracts: state_changes_keys.modified_contracts.into_iter().collect(),
+            accessed_class_hashes: state_changes_keys
+                .compiled_class_hash_keys
+                .into_iter()
+                .collect(),
+        }
+    }
+}
+
 impl AccessedKeys {
     /// Builds the [`AccessedKeys`] the OS needs to read at the execution of a block.
     pub fn new<'a>(

crates/blockifier/src/state/accessed_keys.rs

@@ -15,6 +15,7 @@ apollo_config.workspace = true
 async-trait.workspace = true
 bincode = { workspace = true, optional = true }
 blake2 = { workspace = true, optional = true }
+blockifier.workspace = true
 derive_more = { workspace = true, features = ["as_ref", "from", "into"] }
 digest = { workspace = true, optional = true }
 ethnum.workspace = true

crates/starknet_committer/Cargo.toml

@@ -1,6 +1,7 @@
 use std::collections::HashMap;
 
-use starknet_api::core::{ClassHash, ContractAddress};
+use blockifier::state::accessed_keys::AccessedKeys;
+use starknet_api::core::ContractAddress;
 use starknet_api::hash::HashOutput;
 use starknet_patricia::patricia_merkle_tree::original_skeleton_tree::config::OriginalSkeletonTreeConfig;
 use starknet_patricia::patricia_merkle_tree::traversal::TraversalResult;
@@ -52,7 +53,7 @@ impl OriginalSkeletonTreeConfig for OriginalSkeletonTrieConfig {
 }
 
 /// Requested trie leaves for Patricia witness collection (classes trie, contracts trie, and
-/// per-contract storage leaves). Built via [`LeavesRequest::from_accessed_leaves`].
+/// per-contract storage leaves). Built via [`From`] from [`AccessedKeys`].
 #[derive(Clone)]
 pub struct LeavesRequest {
     pub class_leaf_indices: Vec<NodeIndex>,
@@ -67,31 +68,35 @@ pub struct SortedLeavesRequest<'a> {
     pub storage_sorted: HashMap<NodeIndex, SortedLeafIndices<'a>>,
 }
 
-impl LeavesRequest {
-    /// Builds index buffers expected by [`fetch_all_patricia_paths`].
-    pub fn from_accessed_leaves(
-        class_hashes: &[ClassHash],
-        contract_addresses: &[ContractAddress],
-        contract_storage_keys: &HashMap<ContractAddress, Vec<StarknetStorageKey>>,
-    ) -> Self {
-        let contract_leaf_indices: Vec<NodeIndex> =
-            contract_addresses.iter().map(contract_address_into_node_index).collect();
-        let contract_storage_leaf_indices: HashMap<NodeIndex, Vec<NodeIndex>> =
-            contract_storage_keys
-                .iter()
-                .map(|(address, keys)| {
-                    let node_index = contract_address_into_node_index(address);
-                    let leaf_indices: Vec<_> = keys.iter().map(NodeIndex::from).collect();
-                    (node_index, leaf_indices)
-                })
-                .collect();
+impl From<&AccessedKeys> for LeavesRequest {
+    fn from(accessed_keys: &AccessedKeys) -> Self {
+        let contract_storage_leaf_indices: HashMap<NodeIndex, Vec<NodeIndex>> = accessed_keys
+            .storage_keys
+            .iter()
+            .fold(HashMap::new(), |mut accumulator, (address, key)| {
+                accumulator
+                    .entry(contract_address_into_node_index(address))
+                    .or_default()
+                    .push(NodeIndex::from(&StarknetStorageKey(*key)));
+                accumulator
+            });
         Self {
-            class_leaf_indices: class_hashes.iter().map(class_hash_into_node_index).collect(),
-            contract_leaf_indices,
+            class_leaf_indices: accessed_keys
+                .accessed_class_hashes
+                .iter()
+                .map(class_hash_into_node_index)
+                .collect(),
+            contract_leaf_indices: accessed_keys
+                .accessed_contracts
+                .iter()
+                .map(contract_address_into_node_index)
+                .collect(),
             contract_storage_leaf_indices,
         }
     }
+}
 
+impl LeavesRequest {
     /// Total number of trie leaves requested (classes, contracts, and storage slots).
     pub fn total_leaf_count(&self) -> usize {
         self.class_leaf_indices.len()
@@ -101,18 +106,17 @@ impl LeavesRequest {
                 .values()
                 .fold(0, |count, leaf_indices| count + leaf_indices.len())
     }
-}
 
-impl<'a> From<&'a mut LeavesRequest> for SortedLeavesRequest<'a> {
-    fn from(leaves_request: &'a mut LeavesRequest) -> Self {
-        let class_sorted = SortedLeafIndices::new(&mut leaves_request.class_leaf_indices);
-        let contract_sorted = SortedLeafIndices::new(&mut leaves_request.contract_leaf_indices);
-        let storage_sorted: HashMap<_, _> = leaves_request
+    /// Sorts leaf indices in place and returns borrowed views for Patricia traversal.
+    pub fn sorted(&mut self) -> SortedLeavesRequest<'_> {
+        let class_sorted = SortedLeafIndices::new(&mut self.class_leaf_indices);
+        let contract_sorted = SortedLeafIndices::new(&mut self.contract_leaf_indices);
+        let storage_sorted: HashMap<_, _> = self
             .contract_storage_leaf_indices
             .iter_mut()
             .map(|(address, leaf_indices)| (*address, SortedLeafIndices::new(leaf_indices)))
             .collect();
-        Self { class_sorted, contract_sorted, storage_sorted }
+        SortedLeavesRequest { class_sorted, contract_sorted, storage_sorted }
     }
 }
 
@@ -241,18 +245,12 @@ pub async fn fetch_previous_and_new_patricia_paths(
     storage: &mut impl ReadOnlyStorage,
     classes_trie_root_hashes: RootHashes,
     contracts_trie_root_hashes: RootHashes,
-    class_hashes: &[ClassHash],
-    contract_addresses: &[ContractAddress],
-    contract_storage_keys: &HashMap<ContractAddress, Vec<StarknetStorageKey>>,
+    accessed_keys: &AccessedKeys,
 ) -> TraversalResult<StarknetForestProofs> {
-    let mut leaves_request = LeavesRequest::from_accessed_leaves(
-        class_hashes,
-        contract_addresses,
-        contract_storage_keys,
-    );
+    let mut leaves_request = LeavesRequest::from(accessed_keys);
 
     let SortedLeavesRequest { class_sorted, contract_sorted, storage_sorted } =
-        SortedLeavesRequest::from(&mut leaves_request);
+        leaves_request.sorted();
     let prev_proofs = fetch_all_patricia_paths::<FactsNodeLayout>(
         storage,
         classes_trie_root_hashes.previous_root_hash,

crates/starknet_committer/src/patricia_merkle_tree/tree.rs

@@ -81,7 +81,6 @@ expect-test.workspace = true
 itertools.workspace = true
 num-bigint.workspace = true
 rstest.workspace = true
-starknet_committer.workspace = true
 starknet_patricia = { workspace = true, features = ["testing"] }
 tempfile.workspace = true
 tokio = { workspace = true, features = ["rt-multi-thread"] }

crates/starknet_os/Cargo.toml

@@ -1,17 +1,16 @@
 use std::collections::HashMap;
 
-use blockifier::state::cached_state::StateChangesKeys;
-use starknet_api::core::{ClassHash, ContractAddress};
+use blockifier::state::accessed_keys::AccessedKeys;
+use starknet_api::core::ContractAddress;
 use starknet_api::hash::{HashOutput, StateRoots};
 use starknet_committer::block_committer::input::{
     contract_address_into_node_index,
     try_node_index_into_contract_address,
-    StarknetStorageKey,
 };
 use starknet_committer::db::facts_db::create_facts_tree::get_leaves;
 use starknet_committer::patricia_merkle_tree::leaf::leaf_impl::ContractState;
 use starknet_committer::patricia_merkle_tree::tree::fetch_previous_and_new_patricia_paths;
-use starknet_committer::patricia_merkle_tree::types::{RootHashes, StarknetForestProofs};
+use starknet_committer::patricia_merkle_tree::types::RootHashes;
 use starknet_patricia::patricia_merkle_tree::node_data::inner_node::flatten_preimages;
 use starknet_patricia::patricia_merkle_tree::original_skeleton_tree::errors::OriginalSkeletonTreeError;
 use starknet_patricia::patricia_merkle_tree::traversal::TraversalError;
@@ -74,10 +73,10 @@ impl StateCommitmentInfos {
         previous_state_roots: &StateRoots,
         new_state_roots: &StateRoots,
         commitments: &mut MapStorage,
-        initial_reads_keys: &StateChangesKeys,
+        accessed_keys: &AccessedKeys,
     ) -> Result<Self, CommitmentInfosError> {
         let addresses: Vec<ContractAddress> =
-            initial_reads_keys.modified_contracts.iter().copied().collect();
+            accessed_keys.accessed_contracts.iter().copied().collect();
 
         let previous_storage_roots = get_storage_roots(
             &addresses,
@@ -89,8 +88,7 @@ impl StateCommitmentInfos {
             get_storage_roots(&addresses, new_state_roots.contracts_trie_root_hash, commitments)
                 .await?;
 
-        let storage_proofs = fetch_storage_proofs_from_state_changes_keys(
-            initial_reads_keys,
+        let storage_proofs = fetch_previous_and_new_patricia_paths(
             commitments,
             RootHashes {
                 previous_root_hash: previous_state_roots.classes_trie_root_hash,
@@ -100,6 +98,7 @@ impl StateCommitmentInfos {
                 previous_root_hash: previous_state_roots.contracts_trie_root_hash,
                 new_root_hash: new_state_roots.contracts_trie_root_hash,
             },
+            accessed_keys,
         )
         .await?;
 
@@ -172,32 +171,3 @@ async fn get_storage_roots(
         })
         .collect()
 }
-
-async fn fetch_storage_proofs_from_state_changes_keys(
-    initial_reads_keys: &StateChangesKeys,
-    storage: &mut MapStorage,
-    classes_trie_root_hashes: RootHashes,
-    contracts_trie_root_hashes: RootHashes,
-) -> Result<StarknetForestProofs, CommitmentInfosError> {
-    let class_hashes: Vec<ClassHash> =
-        initial_reads_keys.compiled_class_hash_keys.iter().cloned().collect();
-    let contract_addresses =
-        &initial_reads_keys.modified_contracts.iter().cloned().collect::<Vec<_>>();
-    let contract_storage_keys = initial_reads_keys.storage_keys.iter().fold(
-        HashMap::<ContractAddress, Vec<StarknetStorageKey>>::new(),
-        |mut acc, (address, key)| {
-            acc.entry(*address).or_default().push(StarknetStorageKey(*key));
-            acc
-        },
-    );
-
-    Ok(fetch_previous_and_new_patricia_paths(
-        storage,
-        classes_trie_root_hashes,
-        contracts_trie_root_hashes,
-        &class_hashes,
-        contract_addresses,
-        &contract_storage_keys,
-    )
-    .await?)
-}

crates/starknet_os/src/commitment_infos.rs

@@ -944,7 +944,7 @@ impl<S: FlowTestState> TestBuilder<S> {
                 &previous_state_roots,
                 &new_state_roots,
                 &mut map_storage,
-                &accessed_keys.into(),
+                &accessed_keys,
             )
             .await
             .unwrap();

crates/starknet_os_flow_tests/src/test_manager.rs

@@ -1,6 +1,7 @@
 use std::collections::HashMap;
 
 use async_trait::async_trait;
+use blockifier::state::accessed_keys::AccessedKeys;
 use blockifier::state::cached_state::StateMaps;
 use indexmap::IndexMap;
 use serde::{Deserialize, Serialize};
@@ -343,15 +344,15 @@ impl RpcStorageProofsProvider {
         let mut map_storage: MapStorage = facts_db.consume_storage();
 
         // Get extended initial reads keys.
-        let initial_reads_keys = extended_initial_reads.keys();
+        let accessed_keys: AccessedKeys = extended_initial_reads.keys().into();
 
         // TODO(Aviv): Try to undertand if we can create classes trie commitment info
         // without the compiled class hashes.
         let mut commitment_infos = StateCommitmentInfos::new(
             &previous_state_roots,
             &new_roots,
             &mut map_storage,
-            &initial_reads_keys,
+            &accessed_keys,
         )
         .await
         .map_err(|e| ProofProviderError::BlockCommitmentError(e.to_string()))?;