feat(order): withdraw imbalance

lib/amm_dex_v2/math.ak

@@ -173,7 +173,7 @@ pub fn calculate_amount_in(
 //      + swap_amount * ( 2 - f ) * reserve_in * ( amount_out + reserve_out )
 //      + ( reserve_in ^ 2 * amount_out + reserve_in * reserve_out * amount_out )
 //    = 0
-pub fn calculate_swap_amount(
+pub fn calculate_deposit_swap_amount(
   amount_in: Int,
   amount_out: Int,
   reserve_in: Int,
@@ -212,7 +212,7 @@ pub fn calculate_deposit_amount(
   if ratio_a > ratio_b {
     // swap a part of A to B
     let (swap_amount_a_numerator, swap_amount_a_denominator) =
-      calculate_swap_amount(
+      calculate_deposit_swap_amount(
         amount_in: amount_a,
         amount_out: amount_b,
         reserve_in: reserve_a,
@@ -250,7 +250,7 @@ pub fn calculate_deposit_amount(
   } else if ratio_a < ratio_b {
     // swap a part of B to A
     let (swap_amount_b_numerator, swap_amount_b_denominator) =
-      calculate_swap_amount(
+      calculate_deposit_swap_amount(
         amount_in: amount_b,
         amount_out: amount_a,
         reserve_in: reserve_b,
@@ -366,3 +366,205 @@ pub fn calculate_zap_out(
     (0, earned_fee_b, withdrawal_asset_a_amount + extra_amount_out)
   }
 }
+
+// Base formula:
+// ( amount_in - swap_amount_in) / (amount_out + receive_amount_out) = expect_io_ratio_numerator / expect_io_ratio_denominator (1)
+// receive_amount_out = fee_diff * swap_amount_in * reserve_out / (fee_denominator) * reserve_in + fee_diff * swap_amount_in (2)
+// (1) & (2) =>>>>>
+// a * swap_amount_in ^ 2 + b * swap_amount_in + c = 0
+// Where:
+// - a = fee_diff * expect_io_ratio_denominator
+// - b = expect_io_ratio_numerator * fee_diff ( reserve_out + amount_out ) * expect_io_ratio_denominator * (fee_denominator * reserve_in - fee_diff * amount_in)
+// - c = fee_denominator * (expect_io_ratio_numerator * amount_out - expect_io_ratio_denominator * amount_in) 
+fn calculate_withdraw_swap_amount(
+  amount_in: Int,
+  amount_out: Int,
+  reserve_in: Int,
+  reserve_out: Int,
+  expect_io_ratio_numerator: Int,
+  expect_io_ratio_denominator: Int,
+  trading_fee_numerator: Int,
+  trading_fee_denominator: Int,
+) -> Int {
+  let diff = trading_fee_denominator - trading_fee_numerator
+  let a = expect_io_ratio_denominator * diff
+  let b =
+    expect_io_ratio_numerator * diff * ( reserve_out + amount_out ) + expect_io_ratio_denominator * (
+      reserve_in * trading_fee_denominator - diff * amount_in
+    )
+  let c =
+    trading_fee_denominator * reserve_in * (
+      expect_io_ratio_numerator * amount_out - expect_io_ratio_denominator * amount_in
+    )
+  // a*x^2+b*x+c=0
+  // delta = b^2 - 4ac
+  // x = (-b +sqrt(delta))/(2*a) or (-b - sqrt(delta))/(2*a)
+  let delta = b * b - 4 * a * c
+  ( calculate_sqrt(delta) - b ) / ( 2 * a )
+}
+
+// return earning fee(a,b)+ amount out(a,b)
+pub fn calculate_withdraw_imbalance(
+  expect_ab_ratio_numerator: Int,
+  expect_ab_ratio_denominator: Int,
+  reserve_a: Int,
+  reserve_b: Int,
+  trading_fee_numerator: Int,
+  trading_fee_denominator: Int,
+  withdrawal_lp_amount: Int,
+  total_liquidity: Int,
+  profit_sharing_opt: Option<(Int, Int)>,
+) -> (Int, Int, Int, Int) {
+  let (withdrawal_asset_a_amount, withdrawal_asset_b_amount) =
+    calculate_withdraw(
+      reserve_a: reserve_a,
+      reserve_b: reserve_b,
+      withdrawal_lp_amount: withdrawal_lp_amount,
+      total_liquidity: total_liquidity,
+    )
+  let reserve_a_after_withdraw = reserve_a - withdrawal_asset_a_amount
+  let reserve_b_after_withdraw = reserve_b - withdrawal_asset_b_amount
+  let ratio_a = withdrawal_asset_a_amount * expect_ab_ratio_denominator
+  let ratio_b = withdrawal_asset_b_amount * expect_ab_ratio_numerator
+  if ratio_a > ratio_b {
+    // swap a-> b
+    let swap_amount_a =
+      calculate_withdraw_swap_amount(
+        amount_in: withdrawal_asset_a_amount,
+        amount_out: withdrawal_asset_b_amount,
+        reserve_in: reserve_a_after_withdraw,
+        reserve_out: reserve_b_after_withdraw,
+        expect_io_ratio_numerator: expect_ab_ratio_numerator,
+        expect_io_ratio_denominator: expect_ab_ratio_denominator,
+        trading_fee_numerator: trading_fee_numerator,
+        trading_fee_denominator: trading_fee_denominator,
+      )
+
+    let receive_amount_b =
+      calculate_amount_out(
+        reserve_in: reserve_a_after_withdraw,
+        reserve_out: reserve_b_after_withdraw,
+        amount_in: swap_amount_a,
+        trading_fee_numerator: trading_fee_numerator,
+        trading_fee_denominator: trading_fee_denominator,
+      )
+    let earned_fee_a =
+      when profit_sharing_opt is {
+        Some((fee_sharing_numerator, fee_sharing_denominator)) ->
+          calculate_earned_fee_in(
+            amount_in: swap_amount_a,
+            trading_fee_numerator: trading_fee_numerator,
+            trading_fee_denominator: trading_fee_denominator,
+            fee_sharing_numerator: fee_sharing_numerator,
+            fee_sharing_denominator: fee_sharing_denominator,
+          )
+        None -> 0
+      }
+    (
+      earned_fee_a,
+      0,
+      withdrawal_asset_a_amount - swap_amount_a,
+      withdrawal_asset_b_amount + receive_amount_b,
+    )
+  } else if ratio_a < ratio_b {
+    // swap b-> a
+    let swap_amount_b =
+      calculate_withdraw_swap_amount(
+        amount_in: withdrawal_asset_b_amount,
+        amount_out: withdrawal_asset_a_amount,
+        reserve_in: reserve_b_after_withdraw,
+        reserve_out: reserve_a_after_withdraw,
+        expect_io_ratio_numerator: expect_ab_ratio_denominator,
+        expect_io_ratio_denominator: expect_ab_ratio_numerator,
+        trading_fee_numerator: trading_fee_numerator,
+        trading_fee_denominator: trading_fee_denominator,
+      )
+
+    let receive_amount_a =
+      calculate_amount_out(
+        reserve_in: reserve_b_after_withdraw,
+        reserve_out: reserve_a_after_withdraw,
+        amount_in: swap_amount_b,
+        trading_fee_numerator: trading_fee_numerator,
+        trading_fee_denominator: trading_fee_denominator,
+      )
+    let earned_fee_b =
+      when profit_sharing_opt is {
+        Some((fee_sharing_numerator, fee_sharing_denominator)) ->
+          calculate_earned_fee_in(
+            amount_in: swap_amount_b,
+            trading_fee_numerator: trading_fee_numerator,
+            trading_fee_denominator: trading_fee_denominator,
+            fee_sharing_numerator: fee_sharing_numerator,
+            fee_sharing_denominator: fee_sharing_denominator,
+          )
+        None -> 0
+      }
+    (
+      0,
+      earned_fee_b,
+      withdrawal_asset_a_amount + receive_amount_a,
+      withdrawal_asset_b_amount - swap_amount_b,
+    )
+  } else {
+    (0, 0, withdrawal_asset_a_amount, withdrawal_asset_b_amount)
+  }
+}
+
+test test_calculate_withdraw_swap_amount() {
+  let reserve_a = 1_000_000_000
+  let reserve_b = 1_000_000_000
+  let total_liquidity = 1_000_000_000
+  let trading_fee_numerator = 3
+  let trading_fee_denominator = 1000
+  let withdrawal_lp_amount = 1000
+  let asset_a_ratio = 3
+  let asset_b_ratio = 1
+
+  let (withdrawal_asset_a_amount, withdrawal_asset_b_amount) =
+    calculate_withdraw(
+      reserve_a: reserve_a,
+      reserve_b: reserve_b,
+      withdrawal_lp_amount: withdrawal_lp_amount,
+      total_liquidity: total_liquidity,
+    )
+  let reserve_a_after_withdraw = reserve_a - withdrawal_asset_a_amount
+  let reserve_b_after_withdraw = reserve_b - withdrawal_asset_b_amount
+  // b -> a
+  let swap_amount_b =
+    calculate_withdraw_swap_amount(
+      amount_in: withdrawal_asset_b_amount,
+      amount_out: withdrawal_asset_a_amount,
+      reserve_in: reserve_b_after_withdraw,
+      reserve_out: reserve_a_after_withdraw,
+      expect_io_ratio_numerator: asset_b_ratio,
+      expect_io_ratio_denominator: asset_a_ratio,
+      trading_fee_numerator: trading_fee_numerator,
+      trading_fee_denominator: trading_fee_denominator,
+    )
+  swap_amount_b == 500 && withdrawal_asset_a_amount == 1000 && withdrawal_asset_a_amount == 1000
+}
+
+test test_calculate_withdraw_imbalance() {
+  let reserve_a = 1_000_000_000
+  let reserve_b = 1_000_000_000
+  let total_liquidity = 1_000_000_000
+  let trading_fee_numerator = 3
+  let trading_fee_denominator = 1000
+  let withdrawal_lp_amount = 1000
+  let asset_a_ratio = 3
+  let asset_b_ratio = 1
+  let profit_sharing_opt = None
+
+  calculate_withdraw_imbalance(
+    expect_ab_ratio_numerator: asset_a_ratio,
+    expect_ab_ratio_denominator: asset_b_ratio,
+    reserve_a: reserve_a,
+    reserve_b: reserve_b,
+    trading_fee_numerator: trading_fee_numerator,
+    trading_fee_denominator: trading_fee_denominator,
+    withdrawal_lp_amount: withdrawal_lp_amount,
+    total_liquidity: total_liquidity,
+    profit_sharing_opt: profit_sharing_opt,
+  ) == (0, 0, 1498, 500)
+}

lib/amm_dex_v2/order_validation.ak

@@ -8,7 +8,7 @@ use amm_dex_v2/math
 use amm_dex_v2/types.{
   AToB, Asset, BToA, BatchingPool, CustomDatumHash, DatumMap, Deposit, OCO,
   OrderDatum, PartialSwap, PoolState, StopLoss, SwapExactIn, SwapExactOut,
-  SwapMultiRouting, SwapRouting, Withdraw, ZapOut,
+  SwapMultiRouting, SwapRouting, Withdraw, WithdrawImbalance, ZapOut,
 }
 use amm_dex_v2/utils
 
@@ -544,6 +544,71 @@ fn validate_partial_swap(
   }
 }
 
+fn validate_withdraw_imbalance(
+  order_in_value: Value,
+  order_out_value: Value,
+  ratio_asset_a: Int,
+  ratio_asset_b: Int,
+  minimum_amount_a: Int,
+  asset_a: Asset,
+  asset_b: Asset,
+  lp_asset: Asset,
+  output_ada: Int,
+  trading_fee_numerator: Int,
+  trading_fee_denominator: Int,
+  profit_sharing_opt: Option<(Int, Int)>,
+  pool_state: PoolState,
+) -> PoolState {
+  let (
+    datum_reserve_a,
+    datum_reserve_b,
+    value_reserve_a,
+    value_reserve_b,
+    total_liquidity,
+  ) = pool_state
+  let Asset { policy_id: asset_a_policy_id, asset_name: asset_a_asset_name } =
+    asset_a
+  let Asset { policy_id: asset_b_policy_id, asset_name: asset_b_asset_name } =
+    asset_b
+  let Asset { policy_id: lp_policy_id, asset_name: lp_asset_name } = lp_asset
+  let withdrawal_lp_amount =
+    value.quantity_of(order_in_value, lp_policy_id, lp_asset_name)
+  let has_enough_amount_in = withdrawal_lp_amount > 0
+  let (earned_fee_a, earned_fee_b, amount_a_out, amount_b_out) =
+    math.calculate_withdraw_imbalance(
+      expect_ab_ratio_numerator: ratio_asset_a,
+      expect_ab_ratio_denominator: ratio_asset_b,
+      reserve_a: datum_reserve_a,
+      reserve_b: datum_reserve_b,
+      trading_fee_numerator: trading_fee_numerator,
+      trading_fee_denominator: trading_fee_denominator,
+      withdrawal_lp_amount: withdrawal_lp_amount,
+      total_liquidity: total_liquidity,
+      profit_sharing_opt: profit_sharing_opt,
+    )
+  let not_over_slippage = amount_a_out >= minimum_amount_a
+  let actual_amount_a_out =
+    value.quantity_of(order_out_value, asset_a_policy_id, asset_a_asset_name)
+  let actual_amount_b_out =
+    value.quantity_of(order_out_value, asset_b_policy_id, asset_b_asset_name)
+  let ada_amount_out =
+    value.quantity_of(order_out_value, ada_policy_id, ada_asset_name)
+  let is_valid_amount_out =
+    if utils.is_ada_asset(asset_a_policy_id, asset_a_asset_name) {
+      actual_amount_a_out == amount_a_out + output_ada && actual_amount_b_out == amount_b_out
+    } else {
+      ada_amount_out == output_ada && actual_amount_a_out == amount_a_out && actual_amount_b_out == amount_b_out
+    }
+  expect has_enough_amount_in && not_over_slippage && is_valid_amount_out
+  (
+    datum_reserve_a - amount_a_out - earned_fee_a,
+    datum_reserve_b - amount_b_out - earned_fee_b,
+    value_reserve_a - amount_a_out,
+    value_reserve_b - amount_b_out,
+    total_liquidity - withdrawal_lp_amount,
+  )
+}
+
 fn get_multi_routing_swap_amount_out(
   amount_in: Int,
   pool: BatchingPool,
@@ -1071,6 +1136,31 @@ pub fn apply_orders(
                   new_root_k_last,
                 )
               }
+              WithdrawImbalance(ratio_asset_a, ratio_asset_b, minimum_asset_a) -> {
+                expect
+                  ratio_asset_a > 0 && ratio_asset_b > 0 && minimum_asset_a > 0
+                expect
+                  validate_order_receiver(
+                    receiver: receiver,
+                    receiver_datum_hash_opt: receiver_datum_hash_opt,
+                    output: output,
+                  )
+                validate_withdraw_imbalance(
+                  order_in_value: order_in_value,
+                  order_out_value: order_out_value,
+                  ratio_asset_a: ratio_asset_a,
+                  ratio_asset_b: ratio_asset_b,
+                  minimum_amount_a: minimum_asset_a,
+                  asset_a: asset_a,
+                  asset_b: asset_b,
+                  lp_asset: lp_asset,
+                  output_ada: output_ada,
+                  trading_fee_numerator: trading_fee_numerator,
+                  trading_fee_denominator: trading_fee_denominator,
+                  profit_sharing_opt: profit_sharing_opt,
+                  pool_state: pool_state,
+                )
+              }
               SwapMultiRouting(_, _) -> fail
             }
           apply_orders(

lib/amm_dex_v2/types.ak

@@ -87,6 +87,11 @@ pub type OrderStep {
     hops: Int,
     minimum_swap_amount_required: Int,
   }
+  WithdrawImbalance {
+    ratio_asset_a: Int,
+    ratio_asset_b: Int,
+    minimum_asset_a: Int,
+  }
   SwapMultiRouting { routings: List<SwapRouting>, minimum_receive: Int }
 }
 

validators/pool_validator.ak

@@ -1,3 +1,4 @@
+use aiken/builtin
 use aiken/list
 use aiken/transaction.{
   Input, Output, OutputReference, ScriptContext, Spend, Transaction,
@@ -29,7 +30,7 @@ validator(
     expect Spend(pool_input_ref) = purpose
     when redeemer is {
       Batching(batcher_address, input_indexes, license_index) -> {
-        expect license_index >= 0 && list.length(input_indexes) > 0
+        expect license_index >= 0 && !builtin.null_list(input_indexes)
         let Transaction {
           inputs,
           outputs,