feat(ext): add state override support for ERC20 balances and allowances

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "uniswap-v3-sdk"
-version = "3.1.0"
+version = "3.1.1"
 edition = "2021"
 authors = ["Shuhui Luo <twitter.com/aureliano_law>"]
 description = "Uniswap V3 SDK for Rust"

src/error.rs

@@ -38,10 +38,6 @@ pub enum Error {
     #[error("Invalid price")]
     InvalidPrice,
 
-    #[cfg(feature = "extensions")]
-    #[error("Invalid tick range")]
-    InvalidRange,
-
     #[error("Overflow in full math mulDiv")]
     MulDivOverflow,
 
@@ -60,6 +56,13 @@ pub enum Error {
     #[error("No tick data provider was given")]
     NoTickDataError,
 
+    #[error("{0}")]
+    TickListError(#[from] TickListError),
+
+    #[cfg(feature = "extensions")]
+    #[error("Invalid tick range")]
+    InvalidRange,
+
     #[cfg(feature = "extensions")]
     #[error("{0}")]
     ContractError(#[from] ContractError),
@@ -68,8 +71,9 @@ pub enum Error {
     #[error("{0}")]
     LensError(#[from] LensError),
 
-    #[error("{0}")]
-    TickListError(#[from] TickListError),
+    #[cfg(feature = "extensions")]
+    #[error("Invalid access list")]
+    InvalidAccessList,
 }
 
 #[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, thiserror::Error)]
@@ -81,3 +85,10 @@ pub enum TickListError {
     #[error("Not contained in tick list")]
     NotContained,
 }
+
+#[cfg(feature = "extensions")]
+impl From<alloy::transports::TransportError> for Error {
+    fn from(e: alloy::transports::TransportError) -> Self {
+        Self::ContractError(ContractError::TransportError(e))
+    }
+}

src/extensions/mod.rs

@@ -5,6 +5,7 @@ mod ephemeral_tick_map_data_provider;
 mod pool;
 mod position;
 mod price_tick_conversions;
+mod state_overrides;
 mod tick_bit_map;
 mod tick_map;
 
@@ -13,5 +14,6 @@ pub use ephemeral_tick_map_data_provider::EphemeralTickMapDataProvider;
 pub use pool::*;
 pub use position::*;
 pub use price_tick_conversions::*;
+pub use state_overrides::*;
 pub use tick_bit_map::*;
 pub use tick_map::*;

src/extensions/state_overrides.rs

@@ -0,0 +1,115 @@
+use crate::prelude::Error;
+use alloy::{
+    eips::eip2930::{AccessList, AccessListItem},
+    providers::Provider,
+    rpc::types::{
+        state::{AccountOverride, StateOverride},
+        TransactionRequest,
+    },
+    transports::Transport,
+};
+use alloy_primitives::{
+    map::{B256HashMap, B256HashSet},
+    Address, B256, U256,
+};
+use alloy_sol_types::SolCall;
+use uniswap_lens::bindings::ierc20::IERC20;
+
+#[inline]
+pub async fn get_erc20_state_overrides<T, P>(
+    token: Address,
+    owner: Address,
+    spender: Address,
+    amount: U256,
+    provider: P,
+) -> Result<StateOverride, Error>
+where
+    T: Transport + Clone,
+    P: Provider<T>,
+{
+    let balance_tx = TransactionRequest::default()
+        .to(token)
+        .gas_limit(0x11E1A300) // avoids "intrinsic gas too low" error
+        .input(IERC20::balanceOfCall { account: owner }.abi_encode().into());
+    let allowance_tx = TransactionRequest::default()
+        .to(token)
+        .gas_limit(0x11E1A300)
+        .input(IERC20::allowanceCall { owner, spender }.abi_encode().into());
+    let balance_access_list = provider.create_access_list(&balance_tx).await?.access_list;
+    let allowance_access_list = provider
+        .create_access_list(&allowance_tx)
+        .await?
+        .access_list;
+    // tokens on L2 and those with a proxy will have more than one access list entry
+    let filtered_balance_access_list = filter_access_list(balance_access_list, token);
+    let filtered_allowance_access_list = filter_access_list(allowance_access_list, token);
+    if filtered_balance_access_list.len() != 1 || filtered_allowance_access_list.len() != 1 {
+        return Err(Error::InvalidAccessList);
+    }
+    // get rid of the storage key of implementation address
+    let balance_slots_set =
+        B256HashSet::from_iter(filtered_balance_access_list[0].storage_keys.clone());
+    let allowance_slots_set =
+        B256HashSet::from_iter(filtered_allowance_access_list[0].storage_keys.clone());
+    let state_diff = B256HashMap::from_iter(
+        balance_slots_set
+            .symmetric_difference(&allowance_slots_set)
+            .cloned()
+            .map(|slot| (slot, B256::from(amount))),
+    );
+    if state_diff.len() != 2 {
+        return Err(Error::InvalidAccessList);
+    }
+    Ok(StateOverride::from_iter([(
+        token,
+        AccountOverride {
+            state_diff: Some(state_diff),
+            ..Default::default()
+        },
+    )]))
+}
+
+fn filter_access_list(access_list: AccessList, token: Address) -> Vec<AccessListItem> {
+    access_list
+        .0
+        .into_iter()
+        .filter(|item| item.address == token)
+        .collect()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::tests::*;
+    use alloy_primitives::{address, U256};
+    use uniswap_sdk_core::prelude::{BaseCurrency, NONFUNGIBLE_POSITION_MANAGER_ADDRESSES};
+
+    #[tokio::test]
+    async fn test_get_erc20_overrides() {
+        let provider = PROVIDER.clone();
+        let owner = address!("88e6A0c2dDD26FEEb64F039a2c41296FcB3f5640");
+        let npm = *NONFUNGIBLE_POSITION_MANAGER_ADDRESSES.get(&1).unwrap();
+        let amount = U256::from(1_000_000);
+        let overrides =
+            get_erc20_state_overrides(USDC.address(), owner, npm, amount, provider.clone())
+                .await
+                .unwrap();
+        let usdc = IERC20::new(USDC.address(), provider);
+        let balance = usdc
+            .balanceOf(owner)
+            .call()
+            .overrides(&overrides)
+            .await
+            .unwrap()
+            ._0;
+        assert_eq!(balance, amount);
+        let allowance = usdc
+            .allowance(owner, npm)
+            .call()
+            .overrides(&overrides)
+            .await
+            .unwrap()
+            ._0;
+        assert_eq!(allowance, amount);
+    }
+}