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,49 @@ 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()
-                }),
-            };
-            acc.insert(trade.order().uid, order);
+        for trade in &self.solution.trades {
+            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)],
+                    };
+                    let order = 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, "sell_amount")),
+                        executed_buy: trade
+                            .buy_amount(&prices)
+                            .unwrap_or_else(|err| log_err(trade, err, "buy_amount")),
+                    };
+                    acc.insert(trade.uid(), order);
+                }
+                Trade::Jit(jit) => {
+                    let order = competition::Amounts {
+                        side: trade.side(),
+                        sell: trade.sell(),
+                        buy: trade.buy(),
+                        executed_sell: jit
+                            .executed_sell()
+                            .unwrap_or_else(|err| log_err(trade, err, "sell_amount")),
+                        executed_buy: jit
+                            .executed_buy()
+                            .unwrap_or_else(|err| log_err(trade, err, "buy_amount")),
+                    };
+                    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

@@ -54,7 +54,7 @@ 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()
@@ -83,6 +83,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 +105,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

@@ -1224,6 +1224,41 @@ 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() {
+            for trade in trades.values() {
+                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 buy_token = trade.get("buyToken").unwrap().to_string();
+
+                if sell_token == expected.order.sell_token && buy_token == expected.order.buy_token
+                {
+                    assert!(is_approximately_equal(
+                        expected.order.sell_amount,
+                        u256(trade.get("executedSell").unwrap())
+                    ));
+                    assert!(is_approximately_equal(
+                        expected.order.buy_amount.unwrap(),
+                        u256(trade.get("executedBuy").unwrap())
+                    ));
+                }
+            }
+        }
+        self
+    }
+
     /// 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

@@ -270,18 +270,25 @@ 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)
+                        {
+                            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(),
+                            );
+                        }
                         {
                             let executed_amount = match jit.quoted_order.order.executed {
                                 Some(executed) => executed.to_string(),