apollo_storage: add scan_storage_keys_for_contract (#13554)

Author: yoavGrs

crates/apollo_storage/src/state/mod.rs

@@ -564,6 +564,21 @@ impl<'env, Mode: TransactionKind> StateReader<'env, Mode> {
         let mut cursor = self.deployed_contracts_table.cursor(self.txn)?;
         scan_at_block(&mut cursor, start, end, block_target, limit)
     }
+
+    /// Returns all storage values for `addr` in the key range [`start_key`, `end`] at
+    /// `scan_at_block`, up to `limit` entries.
+    pub fn scan_storage_keys_for_contract(
+        &self,
+        addr: ContractAddress,
+        start: StorageKey,
+        end: StorageKey,
+        block_target: BlockNumber,
+        limit: usize,
+    ) -> StorageResult<Vec<(StorageKey, Felt)>> {
+        let mut cursor = self.storage_table.cursor(self.txn)?;
+        let entries = scan_at_block(&mut cursor, (addr, start), (addr, end), block_target, limit)?;
+        Ok(entries.into_iter().map(|((_, key), value)| (key, value)).collect())
+    }
 }
 
 impl StateStorageWriter for StorageTxn<'_, RW> {

crates/apollo_storage/src/state/state_test.rs

@@ -8,7 +8,7 @@ use starknet_api::block::BlockNumber;
 use starknet_api::core::{ClassHash, ContractAddress, Nonce};
 use starknet_api::deprecated_contract_class::ContractClass as DeprecatedContractClass;
 use starknet_api::hash::StarkHash;
-use starknet_api::state::{SierraContractClass, StateNumber, ThinStateDiff};
+use starknet_api::state::{SierraContractClass, StateNumber, StorageKey, ThinStateDiff};
 use starknet_api::{class_hash, compiled_class_hash, contract_address, felt, storage_key};
 use starknet_types_core::felt::Felt;
 
@@ -1062,3 +1062,37 @@ scan_cases!(
         );
     }
 );
+
+scan_cases!(
+    key_macro = storage_key,
+    value_macro = felt,
+    fn test_scan_storage_keys_for_contract(
+        #[case] start: StorageKey,
+        #[case] end: StorageKey,
+        #[case] synced_block: BlockNumber,
+        #[case] limit: usize,
+        #[case] expected: Vec<(StorageKey, Felt)>,
+    ) {
+        let addr = contract_address!("0x1");
+        let ((reader, mut writer), _temp_dir) = get_test_storage();
+        write_two_block_state_diffs(
+            &mut writer,
+            ThinStateDiff {
+                storage_diffs: IndexMap::from([(addr, block_0_entries!())]),
+                ..Default::default()
+            },
+            ThinStateDiff {
+                storage_diffs: IndexMap::from([(addr, block_1_entries!())]),
+                ..Default::default()
+            },
+        );
+        let txn = reader.begin_ro_txn().unwrap();
+        let state_reader = txn.get_state_reader().unwrap();
+        assert_eq!(
+            state_reader
+                .scan_storage_keys_for_contract(addr, start, end, synced_block, limit)
+                .unwrap(),
+            expected,
+        );
+    }
+);