feat: debug_readAllSlotsFromAccount rpc call (#939)

src/eth/rpc/rpc_server.rs

@@ -131,6 +131,7 @@ fn register_methods(mut module: RpcModule<RpcContext>) -> anyhow::Result<RpcModu
         module.register_async_method("evm_setNextBlockTimestamp", evm_set_next_block_timestamp)?;
         module.register_async_method("evm_mine", evm_mine)?;
         module.register_async_method("debug_setHead", debug_set_head)?;
+        module.register_async_method("debug_readAllSlotsFromAccount", debug_read_all_slots)?;
     }
 
     // stratus health check
@@ -212,6 +213,12 @@ async fn evm_set_next_block_timestamp(params: Params<'_>, ctx: Arc<RpcContext>)
     Ok(serde_json::to_value(timestamp).expect_infallible())
 }
 
+#[cfg(feature = "dev")]
+async fn debug_read_all_slots(params: Params<'_>, ctx: Arc<RpcContext>) -> anyhow::Result<JsonValue, RpcError> {
+    let (_, address) = next_rpc_param::<Address>(params.sequence())?;
+    Ok(serde_json::to_value(ctx.storage.read_all_slots(&address).await?).expect_infallible())
+}
+
 // Status
 async fn net_listening(params: Params<'_>, arc: Arc<RpcContext>) -> anyhow::Result<JsonValue, RpcError> {
     stratus_readiness(params, arc).await

src/eth/storage/inmemory/inmemory_permanent.rs

@@ -219,6 +219,17 @@ impl PermanentStorage for InMemoryPermanentStorage {
         Ok(slots)
     }
 
+    async fn read_all_slots(&self, address: &Address) -> anyhow::Result<Vec<Slot>> {
+        let state = self.lock_read().await;
+
+        let Some(account) = state.accounts.get(address) else {
+            tracing::trace!(%address, "account not found in permanent");
+            return Ok(Default::default());
+        };
+
+        Ok(account.slots.clone().into_values().map(|slot| slot.get_current()).collect())
+    }
+
     async fn read_block(&self, selection: &BlockSelection) -> anyhow::Result<Option<Block>> {
         tracing::debug!(?selection, "reading block");
 

src/eth/storage/permanent_storage.rs

@@ -81,6 +81,9 @@ pub trait PermanentStorage: Send + Sync {
     /// Retrieves a random sample of slots, from the provided start and end blocks.
     async fn read_slots_sample(&self, start: BlockNumber, end: BlockNumber, max_samples: u64, seed: u64) -> anyhow::Result<Vec<SlotSample>>;
 
+    /// Retrieves all current slots associated to an address.
+    async fn read_all_slots(&self, address: &Address) -> anyhow::Result<Vec<Slot>>;
+
     // -------------------------------------------------------------------------
     // Global state
     // -------------------------------------------------------------------------

src/eth/storage/postgres_permanent/postgres_permanent.rs

@@ -97,6 +97,10 @@ impl PermanentStorage for PostgresPermanentStorage {
         PermanentStorageKind::Postgres { url: self.url.clone() }
     }
 
+    async fn read_all_slots(&self, _address: &Address) -> anyhow::Result<Vec<Slot>> {
+        todo!();
+    }
+
     async fn allocate_evm_thread_resources(&self) -> anyhow::Result<()> {
         let conn = self.pool.acquire().await?;
         let conn = conn.leak();

src/eth/storage/rocks/rocks_permanent.rs

@@ -9,6 +9,8 @@ use async_trait::async_trait;
 use futures::future::join_all;
 
 use super::rocks_state::RocksStorageState;
+use super::types::AddressRocksdb;
+use super::types::SlotIndexRocksdb;
 use crate::config::PermanentStorageKind;
 use crate::eth::primitives::Account;
 use crate::eth::primitives::Address;
@@ -229,4 +231,18 @@ impl PermanentStorage for RocksPermanentStorage {
     async fn read_slots_sample(&self, _start: BlockNumber, _end: BlockNumber, _max_samples: u64, _seed: u64) -> anyhow::Result<Vec<SlotSample>> {
         todo!()
     }
+
+    async fn read_all_slots(&self, address: &Address) -> anyhow::Result<Vec<Slot>> {
+        let address: AddressRocksdb = (*address).into();
+        Ok(self
+            .state
+            .account_slots
+            .iter_from((address, SlotIndexRocksdb::from(0)), rocksdb::Direction::Forward)
+            .take_while(|((addr, _), _)| &address == addr)
+            .map(|((_, idx), value)| Slot {
+                index: idx.into(),
+                value: value.into(),
+            })
+            .collect())
+    }
 }

src/eth/storage/rocks/types.rs

@@ -238,6 +238,8 @@ impl From<BlockNumberRocksdb> for u64 {
 #[derive(Clone, Default, Hash, Eq, PartialEq, PartialOrd, Ord, serde::Serialize, serde::Deserialize)]
 pub struct SlotIndexRocksdb(U256);
 
+gen_newtype_from!(self = SlotIndexRocksdb, other = u64);
+
 impl SlotIndexRocksdb {
     pub fn inner_value(&self) -> U256 {
         self.0

src/eth/storage/stratus_storage.rs

@@ -327,6 +327,11 @@ impl StratusStorage {
         Ok(slots)
     }
 
+    #[tracing::instrument(skip_all)]
+    pub async fn read_all_slots(&self, address: &Address) -> anyhow::Result<Vec<Slot>> {
+        self.perm.read_all_slots(address).await
+    }
+
     // -------------------------------------------------------------------------
     // Blocks
     // -------------------------------------------------------------------------