Skip to content

Commit

Permalink
feat: update automated market maker trait to support multi token pools
Browse files Browse the repository at this point in the history
  • Loading branch information
@0xOsiris
0xOsiris committed Aug 7, 2024
1 parent 643b049 commit 684133a
Show file tree
Hide file tree
Showing 7 changed files with 200 additions and 110 deletions.
2 changes: 2 additions & 0 deletions examples/filter-value.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@ async fn main() -> eyre::Result<()> {

// Filter out pools below usd threshold
let weth_address = address!("0d500B1d8E8eF31E21C99d1Db9A6444d3ADf1270");
let usdc = address!("3c499c542cEF5E3811e1192ce70d8cC03d5c3359");
let usd_weth_pair_address = address!("cd353F79d9FADe311fC3119B841e1f456b54e858");
let usd_weth_pool = AMM::UniswapV2Pool(
UniswapV2Pool::new_from_address(usd_weth_pair_address, 300, provider.clone()).await?,
Expand All @@ -61,6 +62,7 @@ async fn main() -> eyre::Result<()> {
usd_weth_pool,
15000.00, //Setting usd_threshold to 15000 filters out any pool that contains less than $15000.00 USD value
weth_address,
usdc,
weth_value_in_token_to_weth_pool_threshold,
200,
provider.clone(),
Expand Down
8 changes: 6 additions & 2 deletions examples/simulate-swap.rs
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
use std::sync::Arc;

use alloy::{
primitives::{address, U256},
primitives::{address, Address, U256},
providers::ProviderBuilder,
};

Expand All @@ -20,7 +20,11 @@ async fn main() -> eyre::Result<()> {

// Simulate a swap
let token_in = address!("C02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2");
let amount_out = pool.simulate_swap(token_in, U256::from(1000000000000000000_u128))?;
let amount_out = pool.simulate_swap(
token_in,
Address::default(),
U256::from(1000000000000000000_u128),
)?;

println!("Amount out: {amount_out}");

Expand Down
64 changes: 41 additions & 23 deletions src/amm/erc_4626/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,11 @@ impl AutomatedMarketMaker for ERC4626Vault {
vec![self.vault_token, self.asset_token]
}

fn calculate_price(&self, base_token: Address) -> Result<f64, ArithmeticError> {
fn calculate_price(
&self,
base_token: Address,
_quote_token: Address,
) -> Result<f64, ArithmeticError> {
Ok(q64_to_f64(self.calculate_price_64_x_64(base_token)?))
}

Expand Down Expand Up @@ -128,10 +132,11 @@ impl AutomatedMarketMaker for ERC4626Vault {

fn simulate_swap(
&self,
token_in: Address,
base_token: Address,
_quote_token: Address,
amount_in: U256,
) -> Result<U256, SwapSimulationError> {
if self.vault_token == token_in {
if self.vault_token == base_token {
Ok(self.get_amount_out(amount_in, self.vault_reserve, self.asset_reserve))
} else {
Ok(self.get_amount_out(amount_in, self.asset_reserve, self.vault_reserve))
Expand All @@ -140,10 +145,11 @@ impl AutomatedMarketMaker for ERC4626Vault {

fn simulate_swap_mut(
&mut self,
token_in: Address,
base_token: Address,
_quote_token: Address,
amount_in: U256,
) -> Result<U256, SwapSimulationError> {
if self.vault_token == token_in {
if self.vault_token == base_token {
let amount_out = self.get_amount_out(amount_in, self.vault_reserve, self.asset_reserve);

self.vault_reserve -= amount_in;
Expand All @@ -159,14 +165,6 @@ impl AutomatedMarketMaker for ERC4626Vault {
Ok(amount_out)
}
}

fn get_token_out(&self, token_in: Address) -> Address {
if self.vault_token == token_in {
self.asset_token
} else {
self.vault_token
}
}
}

impl ERC4626Vault {
Expand Down Expand Up @@ -308,10 +306,10 @@ impl ERC4626Vault {

#[cfg(test)]
mod tests {
use std::sync::Arc;
use std::{ops::Add, sync::Arc};

use alloy::{
primitives::{address, U256},
primitives::{address, Address, U256},
providers::ProviderBuilder,
};

Expand Down Expand Up @@ -357,8 +355,12 @@ mod tests {
vault.asset_token_decimals = 6;
vault.asset_reserve = U256::from(505434849031_u64);

let price_v_64_x = vault.calculate_price(vault.vault_token).unwrap();
let price_a_64_x = vault.calculate_price(vault.asset_token).unwrap();
let price_v_64_x = vault
.calculate_price(vault.vault_token, vault.asset_token)
.unwrap();
let price_a_64_x = vault
.calculate_price(vault.asset_token, vault.vault_token)
.unwrap();

assert_eq!(price_v_64_x, 1.0070222372637234);
assert_eq!(price_a_64_x, 0.99302673068789);
Expand All @@ -379,8 +381,12 @@ mod tests {
vault.vault_reserve = U256::ZERO;
vault.asset_reserve = U256::ZERO;

let price_v_64_x = vault.calculate_price(vault.vault_token).unwrap();
let price_a_64_x = vault.calculate_price(vault.asset_token).unwrap();
let price_v_64_x = vault
.calculate_price(vault.vault_token, Address::default())
.unwrap();
let price_a_64_x = vault
.calculate_price(vault.asset_token, Address::default())
.unwrap();

assert_eq!(price_v_64_x, 1.0);
assert_eq!(price_a_64_x, 1.0);
Expand All @@ -401,8 +407,12 @@ mod tests {
vault.vault_reserve = U256::from(501910315708981197269904_u128);
vault.asset_reserve = U256::from(505434849031054568651911_u128);

let price_v_64_x = vault.calculate_price(vault.vault_token).unwrap();
let price_a_64_x = vault.calculate_price(vault.asset_token).unwrap();
let price_v_64_x = vault
.calculate_price(vault.vault_token, Address::default())
.unwrap();
let price_a_64_x = vault
.calculate_price(vault.asset_token, Address::default())
.unwrap();

assert_eq!(price_v_64_x, 1.0070222372638322);
assert_eq!(price_a_64_x, 0.9930267306877828);
Expand Down Expand Up @@ -446,10 +456,18 @@ mod tests {
vault.asset_reserve = U256::from(505434849031054568651911_u128);

let assets_out = vault
.simulate_swap(vault.vault_token, U256::from(3000000000000000000_u128))
.simulate_swap(
vault.vault_token,
vault.asset_token,
U256::from(3000000000000000000_u128),
)
.unwrap();
let shares_out = vault
.simulate_swap(vault.asset_token, U256::from(3000000000000000000_u128))
.simulate_swap(
vault.asset_token,
vault.vault_token,
U256::from(3000000000000000000_u128),
)
.unwrap();

assert_eq!(assets_out, U256::from(3021066711791496478_u128));
Expand Down
33 changes: 15 additions & 18 deletions src/amm/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,11 @@ pub trait AutomatedMarketMaker {
fn tokens(&self) -> Vec<Address>;

/// Calculates a f64 representation of base token price in the AMM.
fn calculate_price(&self, base_token: Address) -> Result<f64, ArithmeticError>;
fn calculate_price(
&self,
base_token: Address,
quote_token: Address,
) -> Result<f64, ArithmeticError>;

/// Updates the AMM data from a log.
fn sync_from_log(&mut self, log: Log) -> Result<(), EventLogError>;
Expand All @@ -71,7 +75,8 @@ pub trait AutomatedMarketMaker {
/// Returns the amount received for `amount_in` of `token_in`.
fn simulate_swap(
&self,
token_in: Address,
base_token: Address,
quote_token: Address,
amount_in: U256,
) -> Result<U256, SwapSimulationError>;

Expand All @@ -80,12 +85,10 @@ pub trait AutomatedMarketMaker {
/// Returns the amount received for `amount_in` of `token_in`.
fn simulate_swap_mut(
&mut self,
token_in: Address,
base_token: Address,
quote_token: Address,
amount_in: U256,
) -> Result<U256, SwapSimulationError>;

/// Returns the token out of the AMM for a given `token_in`.
fn get_token_out(&self, token_in: Address) -> Address;
}

macro_rules! amm {
Expand Down Expand Up @@ -126,21 +129,15 @@ macro_rules! amm {
}
}

fn simulate_swap(&self, token_in: Address, amount_in: U256) -> Result<U256, SwapSimulationError> {
match self {
$(AMM::$pool_type(pool) => pool.simulate_swap(token_in, amount_in),)+
}
}

fn simulate_swap_mut(&mut self, token_in: Address, amount_in: U256) -> Result<U256, SwapSimulationError> {
fn simulate_swap(&self, base_token: Address, quote_token: Address,amount_in: U256) -> Result<U256, SwapSimulationError> {
match self {
$(AMM::$pool_type(pool) => pool.simulate_swap_mut(token_in, amount_in),)+
$(AMM::$pool_type(pool) => pool.simulate_swap(base_token, quote_token, amount_in),)+
}
}

fn get_token_out(&self, token_in: Address) -> Address {
fn simulate_swap_mut(&mut self, base_token: Address, quote_token: Address, amount_in: U256) -> Result<U256, SwapSimulationError> {
match self {
$(AMM::$pool_type(pool) => pool.get_token_out(token_in),)+
$(AMM::$pool_type(pool) => pool.simulate_swap_mut(base_token, quote_token, amount_in),)+
}
}

Expand All @@ -161,9 +158,9 @@ macro_rules! amm {
}
}

fn calculate_price(&self, base_token: Address) -> Result<f64, ArithmeticError> {
fn calculate_price(&self, base_token: Address, quote_token: Address) -> Result<f64, ArithmeticError> {
match self {
$(AMM::$pool_type(pool) => pool.calculate_price(base_token),)+
$(AMM::$pool_type(pool) => pool.calculate_price(base_token, quote_token),)+
}
}
}
Expand Down
40 changes: 21 additions & 19 deletions src/amm/uniswap_v2/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -108,7 +108,11 @@ impl AutomatedMarketMaker for UniswapV2Pool {
}

// Calculates base/quote, meaning the price of base token per quote (ie. exchange rate is X base per 1 quote)
fn calculate_price(&self, base_token: Address) -> Result<f64, ArithmeticError> {
fn calculate_price(
&self,
base_token: Address,
_quote_token: Address,
) -> Result<f64, ArithmeticError> {
Ok(q64_to_f64(self.calculate_price_64_x_64(base_token)?))
}

Expand All @@ -118,10 +122,11 @@ impl AutomatedMarketMaker for UniswapV2Pool {

fn simulate_swap(
&self,
token_in: Address,
base_token: Address,
_quote_token: Address,
amount_in: U256,
) -> Result<U256, SwapSimulationError> {
if self.token_a == token_in {
if self.token_a == base_token {
Ok(self.get_amount_out(
amount_in,
U256::from(self.reserve_0),
Expand All @@ -138,10 +143,11 @@ impl AutomatedMarketMaker for UniswapV2Pool {

fn simulate_swap_mut(
&mut self,
token_in: Address,
base_token: Address,
_quote_token: Address,
amount_in: U256,
) -> Result<U256, SwapSimulationError> {
if self.token_a == token_in {
if self.token_a == base_token {
let amount_out = self.get_amount_out(
amount_in,
U256::from(self.reserve_0),
Expand Down Expand Up @@ -175,14 +181,6 @@ impl AutomatedMarketMaker for UniswapV2Pool {
Ok(amount_out)
}
}

fn get_token_out(&self, token_in: Address) -> Address {
if self.token_a == token_in {
self.token_b
} else {
self.token_a
}
}
}

impl UniswapV2Pool {
Expand Down Expand Up @@ -555,10 +553,10 @@ pub fn q64_to_f64(x: u128) -> f64 {

#[cfg(test)]
mod tests {
use std::sync::Arc;
use std::{ops::Add, sync::Arc};

use alloy::{
primitives::{address, U256},
primitives::{address, Address, U256},
providers::ProviderBuilder,
};

Expand Down Expand Up @@ -651,8 +649,8 @@ mod tests {
fee: 300,
};

assert!(x.calculate_price(token_a).unwrap() != 0.0);
assert!(x.calculate_price(token_b).unwrap() != 0.0);
assert!(x.calculate_price(token_a, Address::default()).unwrap() != 0.0);
assert!(x.calculate_price(token_b, Address::default()).unwrap() != 0.0);
}

#[tokio::test]
Expand All @@ -670,8 +668,12 @@ mod tests {
pool.reserve_0 = 47092140895915;
pool.reserve_1 = 28396598565590008529300;

let price_a_64_x = pool.calculate_price(pool.token_a).unwrap();
let price_b_64_x = pool.calculate_price(pool.token_b).unwrap();
let price_a_64_x = pool
.calculate_price(pool.token_a, Address::default())
.unwrap();
let price_b_64_x = pool
.calculate_price(pool.token_b, Address::default())
.unwrap();

// No precision loss: 30591574867092394336528 / 2**64
assert_eq!(1658.3725965327264, price_b_64_x);
Expand Down
Loading

0 comments on commit 684133a

Please sign in to comment.