retrieve extra keys for 0x1 contract

crates/bin/prove_block/src/rpc_utils.rs

@@ -118,9 +118,20 @@ pub(crate) async fn get_storage_proofs(
     old_block_number: Felt,
 ) -> Result<HashMap<Felt, PathfinderProof>, ClientError> {
     let accessed_keys_by_address = {
-        let mut keys = get_all_accessed_keys(tx_execution_infos);
+        let (mut keys, storage_reads) = get_all_accessed_keys(tx_execution_infos);
+
         // We need to fetch the storage proof for the block hash contract
         keys.entry(contract_address!("0x1")).or_default().insert(old_block_number.try_into().unwrap());
+
+        // Within the Starknet architecture, the address 0x1 is a special address that maps block numbers to their corresponding block hashes. As some contracts might access storage reads using `get_block_hash_syscall`, it is necessary to add some extra keys here.
+        // By leveraging the structure of this special contract, we filter out storage_read_values that are greater than old_block_number and add these extra values to the necessary keys from the 0x1 contract address.
+        // It is worth noting that this approach incurs some overhead due to the retrieval of additional data.
+        let additional_storage_reads: Vec<Felt> =
+            storage_reads.values().flat_map(|vec| vec.iter().filter(|&x| x < &old_block_number)).cloned().collect();
+        keys.entry(contract_address!("0x1"))
+            .or_default()
+            .extend(additional_storage_reads.into_iter().map(|key| StorageKey::try_from(key).unwrap()));
+
         keys
     };
 

crates/bin/prove_block/src/utils.rs

@@ -84,23 +84,26 @@ fn process_function_invocations(
 /// We need this as the OS require proofs for all the accessed values
 pub(crate) fn get_all_accessed_keys(
     tx_execution_infos: &[TransactionExecutionInfo],
-) -> HashMap<ContractAddress, HashSet<StorageKey>> {
+) -> (HashMap<ContractAddress, HashSet<StorageKey>>, HashMap<ContractAddress, Vec<Felt252>>) {
     let mut accessed_keys_by_address: HashMap<ContractAddress, HashSet<StorageKey>> = HashMap::new();
+    let mut call_storage_reads: HashMap<ContractAddress, Vec<Felt252>> = HashMap::new();
 
     for tx_execution_info in tx_execution_infos {
-        let accessed_keys_in_tx = get_accessed_keys_in_tx(tx_execution_info);
+        let (accessed_keys_in_tx, accessed_storage_reads_in_tx) = get_accessed_keys_in_tx(tx_execution_info);
         for (contract_address, storage_keys) in accessed_keys_in_tx {
             accessed_keys_by_address.entry(contract_address).or_default().extend(storage_keys);
         }
+        call_storage_reads.extend(accessed_storage_reads_in_tx);
     }
 
-    accessed_keys_by_address
+    (accessed_keys_by_address, call_storage_reads)
 }
 
 fn get_accessed_keys_in_tx(
     tx_execution_info: &TransactionExecutionInfo,
-) -> HashMap<ContractAddress, HashSet<StorageKey>> {
+) -> (HashMap<ContractAddress, HashSet<StorageKey>>, HashMap<ContractAddress, Vec<Felt252>>) {
     let mut accessed_keys_by_address: HashMap<ContractAddress, HashSet<StorageKey>> = HashMap::new();
+    let mut accessed_storage_reads: HashMap<ContractAddress, Vec<Felt252>> = HashMap::new();
 
     for call_info in [
         &tx_execution_info.validate_call_info,
@@ -110,30 +113,42 @@ fn get_accessed_keys_in_tx(
     .into_iter()
     .flatten()
     {
-        let call_storage_keys = get_accessed_storage_keys(call_info);
+        let (call_storage_keys, call_storage_reads) = get_accessed_storage_keys(call_info);
         for (contract_address, storage_keys) in call_storage_keys {
             accessed_keys_by_address.entry(contract_address).or_default().extend(storage_keys);
         }
+        accessed_storage_reads.extend(call_storage_reads);
     }
 
-    accessed_keys_by_address
+    (accessed_keys_by_address, accessed_storage_reads)
 }
 
-fn get_accessed_storage_keys(call_info: &CallInfo) -> HashMap<ContractAddress, HashSet<StorageKey>> {
+fn get_accessed_storage_keys(
+    call_info: &CallInfo,
+) -> (HashMap<ContractAddress, HashSet<StorageKey>>, HashMap<ContractAddress, Vec<Felt252>>) {
     let mut accessed_keys_by_address: HashMap<ContractAddress, HashSet<StorageKey>> = HashMap::new();
+    let mut accessed_storage_read_values: HashMap<ContractAddress, Vec<Felt252>> = HashMap::new();
 
     let contract_address = &call_info.call.storage_address;
     accessed_keys_by_address
         .entry(*contract_address)
         .or_default()
         .extend(call_info.accessed_storage_keys.iter().copied());
 
+    accessed_storage_read_values
+        .entry(*contract_address)
+        .or_default()
+        .extend(call_info.storage_read_values.iter().copied());
+
     for inner_call in &call_info.inner_calls {
-        let inner_call_storage_keys = get_accessed_storage_keys(inner_call);
+        let (inner_call_storage_keys, inner_call_storage_read_values) = get_accessed_storage_keys(inner_call);
         for (contract_address, storage_keys) in inner_call_storage_keys {
             accessed_keys_by_address.entry(contract_address).or_default().extend(storage_keys);
         }
+        for (contract_address, storage_read_values) in inner_call_storage_read_values {
+            accessed_storage_read_values.entry(contract_address).or_default().extend(storage_read_values);
+        }
     }
 
-    accessed_keys_by_address
+    (accessed_keys_by_address, accessed_storage_read_values)
 }