Osiris/checkpoints

Cargo.toml

@@ -38,7 +38,7 @@ future-utils = "0.12"
 async-trait = "0.1"
 
 # serde
-serde = "1.0"
+serde = {version = "1.0", features = ["std"] }
 serde_json = "1.0"
 
 # misc

benches/uniswap_v2.rs

@@ -1,17 +1,12 @@
 use alloy::primitives::{address, U256};
-use amms::amms::{amm::AutomatedMarketMaker, uniswap_v2::UniswapV2Pool};
+use amms::amms::{amm::AutomatedMarketMaker, uniswap_v2::UniswapV2Pool, Token};
 use criterion::{criterion_group, criterion_main, Criterion};
 use rand::Rng;
 
 fn simulate_swap(c: &mut Criterion) {
-    let token_a = address!("c02aaa39b223fe8d0a0e5c4f27ead9083c756cc2");
-    let token_b = address!("fc0d6cf33e38bce7ca7d89c0e292274031b7157a");
-
     let pool = UniswapV2Pool {
-        token_a,
-        token_a_decimals: 18,
-        token_b,
-        token_b_decimals: 18,
+        token_a: Token::new_with_decimals(address!("c02aaa39b223fe8d0a0e5c4f27ead9083c756cc2"), 18),
+        token_b: Token::new_with_decimals(address!("fc0d6cf33e38bce7ca7d89c0e292274031b7157a"), 18),
         reserve_0: 20_000_000_u128,
         reserve_1: 20_000_000_u128,
         fee: 300,
@@ -23,7 +18,9 @@ fn simulate_swap(c: &mut Criterion) {
         b.iter_with_setup(
             || U256::from(rng.gen_range(1_000..=1e24 as u128)),
             |amount| {
-                let _ = pool.simulate_swap(token_a, token_b, amount).unwrap();
+                let _ = pool
+                    .simulate_swap(pool.token_a.address(), pool.token_b.address(), amount)
+                    .unwrap();
             },
         );
     });

benches/uniswap_v3.rs

@@ -11,6 +11,7 @@ use alloy_throttle::ThrottleLayer;
 use amms::amms::{
     amm::{AutomatedMarketMaker, AMM},
     uniswap_v3::{UniswapV3Factory, UniswapV3Pool},
+    Token,
 };
 use criterion::{criterion_group, criterion_main, Criterion};
 use rand::Rng;
@@ -26,15 +27,18 @@ fn simulate_swap(c: &mut Criterion) {
 
     let provider = Arc::new(ProviderBuilder::new().on_client(client));
 
-    let token_a = address!("A0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48");
-    let token_b = address!("C02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2");
-
     let runtime = Runtime::new().expect("Failed to create Tokio runtime");
     let pool = runtime.block_on(async {
         let pool = AMM::UniswapV3Pool(UniswapV3Pool {
             address: address!("88e6A0c2dDD26FEEb64F039a2c41296FcB3f5640"),
-            token_a,
-            token_b,
+            token_a: Token::new_with_decimals(
+                address!("A0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48"),
+                6,
+            ),
+            token_b: Token::new_with_decimals(
+                address!("C02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2"),
+                18,
+            ),
             tick_spacing: 10,
             fee: 500,
             ..Default::default()
@@ -57,7 +61,9 @@ fn simulate_swap(c: &mut Criterion) {
         b.iter_with_setup(
             || U256::from(rng.gen_range(1_000..=1e24 as u128)),
             |amount| {
-                let _ = pool.simulate_swap(token_a, token_b, amount).unwrap();
+                let _ = pool
+                    .simulate_swap(pool.token_a.address(), pool.token_b.address(), amount)
+                    .unwrap();
             },
         );
     });

src/amms/amm.rs

@@ -22,6 +22,9 @@ pub trait AutomatedMarketMaker {
     /// Address of the AMM
     fn address(&self) -> Address;
 
+    /// Whether the AMM has been initialized.
+    fn initialized(&self) -> bool;
+
     /// Event signatures that indicate when the AMM should be synced
     fn sync_events(&self) -> Vec<B256>;
 
@@ -75,6 +78,12 @@ macro_rules! amm {
                 }
             }
 
+            fn initialized(&self) -> bool {
+                match self {
+                    $(AMM::$pool_type(pool) => pool.initialized(),)+
+                }
+            }
+
             fn sync_events(&self) -> Vec<B256> {
                 match self {
                     $(AMM::$pool_type(pool) => pool.sync_events(),)+

src/amms/balancer/mod.rs

@@ -115,6 +115,10 @@ impl AutomatedMarketMaker for BalancerPool {
         self.address
     }
 
+    fn initialized(&self) -> bool {
+        !self.state.is_empty()
+    }
+
     fn sync_events(&self) -> Vec<B256> {
         vec![IBPool::LOG_SWAP::SIGNATURE_HASH]
     }
@@ -365,6 +369,7 @@ impl AutomatedMarketMakerFactory for BalancerFactory {
 impl DiscoverySync for BalancerFactory {
     fn discover<T, N, P>(
         &self,
+        from_block: Option<BlockId>,
         to_block: BlockId,
         provider: Arc<P>,
     ) -> impl Future<Output = Result<Vec<AMM>, AMMError>>
@@ -378,7 +383,7 @@ impl DiscoverySync for BalancerFactory {
             address = ?self.address,
             "Discovering all pools"
         );
-        self.get_all_pools(to_block, provider)
+        self.get_all_pools(from_block, to_block, provider)
     }
 
     fn sync<T, N, P>(
@@ -411,6 +416,7 @@ impl BalancerFactory {
 
     pub async fn get_all_pools<T, N, P>(
         &self,
+        from_block: Option<BlockId>,
         block_number: BlockId,
         provider: Arc<P>,
     ) -> Result<Vec<AMM>, AMMError>
@@ -427,7 +433,8 @@ impl BalancerFactory {
         let mut futures = FuturesUnordered::new();
 
         let sync_step = 100_000;
-        let mut latest_block = self.creation_block;
+        let mut latest_block =
+            from_block.map_or(self.creation_block(), |b| b.as_u64().unwrap_or_default());
         while latest_block < block_number.as_u64().unwrap_or_default() {
             let mut block_filter = disc_filter.clone();
             let from_block = latest_block;

src/amms/erc_4626/mod.rs

@@ -69,6 +69,10 @@ impl AutomatedMarketMaker for ERC4626Vault {
         self.vault_token
     }
 
+    fn initialized(&self) -> bool {
+        !self.vault_token.is_zero() && !self.asset_token.is_zero()
+    }
+
     fn sync_events(&self) -> Vec<B256> {
         vec![
             IERC4626Vault::Deposit::SIGNATURE_HASH,

src/amms/factory.rs

@@ -23,6 +23,7 @@ use std::{
 pub trait DiscoverySync {
     fn discover<T, N, P>(
         &self,
+        from_block: Option<BlockId>,
         to_block: BlockId,
         provider: Arc<P>,
     ) -> impl Future<Output = Result<Vec<AMM>, AMMError>>
@@ -63,6 +64,7 @@ pub trait AutomatedMarketMakerFactory: DiscoverySync {
         Self::PoolVariant::default().sync_events()
     }
 
+    /// Pool variant of the factory
     fn pool_variant(&self) -> Self::PoolVariant {
         Self::PoolVariant::default()
     }
@@ -129,14 +131,14 @@ macro_rules! factory {
 
 
         impl Factory {
-            pub async fn discover<T, N, P>(&self, to_block: BlockId, provider: Arc<P>) -> Result<Vec<AMM>, AMMError>
+            pub async fn discover<T, N, P>(&self, from_block: Option<BlockId>, to_block: BlockId, provider: Arc<P>) -> Result<Vec<AMM>, AMMError>
             where
                 T: Transport + Clone,
                 N: Network,
                 P: Provider<T, N>,
             {
                 match self {
-                    $(Factory::$factory_type(factory) => factory.discover(to_block, provider).await,)+
+                    $(Factory::$factory_type(factory) => factory.discover(from_block, to_block, provider).await,)+
                 }
             }
 
@@ -171,6 +173,10 @@ impl AutomatedMarketMaker for NoopAMM {
         unreachable!()
     }
 
+    fn initialized(&self) -> bool {
+        unreachable!()
+    }
+
     fn sync_events(&self) -> Vec<B256> {
         unreachable!()
     }

src/amms/uniswap_v2/mod.rs

@@ -90,6 +90,10 @@ impl AutomatedMarketMaker for UniswapV2Pool {
         self.address
     }
 
+    fn initialized(&self) -> bool {
+        !self.token_a.address.is_zero() && !self.token_b.address.is_zero()
+    }
+
     fn sync_events(&self) -> Vec<B256> {
         vec![IUniswapV2Pair::Sync::SIGNATURE_HASH]
     }
@@ -397,6 +401,7 @@ impl UniswapV2Factory {
 
     pub async fn get_all_pairs<T, N, P>(
         factory_address: Address,
+        from_block: Option<BlockId>,
         block_number: BlockId,
         provider: Arc<P>,
     ) -> Result<Vec<Address>, AMMError>
@@ -406,6 +411,7 @@ impl UniswapV2Factory {
         P: Provider<T, N>,
     {
         let factory = IUniswapV2FactoryInstance::new(factory_address, provider.clone());
+
         let pairs_length = factory
             .allPairsLength()
             .call()
@@ -414,9 +420,21 @@ impl UniswapV2Factory {
             ._0
             .to::<usize>();
 
+        let start_idx = if let Some(from_block) = from_block {
+            factory
+                .allPairsLength()
+                .call()
+                .block(from_block)
+                .await?
+                ._0
+                .to::<usize>()
+        } else {
+            0
+        };
+
         let step = 766;
         let mut futures_unordered = FuturesUnordered::new();
-        for i in (0..pairs_length).step_by(step) {
+        for i in (start_idx..pairs_length).step_by(step) {
             // Note that the batch contract handles if the step is greater than the pairs length
             // So we can pass the step in as is without checking for this condition
             let deployer = IGetUniswapV2PairsBatchRequest::deploy_builder(
@@ -561,6 +579,7 @@ impl AutomatedMarketMakerFactory for UniswapV2Factory {
 impl DiscoverySync for UniswapV2Factory {
     fn discover<T, N, P>(
         &self,
+        from_block: Option<BlockId>,
         to_block: BlockId,
         provider: Arc<P>,
     ) -> impl Future<Output = Result<Vec<AMM>, AMMError>>
@@ -577,8 +596,13 @@ impl DiscoverySync for UniswapV2Factory {
 
         let provider = provider.clone();
         async move {
-            let pairs =
-                UniswapV2Factory::get_all_pairs(self.address, to_block, provider.clone()).await?;
+            let pairs = UniswapV2Factory::get_all_pairs(
+                self.address,
+                from_block,
+                to_block,
+                provider.clone(),
+            )
+            .await?;
 
             Ok(pairs
                 .into_iter()

src/amms/uniswap_v3/mod.rs

@@ -27,7 +27,6 @@ use std::{
     hash::Hash,
     str::FromStr,
     sync::Arc,
-    u8,
 };
 use thiserror::Error;
 use tracing::info;
@@ -192,6 +191,10 @@ impl AutomatedMarketMaker for UniswapV3Pool {
         self.address
     }
 
+    fn initialized(&self) -> bool {
+        !self.token_a.address.is_zero() && !self.token_b.address.is_zero()
+    }
+
     fn sync_events(&self) -> Vec<B256> {
         vec![
             IUniswapV3PoolEvents::Mint::SIGNATURE_HASH,
@@ -764,6 +767,7 @@ impl UniswapV3Factory {
 
     pub async fn get_all_pools<T, N, P>(
         &self,
+        from_block: Option<BlockId>,
         block_number: BlockId,
         provider: Arc<P>,
     ) -> Result<Vec<AMM>, AMMError>
@@ -780,7 +784,8 @@ impl UniswapV3Factory {
         let mut futures = FuturesUnordered::new();
 
         let sync_step = 100_000;
-        let mut latest_block = self.creation_block;
+        let mut latest_block =
+            from_block.map_or(self.creation_block(), |b| b.as_u64().unwrap_or_default());
         while latest_block < block_number.as_u64().unwrap_or_default() {
             let mut block_filter = disc_filter.clone();
             let from_block = latest_block;
@@ -811,15 +816,15 @@ impl UniswapV3Factory {
     // TODO: update this to use uv3 error and then use thiserror to convert to AMMError
     pub async fn sync_all_pools<T, N, P>(
         mut pools: Vec<AMM>,
-        block_number: BlockId,
+        to_block: BlockId,
         provider: Arc<P>,
     ) -> Result<Vec<AMM>, AMMError>
     where
         T: Transport + Clone,
         N: Network,
         P: Provider<T, N>,
     {
-        UniswapV3Factory::sync_slot_0(&mut pools, block_number, provider.clone()).await?;
+        UniswapV3Factory::sync_slot_0(&mut pools, to_block, provider.clone()).await?;
         UniswapV3Factory::sync_token_decimals(&mut pools, provider.clone()).await;
 
         pools = pools
@@ -834,8 +839,8 @@ impl UniswapV3Factory {
             })
             .collect();
 
-        UniswapV3Factory::sync_tick_bitmaps(&mut pools, block_number, provider.clone()).await?;
-        UniswapV3Factory::sync_tick_data(&mut pools, block_number, provider.clone()).await?;
+        UniswapV3Factory::sync_tick_bitmaps(&mut pools, to_block, provider.clone()).await?;
+        UniswapV3Factory::sync_tick_data(&mut pools, to_block, provider.clone()).await?;
 
         Ok(pools)
     }
@@ -1216,6 +1221,7 @@ impl AutomatedMarketMakerFactory for UniswapV3Factory {
 impl DiscoverySync for UniswapV3Factory {
     fn discover<T, N, P>(
         &self,
+        from_block: Option<BlockId>,
         to_block: BlockId,
         provider: Arc<P>,
     ) -> impl Future<Output = Result<Vec<AMM>, AMMError>>
@@ -1230,7 +1236,7 @@ impl DiscoverySync for UniswapV3Factory {
             "Discovering all pools"
         );
 
-        self.get_all_pools(to_block, provider.clone())
+        self.get_all_pools(from_block, to_block, provider.clone())
     }
 
     fn sync<T, N, P>(

src/state_space/cache.rs

@@ -2,6 +2,7 @@ use std::collections::HashMap;
 
 use crate::amms::amm::{AutomatedMarketMaker, AMM};
 use arraydeque::ArrayDeque;
+use serde::{Deserialize, Serialize};
 
 #[derive(Debug)]
 
@@ -52,7 +53,7 @@ impl<const CAP: usize> StateChangeCache<CAP> {
         // If the block to unwind is greater than the latest state change in the block, exit early
         if cache
             .front()
-            .map_or(true, |latest| block_to_unwind > latest.block_number)
+            .is_none_or(|latest| block_to_unwind > latest.block_number)
         {
             return vec![];
         }
@@ -87,7 +88,7 @@ impl<const CAP: usize> StateChangeCache<CAP> {
 
 // NOTE: we can probably make this more efficient and create a state change struct for each amm rather than
 // cloning each amm when caching
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct StateChange {
     pub state_change: Vec<AMM>,
     pub block_number: u64,

src/state_space/error.rs

@@ -13,4 +13,8 @@ pub enum StateSpaceError {
     JoinError(#[from] tokio::task::JoinError),
     #[error("Block Number Does not Exist")]
     MissingBlockNumber,
+    #[error(transparent)]
+    SerdeError(#[from] serde_json::Error),
+    #[error(transparent)]
+    IOError(#[from] std::io::Error),
 }