apollo_storage: implement contract storage query in storage reader (#11225)

Author: nadin-Starkware

crates/apollo_storage/src/state/mod.rs

@@ -352,6 +352,27 @@ impl<'env, Mode: TransactionKind> StateReader<'env, Mode> {
         }
     }
 
+    /// Returns the storage value for a given contract address, storage key, and block number using
+    /// direct key-value lookup.
+    ///
+    /// # Arguments
+    /// * address - contract address to search for.
+    /// * key - storage key to search for.
+    /// * block_number - block number to search at.
+    ///
+    /// # Errors
+    /// Returns [`StorageError`] if there was an error searching the table.
+    pub fn get_storage_by_key(
+        &self,
+        address: ContractAddress,
+        key: StorageKey,
+        block_number: BlockNumber,
+    ) -> StorageResult<Option<Felt>> {
+        let db_key = ((address, key), block_number);
+        let value = self.storage_table.get(self.txn, &db_key)?;
+        Ok(value)
+    }
+
     /// Returns the class definition at a given state number.
     ///
     /// If class_hash is not found, returns `None`.

crates/apollo_storage/src/storage_reader_types.rs

@@ -176,6 +176,7 @@ impl StorageReaderServerHandler<StorageReaderRequest, StorageReaderResponse>
         request: StorageReaderRequest,
     ) -> Result<StorageReaderResponse, StorageError> {
         let txn = storage_reader.begin_ro_txn()?;
+        let state_reader = txn.get_state_reader()?;
         match request {
             // ============ State-Related Requests ============
             StorageReaderRequest::StateDiffsLocation(block_number) => {
@@ -190,8 +191,11 @@ impl StorageReaderServerHandler<StorageReaderRequest, StorageReaderResponse>
                 let state_diff = txn.get_state_diff_from_location(location)?;
                 Ok(StorageReaderResponse::StateDiffsFromLocation(state_diff))
             }
-            StorageReaderRequest::ContractStorage(_key, _block_number) => {
-                unimplemented!()
+            StorageReaderRequest::ContractStorage((address, key), block_number) => {
+                let value = state_reader
+                    .get_storage_by_key(address, key, block_number)?
+                    .unwrap_or_default();
+                Ok(StorageReaderResponse::ContractStorage(value))
             }
             StorageReaderRequest::Nonces(_address, _block_number) => {
                 unimplemented!()