Driver: Include non-surplus-capturing-jit-orders in the driver /solve

crates/driver/openapi.yml

@@ -237,7 +237,7 @@ components:
           enum: ["erc20", "internal"]
         class:
           type: string
-          enum: ["market", "limit", "liquidity"]
+          enum: ["market", "limit"]
         appData:
           description: 32 bytes encoded as hex with `0x` prefix.
           type: string

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

@@ -2,7 +2,12 @@ use {
     super::{encoding, trade::ClearingPrices, Error, Solution},
     crate::{
         domain::{
-            competition::{self, auction, order, solution},
+            competition::{
+                self,
+                auction,
+                order::{self},
+                solution::{self, error, Trade},
+            },
             eth,
         },
         infra::{blockchain::Ethereum, observe, solver::ManageNativeToken, Simulator},
@@ -254,30 +259,46 @@ impl Settlement {
 
     /// The settled user orders with their in/out amounts.
     pub fn orders(&self) -> HashMap<order::Uid, competition::Amounts> {
+        let log_err = |trade: &Trade, err: error::Math, kind: &str| -> eth::TokenAmount {
+            // This should never happen, returning 0 is better than panicking, but we
+            // should still alert.
+            let msg = format!("could not compute {kind}");
+            tracing::error!(?trade, prices=?self.solution.prices, ?err, msg);
+            0.into()
+        };
         let mut acc: HashMap<order::Uid, competition::Amounts> = HashMap::new();
-        for trade in self.solution.user_trades() {
-            let prices = ClearingPrices {
-                sell: self.solution.prices[&trade.order().sell.token.wrap(self.solution.weth)],
-                buy: self.solution.prices[&trade.order().buy.token.wrap(self.solution.weth)],
-            };
-            let order = competition::Amounts {
-                side: trade.order().side,
-                sell: trade.order().sell,
-                buy: trade.order().buy,
-                executed_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=?self.solution.prices, ?err, "could not compute sell_amount");
-                    0.into()
-                }),
-                executed_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=?self.solution.prices, ?err, "could not compute buy_amount");
-                    0.into()
-                }),
+        for trade in &self.solution.trades {
+            let order = match trade {
+                Trade::Fulfillment(_) => {
+                    let prices = ClearingPrices {
+                        sell: self.solution.prices[&trade.sell().token.wrap(self.solution.weth)],
+                        buy: self.solution.prices[&trade.buy().token.wrap(self.solution.weth)],
+                    };
+                    competition::Amounts {
+                        side: trade.side(),
+                        sell: trade.sell(),
+                        buy: trade.buy(),
+                        executed_sell: trade
+                            .sell_amount(&prices)
+                            .unwrap_or_else(|err| log_err(trade, err, "executed_sell")),
+                        executed_buy: trade
+                            .buy_amount(&prices)
+                            .unwrap_or_else(|err| log_err(trade, err, "executed_buy")),
+                    }
+                }
+                Trade::Jit(jit) => competition::Amounts {
+                    side: trade.side(),
+                    sell: trade.sell(),
+                    buy: trade.buy(),
+                    executed_sell: jit
+                        .executed_sell()
+                        .unwrap_or_else(|err| log_err(trade, err, "executed_sell")),
+                    executed_buy: jit
+                        .executed_buy()
+                        .unwrap_or_else(|err| log_err(trade, err, "executed_buy")),
+                },
             };
-            acc.insert(trade.order().uid, order);
+            acc.insert(trade.uid(), order);
         }
         acc
     }

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

@@ -2,7 +2,7 @@ use crate::{
     domain::{
         competition::{
             self,
-            order::{self, FeePolicy, SellAmount, Side, TargetAmount},
+            order::{self, FeePolicy, SellAmount, Side, TargetAmount, Uid},
             solution::error::{self, Math},
         },
         eth::{self, Asset},
@@ -19,6 +19,13 @@ pub enum Trade {
 }
 
 impl Trade {
+    pub fn uid(&self) -> Uid {
+        match self {
+            Trade::Fulfillment(fulfillment) => fulfillment.order().uid,
+            Trade::Jit(jit) => jit.order().uid,
+        }
+    }
+
     pub fn side(&self) -> Side {
         match self {
             Trade::Fulfillment(fulfillment) => fulfillment.order().side,
@@ -62,7 +69,10 @@ impl Trade {
     }
 
     /// The effective amount that left the user's wallet including all fees.
-    fn sell_amount(&self, prices: &ClearingPrices) -> Result<eth::TokenAmount, error::Math> {
+    pub(crate) fn sell_amount(
+        &self,
+        prices: &ClearingPrices,
+    ) -> Result<eth::TokenAmount, error::Math> {
         let before_fee = match self.side() {
             order::Side::Sell => self.executed().0,
             order::Side::Buy => self
@@ -81,7 +91,10 @@ impl Trade {
     /// The effective amount the user received after all fees.
     ///
     /// Settlement contract uses `ceil` division for buy amount calculation.
-    fn buy_amount(&self, prices: &ClearingPrices) -> Result<eth::TokenAmount, error::Math> {
+    pub(crate) fn buy_amount(
+        &self,
+        prices: &ClearingPrices,
+    ) -> Result<eth::TokenAmount, error::Math> {
         let amount = match self.side() {
             order::Side::Buy => self.executed().0,
             order::Side::Sell => self
@@ -409,6 +422,41 @@ impl Jit {
         self.executed
     }
 
+    pub fn executed_buy(&self) -> Result<eth::TokenAmount, Math> {
+        Ok(match self.order().side {
+            Side::Buy => self.executed().into(),
+            Side::Sell => (self
+                .executed()
+                .0
+                .checked_add(self.fee().0)
+                .ok_or(Math::Overflow)?)
+            .checked_mul(self.order.buy.amount.0)
+            .ok_or(Math::Overflow)?
+            .checked_ceil_div(&self.order.sell.amount.0)
+            .ok_or(Math::DivisionByZero)?
+            .into(),
+        })
+    }
+
+    pub fn executed_sell(&self) -> Result<eth::TokenAmount, Math> {
+        Ok(match self.order().side {
+            Side::Buy => self
+                .executed()
+                .0
+                .checked_mul(self.order.sell.amount.0)
+                .ok_or(Math::Overflow)?
+                .checked_div(self.order.buy.amount.0)
+                .ok_or(Math::DivisionByZero)?
+                .into(),
+            Side::Sell => self
+                .executed()
+                .0
+                .checked_add(self.fee().0)
+                .ok_or(Math::Overflow)?
+                .into(),
+        })
+    }
+
     pub fn fee(&self) -> order::SellAmount {
         self.fee
     }

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

@@ -13,6 +13,7 @@ use crate::{
             ab_order,
             ab_solution,
             test_solver,
+            ExpectedOrderAmounts,
             Test,
         },
     },
@@ -54,14 +55,26 @@ struct TestCase {
 
 #[cfg(test)]
 async fn protocol_fee_test_case(test_case: TestCase) {
-    let test_name = format!("JIT Order: {:?}", test_case.order.side);
+    let test_name = format!("JIT Order: {:?}", test_case.solution.jit_order.order.side);
     // Adjust liquidity pools so that the order is executable at the amounts
     // expected from the solver.
     let quote = ab_liquidity_quote()
         .sell_amount(test_case.execution.solver.sell)
         .buy_amount(test_case.execution.solver.buy);
     let pool = ab_adjusted_pool(quote);
     let solver_fee = test_case.execution.driver.sell / 100;
+    // Amounts expected to be returned by the driver after fee processing
+    let jit_order_expected_amounts = if test_case.is_surplus_capturing_jit_order {
+        ExpectedOrderAmounts {
+            sell: test_case.execution.solver.sell,
+            buy: test_case.execution.solver.buy,
+        }
+    } else {
+        ExpectedOrderAmounts {
+            sell: test_case.solution.jit_order.order.sell_amount,
+            buy: test_case.solution.jit_order.order.buy_amount,
+        }
+    };
 
     let jit_order = setup::JitOrder {
         order: ab_order()
@@ -70,7 +83,8 @@ async fn protocol_fee_test_case(test_case: TestCase) {
             .buy_amount(test_case.solution.jit_order.order.buy_amount)
             .solver_fee(Some(solver_fee))
             .side(test_case.solution.jit_order.order.side)
-            .no_surplus(),
+            .no_surplus()
+            .expected_amounts(jit_order_expected_amounts),
     };
 
     let order = ab_order()
@@ -83,6 +97,7 @@ async fn protocol_fee_test_case(test_case: TestCase) {
         .no_surplus();
 
     let solver = test_solver();
+
     let test: Test = tests::setup()
         .name(test_name)
         .pool(pool)
@@ -104,6 +119,7 @@ async fn protocol_fee_test_case(test_case: TestCase) {
         result.score(),
         test_case.solution.expected_score,
     ));
+    result.jit_orders(&[jit_order]);
 }
 
 #[tokio::test]

crates/driver/src/tests/setup/mod.rs

@@ -27,6 +27,7 @@ use {
                 DEFAULT_SURPLUS_FACTOR,
                 ETH_ORDER_AMOUNT,
             },
+            hex_address,
             setup::blockchain::{Blockchain, Trade},
         },
     },
@@ -1224,6 +1225,44 @@ impl<'a> SolveOk<'a> {
         assert!(self.solutions().is_empty());
     }
 
+    /// Check that the solution contains the expected JIT orders.
+    pub fn jit_orders(self, jit_orders: &[JitOrder]) -> Self {
+        let solution = self.solution();
+        assert!(solution.get("orders").is_some());
+        let trades = serde_json::from_value::<HashMap<String, serde_json::Value>>(
+            solution.get("orders").unwrap().clone(),
+        )
+        .unwrap();
+
+        // Since JIT orders don't have UID at creation time, we need to search for
+        // matching token pair
+        for expected in jit_orders.iter() {
+            let exist = trades
+                .values()
+                .any(|trade| self.trade_matches(trade, expected));
+            assert!(exist, "JIT order {expected:?} not found");
+        }
+        self
+    }
+
+    /// Find for a JIT order, given specific token pair and buy/sell amount,
+    /// return true if the JIT order was found
+    fn trade_matches(&self, trade: &serde_json::Value, expected: &JitOrder) -> bool {
+        let u256 =
+            |value: &serde_json::Value| eth::U256::from_dec_str(value.as_str().unwrap()).unwrap();
+        let sell_token = trade.get("sellToken").unwrap().to_string();
+        let sell_token = sell_token.trim_matches('"');
+        let buy_token = trade.get("buyToken").unwrap().to_string();
+        let buy_token = buy_token.trim_matches('"');
+        let sell_amount = u256(trade.get("executedSell").unwrap());
+        let buy_amount = u256(trade.get("executedBuy").unwrap());
+
+        sell_token == hex_address(self.blockchain.get_token(expected.order.sell_token))
+            && buy_token == hex_address(self.blockchain.get_token(expected.order.buy_token))
+            && expected.order.expected_amounts.clone().unwrap().sell == sell_amount
+            && expected.order.expected_amounts.clone().unwrap().buy == buy_amount
+    }
+
     /// Check that the solution contains the expected orders.
     pub fn orders(self, orders: &[Order]) -> Self {
         let solution = self.solution();

crates/driver/src/tests/setup/solver.rs

@@ -197,20 +197,22 @@ impl Solver {
                                 })
                             },
                         ));
-                        prices_json.insert(
+                        let previous_value = prices_json.insert(
                             config
                                 .blockchain
                                 .get_token_wrapped(fulfillment.quoted_order.order.sell_token),
                             fulfillment.execution.buy.to_string(),
                         );
-                        prices_json.insert(
+                        assert_eq!(previous_value, None, "existing price overwritten");
+                        let previous_value = prices_json.insert(
                             config
                                 .blockchain
                                 .get_token_wrapped(fulfillment.quoted_order.order.buy_token),
                             (fulfillment.execution.sell
                                 - fulfillment.quoted_order.order.surplus_fee())
                             .to_string(),
                         );
+                        assert_eq!(previous_value, None, "existing price overwritten");
                         {
                             // trades have optional field `fee`
                             let order = if config.quote {
@@ -270,18 +272,27 @@ impl Solver {
                                 }).collect_vec(),
                             })
                         }));
-                        prices_json.insert(
-                            config
-                                .blockchain
-                                .get_token_wrapped(jit.quoted_order.order.sell_token),
-                            jit.execution.buy.to_string(),
-                        );
-                        prices_json.insert(
-                            config
-                                .blockchain
-                                .get_token_wrapped(jit.quoted_order.order.buy_token),
-                            (jit.execution.sell - jit.quoted_order.order.surplus_fee()).to_string(),
-                        );
+                        // Skipping the prices for JIT orders (non-surplus-capturing)
+                        if config
+                            .expected_surplus_capturing_jit_order_owners
+                            .contains(&jit.quoted_order.order.owner)
+                        {
+                            let previous_value = prices_json.insert(
+                                config
+                                    .blockchain
+                                    .get_token_wrapped(jit.quoted_order.order.sell_token),
+                                jit.execution.buy.to_string(),
+                            );
+                            assert_eq!(previous_value, None, "existing price overwritten");
+                            let previous_value = prices_json.insert(
+                                config
+                                    .blockchain
+                                    .get_token_wrapped(jit.quoted_order.order.buy_token),
+                                (jit.execution.sell - jit.quoted_order.order.surplus_fee())
+                                    .to_string(),
+                            );
+                            assert_eq!(previous_value, None, "existing price overwritten");
+                        }
                         {
                             let executed_amount = match jit.quoted_order.order.executed {
                                 Some(executed) => executed.to_string(),