Fix volume based cap calculation

crates/driver/src/domain/competition/solution/fee.rs

@@ -24,8 +24,8 @@
 
 use {
     super::{
-        trade,
-        trade::{ClearingPrices, Fee, Fulfillment},
+        trade::{self, ClearingPrices, Fee, Fulfillment},
+        MathError,
     },
     crate::domain::{
         competition::{
@@ -39,22 +39,19 @@ use {
 impl Fulfillment {
     /// Applies the protocol fee to the existing fulfillment creating a new one.
     pub fn with_protocol_fee(&self, prices: ClearingPrices) -> Result<Self, Error> {
-        let protocol_fee = self.protocol_fee(prices)?;
+        let protocol_fee = self.protocol_fee_in_sell_token(prices)?;
 
         // Increase the fee by the protocol fee
         let fee = match self.surplus_fee() {
             None => {
                 if !protocol_fee.is_zero() {
-                    return Err(trade::Error::ProtocolFeeOnStaticOrder.into());
+                    return Err(Error::ProtocolFeeOnStaticOrder);
                 }
                 Fee::Static
             }
-            Some(fee) => Fee::Dynamic(
-                (fee.0
-                    .checked_add(protocol_fee)
-                    .ok_or(trade::Error::Overflow)?)
-                .into(),
-            ),
+            Some(fee) => {
+                Fee::Dynamic((fee.0.checked_add(protocol_fee).ok_or(MathError::Overflow)?).into())
+            }
         };
 
         // Reduce the executed amount by the protocol fee. This is because solvers are
@@ -67,13 +64,14 @@ impl Fulfillment {
                 self.executed()
                     .0
                     .checked_sub(protocol_fee)
-                    .ok_or(trade::Error::Overflow)?,
+                    .ok_or(MathError::Overflow)?,
             ),
         };
 
         Fulfillment::new(order, executed, fee).map_err(Into::into)
     }
 
+    /// Computed protocol fee in surplus token.
     fn protocol_fee(&self, prices: ClearingPrices) -> Result<eth::U256, Error> {
         // TODO: support multiple fee policies
         if self.order().protocol_fees.len() > 1 {
@@ -113,6 +111,8 @@ impl Fulfillment {
 
     /// Computes protocol fee compared to the given limit amounts taken from
     /// the order or a quote.
+    ///
+    /// The protocol fee is computed in surplus token.
     fn calculate_fee(
         &self,
         limit_sell_amount: eth::U256,
@@ -130,6 +130,7 @@ impl Fulfillment {
         Ok(protocol_fee)
     }
 
+    /// Computes the surplus fee in the surplus token.
     fn fee_from_surplus(
         &self,
         sell_amount: eth::U256,
@@ -138,40 +139,40 @@ impl Fulfillment {
         factor: f64,
     ) -> Result<eth::U256, Error> {
         let surplus = self.surplus_over_reference_price(sell_amount, buy_amount, prices)?;
-        let surplus_in_sell_token = self.surplus_in_sell_token(surplus, prices)?;
-        apply_factor(surplus_in_sell_token, factor)
+        apply_factor(surplus, factor)
     }
 
+    /// Computes the volume based fee in surplus token
+    ///
+    /// The volume is defined as a full sell amount (including fees) for buy
+    /// order, or a full buy amount for sell order.
     fn fee_from_volume(&self, prices: ClearingPrices, factor: f64) -> Result<eth::U256, Error> {
-        let executed = self.executed().0;
-        let executed_sell_amount = match self.order().side {
-            Side::Buy => {
-                // How much `sell_token` we need to sell to buy `executed` amount of `buy_token`
-                executed
-                    .checked_mul(prices.buy)
-                    .ok_or(trade::Error::Overflow)?
-                    .checked_div(prices.sell)
-                    .ok_or(trade::Error::DivisionByZero)?
-            }
-            Side::Sell => executed,
+        let volume = match self.order().side {
+            Side::Buy => self.sell_amount(&prices)?,
+            Side::Sell => self.buy_amount(&prices)?,
         };
-        // Sell slightly more `sell_token` to capture the `surplus_fee`
-        let executed_sell_amount_with_fee = executed_sell_amount
-            .checked_add(
-                // surplus_fee is always expressed in sell token
-                self.surplus_fee()
-                    .map(|fee| fee.0)
-                    .ok_or(trade::Error::ProtocolFeeOnStaticOrder)?,
-            )
-            .ok_or(trade::Error::Overflow)?;
-        apply_factor(executed_sell_amount_with_fee, factor)
+        apply_factor(volume.0, factor)
+    }
+
+    /// Returns the protocol fee denominated in the sell token.
+    fn protocol_fee_in_sell_token(&self, prices: ClearingPrices) -> Result<eth::U256, Error> {
+        let fee = self.protocol_fee(prices)?;
+        let fee_in_sell_token = match self.order().side {
+            Side::Buy => fee,
+            Side::Sell => fee
+                .checked_mul(prices.buy)
+                .ok_or(MathError::Overflow)?
+                .checked_div(prices.sell)
+                .ok_or(MathError::DivisionByZero)?,
+        };
+        Ok(fee_in_sell_token)
     }
 }
 
 fn apply_factor(amount: eth::U256, factor: f64) -> Result<eth::U256, Error> {
     Ok(amount
         .checked_mul(eth::U256::from_f64_lossy(factor * 10000.))
-        .ok_or(trade::Error::Overflow)?
+        .ok_or(MathError::Overflow)?
         / 10000)
 }
 
@@ -204,24 +205,24 @@ fn adjust_quote_to_order_limits(
 ) -> Result<(eth::U256, eth::U256), Error> {
     let quote_sell_amount = quote_sell_amount
         .checked_add(quote_fee_amount)
-        .ok_or(trade::Error::Overflow)?;
+        .ok_or(MathError::Overflow)?;
 
     match order_side {
         Side::Sell => {
             let scaled_buy_amount = quote_buy_amount
                 .checked_mul(order_sell_amount)
-                .ok_or(trade::Error::Overflow)?
+                .ok_or(MathError::Overflow)?
                 .checked_div(quote_sell_amount)
-                .ok_or(trade::Error::DivisionByZero)?;
+                .ok_or(MathError::DivisionByZero)?;
             let buy_amount = order_buy_amount.max(scaled_buy_amount);
             Ok((order_sell_amount, buy_amount))
         }
         Side::Buy => {
             let scaled_sell_amount = quote_sell_amount
                 .checked_mul(order_buy_amount)
-                .ok_or(trade::Error::Overflow)?
+                .ok_or(MathError::Overflow)?
                 .checked_div(quote_buy_amount)
-                .ok_or(trade::Error::DivisionByZero)?;
+                .ok_or(MathError::DivisionByZero)?;
             let sell_amount = order_sell_amount.min(scaled_sell_amount);
             Ok((sell_amount, order_buy_amount))
         }
@@ -232,6 +233,10 @@ fn adjust_quote_to_order_limits(
 pub enum Error {
     #[error("multiple fee policies are not supported yet")]
     MultipleFeePolicies,
+    #[error("orders with non solver determined gas cost fees are not supported")]
+    ProtocolFeeOnStaticOrder,
+    #[error(transparent)]
+    Math(#[from] super::MathError),
     #[error(transparent)]
     Fulfillment(#[from] trade::Error),
 }

crates/driver/src/domain/competition/solution/mod.rs

@@ -322,3 +322,11 @@ pub enum SolutionError {
     #[error(transparent)]
     ProtocolFee(#[from] fee::Error),
 }
+
+#[derive(Debug, thiserror::Error)]
+pub enum MathError {
+    #[error("overflow")]
+    Overflow,
+    #[error("division by zero")]
+    DivisionByZero,
+}

crates/driver/src/domain/competition/solution/settlement.rs

@@ -1,5 +1,5 @@
 use {
-    super::{Error, Solution},
+    super::{trade::ClearingPrices, Error, Solution},
     crate::{
         boundary,
         domain::{
@@ -310,16 +310,22 @@ impl Settlement {
             .fold(Default::default(), |mut acc, solution| {
                 for trade in solution.user_trades() {
                     let order = acc.entry(trade.order().uid).or_default();
-                    order.sell = trade.sell_amount(&solution.prices, solution.weth).unwrap_or_else(|| {
+                    let prices = ClearingPrices {
+                        sell: solution.prices
+                            [&trade.order().sell.token.wrap(solution.weth)],
+                        buy: solution.prices
+                            [&trade.order().buy.token.wrap(solution.weth)],
+                    };
+                    order.sell = trade.sell_amount(&prices).unwrap_or_else(|err| {
                         // This should never happen, returning 0 is better than panicking, but we
                         // should still alert.
-                        tracing::error!(?trade, prices=?solution.prices, "could not compute sell_amount");
+                        tracing::error!(?trade, prices=?solution.prices, ?err, "could not compute sell_amount");
                         0.into()
                     });
-                    order.buy = trade.buy_amount(&solution.prices, solution.weth).unwrap_or_else(|| {
+                    order.buy = trade.buy_amount(&prices).unwrap_or_else(|err| {
                         // This should never happen, returning 0 is better than panicking, but we
                         // should still alert.
-                        tracing::error!(?trade, prices=?solution.prices, "could not compute buy_amount");
+                        tracing::error!(?trade, prices=?solution.prices, ?err, "could not compute buy_amount");
                         0.into()
                     });
                 }

crates/driver/src/domain/competition/solution/trade.rs

@@ -1,12 +1,12 @@
 use {
+    super::MathError,
     crate::domain::{
         competition::{
             self,
             order::{self, Side},
         },
         eth,
     },
-    std::collections::HashMap,
 };
 
 /// A trade which executes an order as part of this solution.
@@ -102,37 +102,37 @@ impl Fulfillment {
     }
 
     /// The effective amount that left the user's wallet including all fees.
-    pub fn sell_amount(
-        &self,
-        prices: &HashMap<eth::TokenAddress, eth::U256>,
-        weth: eth::WethAddress,
-    ) -> Option<eth::TokenAmount> {
+    pub fn sell_amount(&self, prices: &ClearingPrices) -> Result<eth::TokenAmount, Error> {
         let before_fee = match self.order.side {
             order::Side::Sell => self.executed.0,
             order::Side::Buy => self
                 .executed
                 .0
-                .checked_mul(*prices.get(&self.order.buy.token.wrap(weth))?)?
-                .checked_div(*prices.get(&self.order.sell.token.wrap(weth))?)?,
+                .checked_mul(prices.buy)
+                .ok_or(MathError::Overflow)?
+                .checked_div(prices.sell)
+                .ok_or(MathError::DivisionByZero)?,
         };
-        Some(eth::TokenAmount(before_fee.checked_add(self.fee().0)?))
+        Ok(eth::TokenAmount(
+            before_fee
+                .checked_add(self.fee().0)
+                .ok_or(MathError::Overflow)?,
+        ))
     }
 
     /// The effective amount the user received after all fees.
-    pub fn buy_amount(
-        &self,
-        prices: &HashMap<eth::TokenAddress, eth::U256>,
-        weth: eth::WethAddress,
-    ) -> Option<eth::TokenAmount> {
+    pub fn buy_amount(&self, prices: &ClearingPrices) -> Result<eth::TokenAmount, Error> {
         let amount = match self.order.side {
             order::Side::Buy => self.executed.0,
             order::Side::Sell => self
                 .executed
                 .0
-                .checked_mul(*prices.get(&self.order.sell.token.wrap(weth))?)?
-                .checked_div(*prices.get(&self.order.buy.token.wrap(weth))?)?,
+                .checked_mul(prices.sell)
+                .ok_or(MathError::Overflow)?
+                .checked_div(prices.buy)
+                .ok_or(MathError::DivisionByZero)?,
         };
-        Some(eth::TokenAmount(amount))
+        Ok(eth::TokenAmount(amount))
     }
 
     /// Returns the surplus denominated in the surplus token.
@@ -151,9 +151,9 @@ impl Fulfillment {
                 // How much `sell_token` we need to sell to buy `executed` amount of `buy_token`
                 executed
                     .checked_mul(prices.buy)
-                    .ok_or(Error::Overflow)?
+                    .ok_or(MathError::Overflow)?
                     .checked_div(prices.sell)
-                    .ok_or(Error::DivisionByZero)?
+                    .ok_or(MathError::DivisionByZero)?
             }
             Side::Sell => executed,
         };
@@ -165,15 +165,15 @@ impl Fulfillment {
                     .map(|fee| fee.0)
                     .ok_or(Error::ProtocolFeeOnStaticOrder)?,
             )
-            .ok_or(Error::Overflow)?;
+            .ok_or(MathError::Overflow)?;
         let surplus = match self.order().side {
             Side::Buy => {
                 // Scale to support partially fillable orders
                 let limit_sell_amount = limit_sell
                     .checked_mul(executed)
-                    .ok_or(Error::Overflow)?
+                    .ok_or(MathError::Overflow)?
                     .checked_div(limit_buy)
-                    .ok_or(Error::DivisionByZero)?;
+                    .ok_or(MathError::DivisionByZero)?;
                 // Remaining surplus after fees
                 // Do not return error if `checked_sub` fails because violated limit prices will
                 // be caught by simulation
@@ -185,15 +185,15 @@ impl Fulfillment {
                 // Scale to support partially fillable orders
                 let limit_buy_amount = limit_buy
                     .checked_mul(executed_sell_amount_with_fee)
-                    .ok_or(Error::Overflow)?
+                    .ok_or(MathError::Overflow)?
                     .checked_div(limit_sell)
-                    .ok_or(Error::DivisionByZero)?;
+                    .ok_or(MathError::DivisionByZero)?;
                 // How much `buy_token` we get for `executed` amount of `sell_token`
                 let executed_buy_amount = executed
                     .checked_mul(prices.sell)
-                    .ok_or(Error::Overflow)?
+                    .ok_or(MathError::Overflow)?
                     .checked_div(prices.buy)
-                    .ok_or(Error::DivisionByZero)?;
+                    .ok_or(MathError::DivisionByZero)?;
                 // Remaining surplus after fees
                 // Do not return error if `checked_sub` fails because violated limit prices will
                 // be caught by simulation
@@ -204,23 +204,6 @@ impl Fulfillment {
         };
         Ok(surplus)
     }
-
-    /// Returns the surplus denominated in the sell token.
-    pub fn surplus_in_sell_token(
-        &self,
-        surplus: eth::U256,
-        prices: ClearingPrices,
-    ) -> Result<eth::U256, Error> {
-        let surplus_in_sell_token = match self.order().side {
-            Side::Buy => surplus,
-            Side::Sell => surplus
-                .checked_mul(prices.buy)
-                .ok_or(Error::Overflow)?
-                .checked_div(prices.sell)
-                .ok_or(Error::DivisionByZero)?,
-        };
-        Ok(surplus_in_sell_token)
-    }
 }
 
 /// A fee that is charged for executing an order.
@@ -291,10 +274,8 @@ pub struct Execution {
 pub enum Error {
     #[error("orders with non solver determined gas cost fees are not supported")]
     ProtocolFeeOnStaticOrder,
-    #[error("overflow error while calculating protocol fee")]
-    Overflow,
-    #[error("division by zero error while calculating protocol fee")]
-    DivisionByZero,
     #[error("invalid executed amount")]
     InvalidExecutedAmount,
+    #[error(transparent)]
+    Math(#[from] super::MathError),
 }

crates/driver/src/tests/cases/protocol_fees.rs

@@ -144,7 +144,7 @@ async fn surplus_protocol_fee_sell_order_capped() {
         quote_buy_amount: 40000000000000000000u128.into(),
         executed: 40000000000000000000u128.into(),
         executed_sell_amount: 50000000000000000000u128.into(),
-        executed_buy_amount: 35000000000000000000u128.into(),
+        executed_buy_amount: 36000000000000000000u128.into(),
     };
 
     protocol_fee_test_case(test_case).await;
@@ -182,7 +182,7 @@ async fn volume_protocol_fee_sell_order() {
         quote_buy_amount: 40000000000000000000u128.into(),
         executed: 40000000000000000000u128.into(),
         executed_sell_amount: 50000000000000000000u128.into(),
-        executed_buy_amount: 15000000000000000000u128.into(),
+        executed_buy_amount: 20000000000000000000u128.into(),
     };
 
     protocol_fee_test_case(test_case).await;