feat(alloy): Revamp get_storage_at function to use Alloy

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "uniswap-lens"
-version = "0.0.2"
+version = "0.1.0"
 edition = "2021"
 authors = ["Shuhui Luo <twitter.com/aureliano_law>"]
 description = "Contains ephemeral lens contracts that can be called without deployment and their Rust interfaces."
@@ -11,11 +11,11 @@ keywords = ["alloy", "ethereum", "solidity", "rust", "uniswap"]
 include = ["src/**/*.rs"]
 
 [dependencies]
-alloy = { version = "0.2.0", features = ["contract", "transports"] }
+alloy = { version = "0.2.0", features = ["contract", "rpc-types", "transports"] }
 anyhow = "1"
 
 [dev-dependencies]
-alloy = { version = "0.2.0", features = ["rpc-types", "transport-http"] }
+alloy = { version = "0.2.0", features = ["transport-http"] }
 dotenv = "0.15.0"
 futures = "0.3"
 once_cell = "1.19"

src/lib.rs

@@ -18,12 +18,11 @@ pub mod bindings;
 pub mod caller;
 pub mod pool_lens;
 pub mod position_lens;
+pub mod storage_lens;
 
 #[cfg(test)]
 mod tests;
 
-// pub mod storage_lens;
-
 pub mod prelude {
-    pub use super::{bindings::*, pool_lens::*, position_lens::*};
+    pub use super::{bindings::*, pool_lens::*, position_lens::*, storage_lens::*};
 }

src/storage_lens.rs

@@ -1,6 +1,14 @@
-use crate::bindings::ephemeral_storage_lens::{EphemeralStorageLens, EPHEMERALSTORAGELENS_DEPLOYED_BYTECODE};
-use ethers::prelude::*;
-use std::sync::Arc;
+use crate::bindings::ephemeralstoragelens::{
+    EphemeralStorageLens, EphemeralStorageLens::EphemeralStorageLensInstance,
+};
+use alloy::{
+    eips::BlockId,
+    primitives::{Address, FixedBytes},
+    providers::Provider,
+    rpc::types::state::{AccountOverride, StateOverride},
+    transports::Transport,
+};
+use anyhow::Result;
 
 /// Batch `eth_getStorageAt` RPC calls in a single `eth_call` by overriding the target contract's
 /// deployed bytecode with `EphemeralStorageLens`
@@ -9,54 +17,69 @@ use std::sync::Arc;
 ///
 /// * `address`: The contract address to fetch storage from
 /// * `slots`: The storage slots to query
-/// * `client`: The ethers client
+/// * `provider`: The alloy provider
 /// * `block_id`: Optional block id to query
 ///
 /// returns: Result<Vec<[u8; 32], Global>, ContractError<M>>
-pub async fn get_storage_at<M: Middleware>(
+pub async fn get_storage_at<T, P>(
     address: Address,
-    slots: Vec<[u8; 32]>,
-    client: Arc<M>,
+    slots: Vec<FixedBytes<32>>,
+    provider: P,
     block_id: Option<BlockId>,
-) -> Result<Vec<[u8; 32]>, ContractError<M>> {
-    let block_id = block_id.unwrap_or(BlockId::from(BlockNumber::Latest));
+) -> Result<Vec<FixedBytes<32>>>
+where
+    T: Transport + Clone,
+    P: Provider<T>,
+{
     // override the deployed bytecode at `address`
-    let mut state = spoof::state();
-    state
-        .account(address)
-        .code(EPHEMERALSTORAGELENS_DEPLOYED_BYTECODE.clone());
-    let lens = EphemeralStorageLens::new(address, client);
-    lens.extsload(slots).call_raw().block(block_id).state(&state).await
+    let mut state = StateOverride::default();
+    state.insert(
+        address,
+        AccountOverride {
+            code: Some(EphemeralStorageLens::DEPLOYED_BYTECODE.clone()),
+            ..Default::default()
+        },
+    );
+    let lens = EphemeralStorageLensInstance::new(address, provider);
+    let call_builder = lens.extsload(slots).state(state);
+    let call_builder = match block_id {
+        Some(block_id) => call_builder.block(block_id),
+        None => call_builder,
+    };
+    Ok(call_builder.call().await?._0)
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::tests::*;
+    use alloy::primitives::{address, U256};
     use anyhow::Result;
     use futures::future::join_all;
 
-    const POOL_ADDRESS: &str = "0x88e6A0c2dDD26FEEb64F039a2c41296FcB3f5640";
-    static BLOCK_NUMBER: Lazy<BlockId> = Lazy::new(|| BlockId::from(17000000));
+    const POOL_ADDRESS: Address = address!("88e6A0c2dDD26FEEb64F039a2c41296FcB3f5640");
 
     #[tokio::test]
     async fn test_get_storage_at() -> Result<()> {
-        let client = Arc::new(MAINNET.provider());
+        let provider = PROVIDER.clone();
         let slots = get_storage_at(
-            POOL_ADDRESS.parse::<Address>()?,
-            (0..10).map(|i| H256::from_low_u64_be(i).to_fixed_bytes()).collect(),
-            client.clone(),
+            POOL_ADDRESS,
+            (0..10)
+                .map(|i| FixedBytes::from(U256::from_limbs([i, 0, 0, 0])))
+                .collect(),
+            provider.clone(),
             Some(*BLOCK_NUMBER),
         )
         .await?;
         let slots_ref = slots.as_slice();
-        let address = POOL_ADDRESS.parse::<Address>()?;
-        let client = client.as_ref();
+        let provider = provider.root();
         let futures = (0..10).map(|i| async move {
-            let slot = client
-                .get_storage_at(address, H256::from_low_u64_be(i), Some(*BLOCK_NUMBER))
+            let slot = provider
+                .get_storage_at(POOL_ADDRESS, U256::from_limbs([i, 0, 0, 0]))
+                .block_id(*BLOCK_NUMBER)
                 .await
                 .unwrap();
-            assert_eq!(slot.to_fixed_bytes(), slots_ref[i as usize]);
+            assert_eq!(slot, U256::from_be_bytes(slots_ref[i as usize].0));
         });
         join_all(futures).await;
         Ok(())