Refactor to use continuations

lib/calculation/deposit.ak

@@ -1,7 +1,7 @@
 use aiken/math
 use aiken/transaction.{NoDatum, Output, InlineDatum}
 use aiken/transaction/credential.{Address, VerificationKeyCredential}
-use aiken/transaction/value.{ada_policy_id, ada_asset_name}
+use aiken/transaction/value.{ada_policy_id, ada_asset_name, PolicyId, AssetName}
 use calculation/shared.{PoolState} as calc_shared
 use sundae/multisig
 use shared.{SingletonValue}
@@ -13,22 +13,31 @@ use types/order.{Destination, Fixed, Self, OrderDatum}
 /// LP tokens are distributed to the destination.
 ///
 pub fn do_deposit(
-  pool_state: PoolState,
+  pool_policy_a: PolicyId,
+  pool_asset_name_a: AssetName,
+  pool_quantity_a: Int,
+  pool_policy_b: PolicyId,
+  pool_asset_name_b: AssetName,
+  pool_quantity_b: Int,
+  pool_policy_lp: PolicyId,
+  pool_asset_name_lp: AssetName,
+  pool_quantity_lp: Int,
   input_utxo: Output,
   assets: (SingletonValue, SingletonValue),
   destination: Destination,
   actual_protocol_fee: Int,
   output: Output,
-) -> PoolState {
+  continuation: fn(Int, Int, Int) -> Bool,
+) -> Bool {
   let (asset_a, asset_b) = assets
   let Output { value: input_value, .. } = input_utxo
 
   // Policy ID and token name of the assets must match the pool, otherwise someone
   // could load the pool with junk tokens and freeze the pool.
-  expect asset_a.1st == pool_state.quantity_a.1st
-  expect asset_a.2nd == pool_state.quantity_a.2nd
-  expect asset_b.1st == pool_state.quantity_b.1st
-  expect asset_b.2nd == pool_state.quantity_b.2nd
+  expect asset_a.1st == pool_policy_a
+  expect asset_a.2nd == pool_asset_name_a
+  expect asset_b.1st == pool_policy_b
+  expect asset_b.2nd == pool_asset_name_b
 
   // A deposit is permitted to have lower funds than the datum claims, so that
   // we can compose it as a step in a chain where the exact amount will be unknown.
@@ -65,19 +74,17 @@ pub fn do_deposit(
   // So some small amount of a or b might be returned to the user
   // So, calculate how much "b" do we have, in units of asset A, so we can check which is greater
   let b_in_units_of_a =
-    user_gives_b * pool_state.quantity_a.3rd / pool_state.quantity_b.3rd
+    user_gives_b * pool_quantity_a / pool_quantity_b
 
   // Amount that user actually deposits in the pool, after giving back change.
   let (deposited_a, deposited_b) =
-    // If we have more b than a, then we can only take up to an quivalent amount of b, and return the rest
+    // If we have more b than a, then we can only take up to an equivalent amount of b, and return the rest
     // otherwise if we have more a than b, we can return some amount of `a` to the user instead
     if b_in_units_of_a > user_gives_a {
-      let num = pool_state.quantity_b.3rd * user_gives_a
-      let denom = pool_state.quantity_a.3rd
       // Make sure to do ceiling division here, to round in favor fo the protocol
       // That is, when in doubt, take up to one more token from the user than the
       // LP tokens we issue would entail
-      let give_b = ((num - 1) / denom) + 1
+      let give_b = ((pool_quantity_b * user_gives_a - 1) / pool_quantity_a) + 1
       (user_gives_a, give_b)
     } else {
       (b_in_units_of_a, user_gives_b)
@@ -92,7 +99,7 @@ pub fn do_deposit(
   // 
   // Solving for `issued_lp_tokens` gives:
   let issued_lp_tokens =
-    deposited_a * pool_state.quantity_lp.3rd / pool_state.quantity_a.3rd
+    deposited_a * pool_quantity_lp / pool_quantity_a
 
   // Make sure we don't ever allow this to round to zero, which would just eat some of the users assets
   expect issued_lp_tokens > 0
@@ -106,8 +113,8 @@ pub fn do_deposit(
       |> value.add(asset_b.1st, asset_b.2nd, -deposited_b)
       |> value.add(ada_policy_id, ada_asset_name, -actual_protocol_fee)
       |> value.add(
-          pool_state.quantity_lp.1st,
-          pool_state.quantity_lp.2nd,
+          pool_policy_lp,
+          pool_asset_name_lp,
           issued_lp_tokens,
         )
 
@@ -132,23 +139,11 @@ pub fn do_deposit(
   }
 
   // And construct the final pool state
-  PoolState {
-    quantity_a: (
-      pool_state.quantity_a.1st,
-      pool_state.quantity_a.2nd,
-      pool_state.quantity_a.3rd + deposited_a,
-    ),
-    quantity_b: (
-      pool_state.quantity_b.1st,
-      pool_state.quantity_b.2nd,
-      pool_state.quantity_b.3rd + deposited_b,
-    ),
-    quantity_lp: (
-      pool_state.quantity_lp.1st,
-      pool_state.quantity_lp.2nd,
-      pool_state.quantity_lp.3rd + issued_lp_tokens,
-    ),
-  }
+  continuation(
+    pool_quantity_a + deposited_a,
+    pool_quantity_b + deposited_b,
+    pool_quantity_lp + issued_lp_tokens,
+  )
 }
 
 test deposit_test() {
@@ -199,8 +194,16 @@ test deposit_test() {
     datum: InlineDatum(order),
     reference_script: None,
   }
-  let final_pool_state = do_deposit(pool_state, input, assets, order.destination, 2_500_000, output)
-  expect final_pool_state.quantity_a.3rd == 1_010_000_000
-  expect final_pool_state.quantity_b.3rd == 1_010_000_000
+  let new_a, new_b, new_lp <- do_deposit(
+    pool_state.quantity_a.1st, pool_state.quantity_a.2nd, pool_state.quantity_a.3rd,
+    pool_state.quantity_b.1st, pool_state.quantity_b.2nd, pool_state.quantity_b.3rd,
+    pool_state.quantity_lp.1st, pool_state.quantity_lp.2nd, pool_state.quantity_lp.3rd,
+    input, assets,
+    order.destination, 2_500_000,
+    output
+  )
+  expect new_a == 1_010_000_000
+  expect new_b == 1_010_000_000
+  expect new_lp == 1_000_000_000 + 10_000_000
   True
 }

lib/calculation/donation.ak

@@ -1,6 +1,5 @@
 use aiken/transaction.{Output}
-use aiken/transaction/value.{ada_policy_id, ada_asset_name}
-use calculation/shared.{PoolState} as calc_shared
+use aiken/transaction/value.{ada_policy_id, ada_asset_name, PolicyId, AssetName}
 use shared.{SingletonValue}
 use types/order.{Destination, Fixed, Self}
 
@@ -13,7 +12,12 @@ use types/order.{Destination, Fixed, Self}
 // Calculates the new pool state, and whether to consume this output or not
 pub fn do_donation(
   /// The pool state as a result of processing all orders before this one
-  pool_state: PoolState,
+  pool_policy_a: PolicyId,
+  pool_asset_name_a: AssetName,
+  pool_quantity_a: Int,
+  pool_policy_b: PolicyId,
+  pool_asset_name_b: AssetName,
+  pool_quantity_b: Int,
   /// The total quantity of tokens on this particular input
   input_utxo: Output,
   /// The amounts of each asset being donated
@@ -26,15 +30,17 @@ pub fn do_donation(
   /// The next output in the list; If there is any change left over from the donation, we would expect this to be that change
   /// If there is no change leftover, this output is ignored and used for the next order
   output: Output,
-) -> (PoolState, Bool) {
+  // A continuation to call with the new pool balances, and whether to skip an output; more efficient than constructing an object each time
+  continuation: fn(Int, Int, Bool) -> Bool
+) -> Bool {
   let Output { value: input_value, .. } = input_utxo
   let ((asset_a_policy_id, asset_a_asset_name, asset_a_qty), (asset_b_policy_id, asset_b_asset_name, asset_b_qty)) = assets
   // Make sure we're actually donating the pool assets; this is to prevent setting
   // poolIdent to None, and then filling the pool UTXO with garbage tokens and eventually locking it
-  expect asset_a_policy_id == pool_state.quantity_a.1st
-  expect asset_a_asset_name == pool_state.quantity_a.2nd
-  expect asset_b_policy_id == pool_state.quantity_b.1st
-  expect asset_b_asset_name == pool_state.quantity_b.2nd
+  expect asset_a_policy_id == pool_policy_a
+  expect asset_a_asset_name == pool_asset_name_a
+  expect asset_b_policy_id == pool_policy_b
+  expect asset_b_asset_name == pool_asset_name_b
   // Compute however much of the UTXO value is *left over* after deducting the donation amount from it; If nonzero, this will need to be returned to the user
   let remainder =
     input_value
@@ -70,24 +76,13 @@ pub fn do_donation(
     }
   }
 
-  // Compute the new pool state, which is exactly the old pool state, plus the 
+  // Continue with the new pool state, which is exactly the old pool state, plus the 
   // quantities of assets donated, plus the protocol fee!
   // We also returned whether we had a remainder or not, so the calling function knows
   // whether to consume the output they gave us or not.
-  (
-    PoolState {
-      quantity_a: (
-        pool_state.quantity_a.1st,
-        pool_state.quantity_a.2nd,
-        pool_state.quantity_a.3rd + assets.1st.3rd,
-      ),
-      quantity_b: (
-        pool_state.quantity_b.1st,
-        pool_state.quantity_b.2nd,
-        pool_state.quantity_b.3rd + assets.2nd.3rd,
-      ),
-      quantity_lp: pool_state.quantity_lp,
-    },
+  continuation(
+    pool_quantity_a + assets.1st.3rd,
+    pool_quantity_b + assets.2nd.3rd,
     has_remainder,
   )
 }