feat: types re-export and error type

anthias-labs · Aug 29, 2024 · bd9ba2e · bd9ba2e
1 parent 7388c41
commit bd9ba2e
Show file tree

Hide file tree

Showing 5 changed files with 56 additions and 23 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -14,6 +14,7 @@ rand = "0.8.5"
 plotly = "0.9.0"
 serde_json = "1.0"
 rand_distr = "0.4.3"
+thiserror = "1.0.63"
 async-trait = "0.1.81"
 alloy-contract = "0.2.1"
 alloy-sol-macro = "0.7.7"

diff --git a/src/arena.rs b/src/arena.rs
@@ -6,6 +6,7 @@ use super::*;
 use crate::{
     config::Config,
     engine::{arbitrageur::Arbitrageur, inspector::Inspector},
+    error::ArenaError,
     feed::Feed,
     strategy::Strategy,
     types::{ArenaToken, Fetcher, LiquidExchange, PoolManager, PoolManager::PoolKey},
@@ -38,14 +39,16 @@ pub struct Arena<V> {
 
 impl<V> Arena<V> {
     /// Run all strategies in the simulation with a given configuration.
-    pub async fn run(&mut self, config: Config) {
+    pub async fn run(&mut self, config: Config) -> Result<(), ArenaError> {
         let admin_provider = self.providers[&0].clone();
 
         let pool_manager = PoolManager::deploy(admin_provider.clone(), U256::from(0))
             .await
-            .unwrap();
+            .map_err(ArenaError::ContractError)?;
 
-        let fetcher = Fetcher::deploy(admin_provider.clone()).await.unwrap();
+        let fetcher = Fetcher::deploy(admin_provider.clone())
+            .await
+            .map_err(ArenaError::ContractError)?;
 
         let currency_0 = ArenaToken::deploy(
             admin_provider.clone(),
@@ -54,7 +57,7 @@ impl<V> Arena<V> {
             18,
         )
         .await
-        .unwrap();
+        .map_err(ArenaError::ContractError)?;
 
         let currency_1 = ArenaToken::deploy(
             admin_provider.clone(),
@@ -63,7 +66,7 @@ impl<V> Arena<V> {
             18,
         )
         .await
-        .unwrap();
+        .map_err(ArenaError::ContractError)?;
 
         let liquid_exchange = LiquidExchange::deploy(
             admin_provider.clone(),
@@ -72,7 +75,7 @@ impl<V> Arena<V> {
             U256::from(1),
         )
         .await
-        .unwrap();
+        .map_err(ArenaError::ContractError)?;
 
         if *currency_1.address() > *currency_0.address() {
             (self.pool.currency0, self.pool.currency1) =
@@ -88,21 +91,25 @@ impl<V> Arena<V> {
             .nonce(5)
             .send()
             .await
-            .unwrap()
+            .map_err(ArenaError::ContractError)?
             .watch()
             .await
-            .unwrap();
+            .map_err(|e| ArenaError::ContractError(alloy_contract::Error::TransportError(e)))?;
 
         let mut signal = Signal::default();
 
         for (idx, strategy) in self.strategies.iter_mut().enumerate() {
-            let id = fetcher.toId(self.pool.clone().into()).call().await.unwrap();
+            let id = fetcher
+                .toId(self.pool.clone().into())
+                .call()
+                .await
+                .map_err(ArenaError::ContractError)?;
 
             let slot = fetcher
                 .getSlot0(*pool_manager.address(), id.poolId)
                 .call()
                 .await
-                .unwrap();
+                .map_err(ArenaError::ContractError)?;
 
             signal = Signal::new(
                 *pool_manager.address(),
@@ -125,13 +132,17 @@ impl<V> Arena<V> {
         self.nonce = 6;
 
         for step in 0..config.steps {
-            let id = fetcher.toId(self.pool.clone().into()).call().await.unwrap();
+            let id = fetcher
+                .toId(self.pool.clone().into())
+                .call()
+                .await
+                .map_err(ArenaError::ContractError)?;
 
             let slot = fetcher
                 .getSlot0(*pool_manager.address(), id.poolId)
                 .call()
                 .await
-                .unwrap();
+                .map_err(ArenaError::ContractError)?;
 
             let signal = Signal::new(
                 *pool_manager.address(),
@@ -145,15 +156,16 @@ impl<V> Arena<V> {
 
             liquid_exchange
                 .setPrice(
-                    alloy::primitives::utils::parse_ether(&self.feed.step().to_string()).unwrap(),
+                    alloy::primitives::utils::parse_ether(&self.feed.step().to_string())
+                        .map_err(ArenaError::ConversionError)?,
                 )
                 .nonce(self.nonce)
                 .send()
                 .await
-                .unwrap()
+                .map_err(ArenaError::ContractError)?
                 .watch()
                 .await
-                .unwrap();
+                .map_err(|e| ArenaError::ContractError(alloy_contract::Error::TransportError(e)))?;
 
             self.nonce += 1;
 
@@ -171,6 +183,8 @@ impl<V> Arena<V> {
 
             self.feed.step();
         }
+
+        Ok(())
     }
 }
 

diff --git a/src/engine/arbitrageur.rs b/src/engine/arbitrageur.rs
@@ -15,7 +15,7 @@ use crate::{
 /// Generic trait allowing user defined arbitrage strategies.
 #[async_trait]
 pub trait Arbitrageur {
-    /// Initialixe arbitrageur agent.
+    /// Initialize arbitrageur agent.
     async fn init(&mut self, signal: &Signal, provider: AnvilProvider);
 
     /// Perform an arbitrage based on a [`Signal`].

diff --git a/src/lib.rs b/src/lib.rs
@@ -16,6 +16,9 @@ pub mod strategy;
 /// Defines core simulation logic types, such as an [`Arbitrageur`].
 pub mod engine;
 
+/// Contains error types for Arena.
+pub mod error;
+
 use alloy::{
     network::{Ethereum, EthereumWallet},
     node_bindings::{Anvil, AnvilInstance},
@@ -26,8 +29,18 @@ use alloy::{
     },
     transports::http::{Client, Http},
 };
-
-use crate::{engine::inspector::Inspector, types::PoolManager::PoolKey};
+use types::PoolManager::PoolKey;
+
+pub use crate::{
+    arena::{Arena, ArenaBuilder},
+    config::Config,
+    engine::{
+        arbitrageur::{Arbitrageur, DefaultArbitrageur, EmptyArbitrageur},
+        inspector::{EmptyInspector, Inspector, LogMessage, Logger},
+    },
+    feed::{Feed, GeometricBrownianMotion, OrnsteinUhlenbeck},
+    strategy::Strategy,
+};
 
 /// Provider type that includes all necessary fillers to execute transactions on an [`Anvil`] node.
 pub type AnvilProvider = FillProvider<
@@ -84,6 +97,13 @@ mod types {
         "src/artifacts/PoolSwapTest.json"
     }
 
+    sol! {
+        #[sol(rpc)]
+        #[derive(Debug)]
+        PoolModifyLiquidityTest,
+        "src/artifacts/PoolModifyLiquidityTest.json"
+    }
+
     impl From<FetcherPoolKey> for ManagerPoolKey {
         fn from(fetcher: FetcherPoolKey) -> Self {
             ManagerPoolKey {
@@ -211,10 +231,7 @@ mod tests {
             _signal: Signal,
             inspector: &mut Box<dyn Inspector<LogMessage>>,
         ) {
-            inspector.log(LogMessage::new(
-                String::from("test_key"),
-                String::from("test_value"),
-            ));
+            inspector.log(LogMessage::new(String::from("test"), String::from("test")));
         }
     }
 
@@ -227,7 +244,7 @@ mod tests {
             .with_fee(4000)
             .with_tick_spacing(2)
             .with_feed(Box::new(OrnsteinUhlenbeck::new(0.1, 0.1, 0.1, 0.1, 0.1)))
-            .with_inspector(Box::new(Logger::new_csv(String::from("test.csv"))))
+            .with_inspector(Box::new(Logger::new_json(String::from("test1.json"))))
             .with_arbitrageur(Box::new(EmptyArbitrageur {}))
             .build();