Merge pull request #211 from darkforestry/osi/balancer

feat(balancer): simulate_swap simulate_swap_mut
darkforestry · Aug 30, 2024 · 4bd8a8a · 4bd8a8a
2 parents 3c0cd86 + 309e64e
commit 4bd8a8a
Show file tree

Hide file tree

Showing 10 changed files with 163 additions and 48 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -36,7 +36,6 @@ alloy = { version = "0.2", features = [
     "rpc-types-eth",
     "signer-local",
 ] }
-reqwest = "0.12"
 rug = "1.24.1"
 
 [features]

diff --git a/src/amm/balancer_v2/batch_request/mod.rs b/src/amm/balancer_v2/batch_request/mod.rs
@@ -118,7 +118,7 @@ where
     let return_data_tokens = constructor_return.abi_decode_sequence(&res)?;
 
     if let Some(tokens_arr) = return_data_tokens.as_array() {
-        for (i, token) in tokens_arr.into_iter().enumerate() {
+        for (i, token) in tokens_arr.iter().enumerate() {
             let pool_data = token
                 .as_tuple()
                 .ok_or(AMMError::BatchRequestError(amms[i].address()))?;

diff --git a/src/amm/balancer_v2/bmath.rs b/src/amm/balancer_v2/bmath.rs
@@ -1,34 +1,116 @@
 use alloy::primitives::U256;
 use rug::Float;
 
-use crate::amm::consts::{BONE, DECIMAL_RADIX, MPFR_T_PRECISION, U256_2};
+use crate::amm::consts::{BONE, DECIMAL_RADIX, MPFR_T_PRECISION, U256_1, U256_10E_10, U256_2};
 
-pub fn badd(a: U256, b: U256) -> U256 {
+use super::error::BMathError;
+
+pub fn btoi(a: U256) -> U256 {
+    a / BONE
+}
+
+#[inline]
+pub fn badd(a: U256, b: U256) -> Result<U256, BMathError> {
     let c = a + b;
-    assert!(c >= a, "ERR_ADD_OVERFLOW");
-    c
+    if c < a {
+        return Err(BMathError::AddOverflow);
+    }
+    Ok(c)
+}
+
+#[inline]
+pub fn bpowi(a: U256, n: U256) -> Result<U256, BMathError> {
+    let mut z = if n % U256_2 != U256::ZERO { a } else { BONE };
+
+    let mut a = a;
+    let mut n = n / U256_2;
+    while n != U256::ZERO {
+        a = bmul(a, a)?;
+        if n % U256_2 != U256::ZERO {
+            z = bmul(z, a)?;
+        }
+        n /= U256_2;
+    }
+    Ok(z)
+}
+
+#[inline]
+pub fn bpow(base: U256, exp: U256) -> Result<U256, BMathError> {
+    let whole = bfloor(exp);
+    let remain = bsub(exp, whole)?;
+    let whole_pow = bpowi(base, btoi(whole))?;
+    if remain == U256::ZERO {
+        return Ok(whole_pow);
+    }
+    let precision = BONE / U256_10E_10;
+    let partial_result = bpow_approx(base, remain, precision)?;
+    bmul(whole_pow, partial_result)
 }
+
+#[inline]
+pub fn bpow_approx(base: U256, exp: U256, precision: U256) -> Result<U256, BMathError> {
+    let a = exp;
+    let (x, xneg) = bsub_sign(base, BONE);
+    let mut term = BONE;
+    let mut sum = term;
+    let mut negative = false;
+    let mut i = U256_1;
+    while term >= precision {
+        let big_k = U256::from(i) * BONE;
+        let (c, cneg) = bsub_sign(a, bsub(big_k, BONE)?);
+        term = bmul(term, bmul(c, x)?)?;
+        term = bdiv(term, big_k)?;
+        if term == U256::ZERO {
+            break;
+        }
+        negative ^= xneg ^ cneg;
+        if negative {
+            sum = bsub(sum, term)?;
+        } else {
+            sum = badd(sum, term)?;
+        }
+        i += U256_1;
+    }
+    Ok(sum)
+}
+
+#[inline]
+pub fn bfloor(a: U256) -> U256 {
+    btoi(a) * BONE
+}
+
 // Reference:
 // https://github.com/balancer/balancer-core/blob/f4ed5d65362a8d6cec21662fb6eae233b0babc1f/contracts/BNum.sol#L75
-pub fn bdiv(a: U256, b: U256) -> U256 {
-    assert!(b != U256::ZERO, "ERR_DIV_ZERO");
+#[inline]
+pub fn bdiv(a: U256, b: U256) -> Result<U256, BMathError> {
+    if b == U256::ZERO {
+        return Err(BMathError::DivZero);
+    }
     let c0 = a * BONE;
-    assert!(a == U256::ZERO || c0 / a == BONE, "ERR_DIV_INTERNAL");
+    if a != U256::ZERO && c0 / a != BONE {
+        return Err(BMathError::DivInternal);
+    }
     let c1 = c0 + (b / U256_2);
-    assert!(c1 >= c0, "ERR_DIV_INTERNAL");
-    c1 / b
+    if c1 < c0 {
+        return Err(BMathError::DivInternal);
+    }
+    Ok(c1 / b)
 }
 
 // Reference:
 // https://github.com/balancer/balancer-core/blob/f4ed5d65362a8d6cec21662fb6eae233b0babc1f/contracts/BNum.sol#L43
-pub fn bsub(a: U256, b: U256) -> U256 {
+#[inline]
+pub fn bsub(a: U256, b: U256) -> Result<U256, BMathError> {
     let (c, flag) = bsub_sign(a, b);
-    assert!(!flag, "ERR_SUB_UNDERFLOW");
-    c
+    if flag {
+        return Err(BMathError::SubUnderflow);
+    }
+    Ok(c)
 }
 
 // Reference:
 // https://github.com/balancer/balancer-core/blob/f4ed5d65362a8d6cec21662fb6eae233b0babc1f/contracts/BNum.sol#L52
+#[inline]
 pub fn bsub_sign(a: U256, b: U256) -> (U256, bool) {
     if a >= b {
         (a - b, false)
@@ -39,12 +121,17 @@ pub fn bsub_sign(a: U256, b: U256) -> (U256, bool) {
 
 // Reference:
 // https://github.com/balancer/balancer-core/blob/f4ed5d65362a8d6cec21662fb6eae233b0babc1f/contracts/BNum.sol#L63C4-L73C6
-pub fn bmul(a: U256, b: U256) -> U256 {
+#[inline]
+pub fn bmul(a: U256, b: U256) -> Result<U256, BMathError> {
     let c0 = a * b;
-    assert!(a == U256::ZERO || c0 / a == b, "ERR_MUL_OVERFLOW");
+    if a != U256::ZERO && c0 / a != b {
+        return Err(BMathError::MulOverflow);
+    }
     let c1 = c0 + (BONE / U256_2);
-    assert!(c1 >= c0, "ERR_MUL_OVERFLOW");
-    c1 / BONE
+    if c1 < c0 {
+        return Err(BMathError::MulOverflow);
+    }
+    Ok(c1 / BONE)
 }
 
 /**********************************************************************************************
@@ -56,11 +143,17 @@ pub fn bmul(a: U256, b: U256) -> U256 {
 // wO = tokenWeightOut                                                                       //
 // sF = swapFee                                                                              //
  **********************************************************************************************/
-pub fn calculate_price(b_i: U256, w_i: U256, b_o: U256, w_o: U256, s_f: U256) -> U256 {
-    let numer = bdiv(b_i, w_i);
-    let denom = bdiv(b_o, w_o);
-    let ratio = bdiv(numer, denom);
-    let scale = bdiv(BONE, bsub(BONE, s_f));
+pub fn calculate_price(
+    b_i: U256,
+    w_i: U256,
+    b_o: U256,
+    w_o: U256,
+    s_f: U256,
+) -> Result<U256, BMathError> {
+    let numer = bdiv(b_i, w_i)?;
+    let denom = bdiv(b_o, w_o)?;
+    let ratio = bdiv(numer, denom)?;
+    let scale = bdiv(BONE, bsub(BONE, s_f)?)?;
     bmul(ratio, scale)
 }
 
@@ -81,14 +174,14 @@ pub fn calculate_out_given_in(
     token_weight_out: U256,
     token_amount_in: U256,
     swap_fee: U256,
-) -> U256 {
-    let weight_ratio = bdiv(token_weight_in, token_weight_out);
-    let adjusted_in = bsub(BONE, swap_fee);
-    let adjusted_in = bmul(token_amount_in, adjusted_in);
-    let y = bdiv(token_balance_in, badd(token_balance_in, adjusted_in));
-    let foo = y.pow(weight_ratio);
-    let bar = bsub(BONE, foo);
-    bmul(token_balance_out, bar)
+) -> Result<U256, BMathError> {
+    let weight_ratio = bdiv(token_weight_in, token_weight_out)?;
+    let adjusted_in = bsub(BONE, swap_fee)?;
+    let adjusted_in = bmul(token_amount_in, adjusted_in)?;
+    let y = bdiv(token_balance_in, badd(token_balance_in, adjusted_in)?)?;
+    let x = bpow(y, weight_ratio)?;
+    let z = bsub(BONE, x)?;
+    bmul(token_balance_out, z)
 }
 
 /// Converts a `U256` into a `Float` with a high precision.

diff --git a/src/amm/balancer_v2/error.rs b/src/amm/balancer_v2/error.rs
@@ -0,0 +1,15 @@
+use thiserror::Error;
+
+#[derive(Error, Debug)]
+pub enum BMathError {
+    #[error("Division by zero")]
+    DivZero,
+    #[error("Error during division")]
+    DivInternal,
+    #[error("Addition overflow")]
+    AddOverflow,
+    #[error("Subtraction underflow")]
+    SubUnderflow,
+    #[error("Multiplication overflow")]
+    MulOverflow,
+}
diff --git a/src/amm/balancer_v2/mod.rs b/src/amm/balancer_v2/mod.rs
@@ -1,12 +1,13 @@
 pub mod batch_request;
-mod bmath;
+pub mod bmath;
+pub mod error;
 pub mod factory;
 
 use std::sync::Arc;
 
 use alloy::{
     network::Network,
-    primitives::{ruint::BaseConvertError, Address, B256, U256},
+    primitives::{Address, B256, U256},
     providers::Provider,
     rpc::types::Log,
     sol,
@@ -147,7 +148,7 @@ impl AutomatedMarketMaker for BalancerV2Pool {
 
         let dividend = (balance_quote / norm_weight_quote) * bone.clone();
         let divisor = (balance_base / norm_weight_base)
-            * (bone - Float::with_val(MPFR_T_PRECISION, self.fee as u32));
+            * (bone - Float::with_val(MPFR_T_PRECISION, self.fee));
         let ratio = dividend / divisor;
         Ok(ratio.to_f64_round(Round::Nearest))
     }
@@ -234,7 +235,7 @@ impl AutomatedMarketMaker for BalancerV2Pool {
             quote_token_weight,
             amount_in,
             swap_fee,
-        ))
+        )?)
     }
 
     /// Locally simulates a swap in the AMM.
@@ -283,9 +284,9 @@ impl AutomatedMarketMaker for BalancerV2Pool {
             quote_token_weight,
             amount_in,
             swap_fee,
-        );
-        self.liquidity[base_token_index] = bmath::badd(base_token_balance, amount_in);
-        self.liquidity[quote_token_index] = bmath::bsub(quote_token_balance, out);
+        )?;
+        self.liquidity[base_token_index] = bmath::badd(base_token_balance, amount_in)?;
+        self.liquidity[quote_token_index] = bmath::bsub(quote_token_balance, out)?;
         Ok(out)
     }
 }
@@ -307,8 +308,9 @@ mod tests {
             address: address!("8a649274E4d777FFC6851F13d23A86BBFA2f2Fbf"),
             ..Default::default()
         };
-        let provider =
-            Arc::new(ProviderBuilder::new().on_http(env!("ETHEREUM_PROVIDER").parse().unwrap()));
+        let provider = Arc::new(
+            ProviderBuilder::new().on_http(env!("ETHEREUM_RPC_ENDPOINT").parse().unwrap()),
+        );
         balancer_v2_pool
             .populate_data(Some(20487793), provider.clone())
             .await
@@ -334,8 +336,9 @@ mod tests {
 
     #[tokio::test]
     pub async fn test_calculate_price() {
-        let provider =
-            Arc::new(ProviderBuilder::new().on_http(env!("ETHEREUM_PROVIDER").parse().unwrap()));
+        let provider = Arc::new(
+            ProviderBuilder::new().on_http(env!("ETHEREUM_RPC_ENDPOINT").parse().unwrap()),
+        );
         let mut balancer_v2_pool = super::BalancerV2Pool {
             address: address!("8a649274E4d777FFC6851F13d23A86BBFA2f2Fbf"),
             ..Default::default()
@@ -357,8 +360,9 @@ mod tests {
 
     #[tokio::test]
     pub async fn test_simulate_swap() {
-        let provider =
-            Arc::new(ProviderBuilder::new().on_http(env!("ETHEREUM_PROVIDER").parse().unwrap()));
+        let provider = Arc::new(
+            ProviderBuilder::new().on_http(env!("ETHEREUM_RPC_ENDPOINT").parse().unwrap()),
+        );
         let mut balancer_v2_pool = super::BalancerV2Pool {
             address: address!("8a649274E4d777FFC6851F13d23A86BBFA2f2Fbf"),
             ..Default::default()
@@ -388,7 +392,7 @@ mod tests {
             .calcOutGivenIn(
                 balancer_v2_pool.liquidity[0],
                 balancer_v2_pool.weights[0],
-                balancer_v2_pool.liquidity[0],
+                balancer_v2_pool.liquidity[1],
                 balancer_v2_pool.weights[1],
                 amount_in,
                 U256::from(balancer_v2_pool.fee),

diff --git a/src/amm/consts.rs b/src/amm/consts.rs
@@ -1,6 +1,7 @@
 use alloy::primitives::U256;
 
 // commonly used U256s
+pub const U256_10E_10: U256 = U256::from_limbs([10000000000, 0, 0, 0]);
 pub const U256_0X100000000: U256 = U256::from_limbs([4294967296, 0, 0, 0]);
 pub const U256_0X10000: U256 = U256::from_limbs([65536, 0, 0, 0]);
 pub const U256_0X100: U256 = U256::from_limbs([256, 0, 0, 0]);

diff --git a/src/amm/erc_4626/mod.rs b/src/amm/erc_4626/mod.rs
@@ -306,7 +306,7 @@ impl ERC4626Vault {
 
 #[cfg(test)]
 mod tests {
-    use std::{ops::Add, sync::Arc};
+    use std::sync::Arc;
 
     use alloy::{
         primitives::{address, Address, U256},

diff --git a/src/amm/uniswap_v2/mod.rs b/src/amm/uniswap_v2/mod.rs
@@ -553,7 +553,7 @@ pub fn q64_to_f64(x: u128) -> f64 {
 
 #[cfg(test)]
 mod tests {
-    use std::{ops::Add, sync::Arc};
+    use std::sync::Arc;
 
     use alloy::{
         primitives::{address, Address, U256},

diff --git a/src/errors.rs b/src/errors.rs
@@ -6,6 +6,8 @@ use thiserror::Error;
 use tokio::task::JoinError;
 use uniswap_v3_math::error::UniswapV3MathError;
 
+use crate::amm::balancer_v2::error::BMathError;
+
 #[derive(Error, Debug)]
 pub enum AMMError {
     #[error(transparent)]
@@ -100,6 +102,8 @@ pub enum SwapSimulationError {
     LiquidityUnderflow,
     #[error(transparent)]
     ArithmeticError(#[from] ArithmeticError),
+    #[error(transparent)]
+    BMathError(#[from] BMathError),
 }
 
 #[derive(Error, Debug)]

diff --git a/src/filters/value.rs b/src/filters/value.rs
@@ -11,7 +11,6 @@ use alloy::{
 
 use crate::{
     amm::{
-        balancer_v2::factory::BalancerV2Factory,
         factory::{AutomatedMarketMakerFactory, Factory},
         AutomatedMarketMaker, AMM,
     },