Merge pull request #52 from SundaeSwap-finance/ignaciodopazo/has_expected_pool_value-optimization

Resolve SSW-312 (`has_expected_pool_value` traverses output value just once)

Author: Quantumplation

validators/pool.ak

@@ -10,13 +10,10 @@ use aiken/transaction.{
 use aiken/transaction/credential.{
   Address, Inline, ScriptCredential,
 }
-use aiken/transaction/value.{MintedValue, PolicyId, Value, ada_policy_id}
+use aiken/transaction/value.{MintedValue, PolicyId, Value, ada_policy_id, ada_asset_name}
 use calculation/process.{pool_input_to_state, process_orders}
 use calculation/shared.{PoolState} as calc_shared
-use shared.{
-  AssetClass, Ident, count_orders, has_exact_token_count, pool_lp_name,
-  pool_nft_name, spent_output,
-}
+use shared.{AssetClass, Ident, spent_output, pool_nft_name, pool_lp_name, count_orders}
 use sundae/multisig
 use types/pool.{
   CreatePool, MintLP, PoolDatum, PoolMintRedeemer, PoolRedeemer, PoolScoop,
@@ -519,7 +516,7 @@ fn minted_correct_pool_tokens(
 
 /// Check that the UTXO contents are correct given a specific pool outcome
 /// In particular, it must have the final A reserves, the final B reserves, the pool NFT, and the protocol fees
-fn has_expected_pool_value(
+pub fn has_expected_pool_value(
   pool_script_hash: PolicyId,
   identifier: Ident,
   output_value: Value,
@@ -531,33 +528,48 @@ fn has_expected_pool_value(
   let (quantity_b_policy_id, quantity_b_name, quantity_b_amt) = quantity_b
   // Asset A *could* be ADA; in which case there should be 3 tokens on the output
   // (ADA, Asset B, and the NFT)
-  // We do this as an optimization, since constructing values is expensive
-  // OPTIMIZATION: is it possible to check this with a single traversal? Each quantity_of represents a small linear scan
   if quantity_a_policy_id == ada_policy_id {
-    and {
-      has_exact_token_count(output_value, 3),
-      value.lovelace_of(output_value) == final_protocol_fees + quantity_a_amt,
-      value.quantity_of(output_value, quantity_b_policy_id, quantity_b_name) == quantity_b_amt,
-      value.quantity_of(
-        output_value,
-        pool_script_hash,
-        pool_nft_name(identifier),
-      ) == 1,
-    }
+    // Rather than constructing a value directly (which can be expensive)
+    // we can just compare the expected token count and amounts with a single pass over the value
+    (3, final_protocol_fees + quantity_a_amt, quantity_b_amt, 1) ==
+      list.foldl(
+        value.flatten(output_value),
+        // (token count, lovelace amount, token b amount, pool nft amount)
+        (0, 0, 0, 0),
+        fn (asset, acc) {
+          let token_count = acc.1st + 1
+          if asset.1st == quantity_a_policy_id {
+            (token_count, acc.2nd + asset.3rd, acc.3rd, acc.4th)
+          } else if asset.1st == quantity_b_policy_id && asset.2nd == quantity_b_name {
+            (token_count, acc.2nd, acc.3rd + asset.3rd, acc.4th)
+          } else {
+            expect asset == (pool_script_hash, pool_nft_name(identifier), 1)
+            (token_count, acc.2nd, acc.3rd, acc.4th + 1)
+          }
+        }
+      )
   } else {
-    // Otherwise, we expect 4 tokens (ADA, Asset A, Asset B, and the NFT)
-    // and the corresponding values
-    and {
-      has_exact_token_count(output_value, 4),
-      value.lovelace_of(output_value) == final_protocol_fees,
-      value.quantity_of(output_value, quantity_a_policy_id, quantity_a_name) == quantity_a_amt,
-      value.quantity_of(output_value, quantity_b_policy_id, quantity_b_name) == quantity_b_amt,
-      value.quantity_of(
-        output_value,
-        pool_script_hash,
-        pool_nft_name(identifier),
-      ) == 1,
-    }
+    // Asset A isn't ADA, Asset B will *never* be ADA; in this case, there should be 4 tokens on the output:
+    // ADA, the Pool NFT, Asset A, and Asset B
+    (4, final_protocol_fees, quantity_a_amt, quantity_b_amt, 1) ==
+      list.foldl(
+        value.flatten(output_value),
+        // (token count, lovelace amount, token a amount, token b amount, pool nft amount)
+        (0, 0, 0, 0, 0),
+        fn (asset, acc) {
+          let token_count = acc.1st + 1
+          if asset.1st == ada_policy_id && asset.2nd == ada_asset_name {
+            (token_count, acc.2nd + asset.3rd, acc.3rd, acc.4th, acc.5th)
+          } else if asset.1st == quantity_a_policy_id && asset.2nd == quantity_a_name {
+            (token_count, acc.2nd, acc.3rd + asset.3rd, acc.4th, acc.5th)
+          } else if asset.1st == quantity_b_policy_id && asset.2nd == quantity_b_name {
+            (token_count, acc.2nd, acc.3rd, acc.4th + asset.3rd, acc.5th)
+          } else {
+            expect asset == (pool_script_hash, pool_nft_name(identifier), 1)
+            (token_count, acc.2nd, acc.3rd, acc.4th, acc.5th + 1)
+          }
+        }
+      )
   }
 }
 
@@ -576,4 +588,4 @@ fn compare_asset_class(a: AssetClass, b: AssetClass) {
 pub fn int_to_ident(n: Int) -> Ident {
   expect n < 256
   bytearray.push(#"", n)
-}
+}

validators/tests/pool.ak

@@ -11,7 +11,7 @@ use aiken/transaction.{
 use aiken/transaction/credential.{
   Address, VerificationKeyCredential, from_verification_key, from_script, with_delegation_key
 }
-use aiken/transaction/value.{Value}
+use aiken/transaction/value.{Value, ada_policy_id, ada_asset_name}
 use shared.{
   pool_lp_name,
 }
@@ -24,6 +24,7 @@ use types/order.{Deposit, Destination, OrderDatum, Swap}
 use types/pool.{
   CreatePool, PoolDatum, PoolScoop,
 }
+use calculation/shared.{PoolState} as calc_shared
 use types/settings.{SettingsDatum}
 use tx_util/builder.{add_asset_to_tx_output, add_tx_input, add_tx_output, add_tx_ref_input, build_txn_context, mint_assets, new_tx_input, new_tx_output}
 use pool as pool_validator
@@ -693,4 +694,61 @@ test mint_test_nonvoid_datum() fail {
   identity
   )
   minted
+}
+
+test has_expected_pool_value_test() {
+  // ADA/rberry pool
+  let rberry_policy_id = #"9a9693a9a37912a5097918f97918d15240c92ab729a0b7c4aa144d77"
+  let rberry_token_name = #"524245525259"
+
+  let pool_script_hash = #"00000000000000000000000000000000000000000000000000000000"
+  let identifier = #"00000000000000000000000000000000000000000000000000000000"
+  let (_, pool_nft_token, pool_lp_token) = shared.pool_token_names(identifier)
+  let pool_value = value.from_lovelace(102_000_000)
+    |> value.add(rberry_policy_id, rberry_token_name, 100_000_000)
+    |> value.add(pool_script_hash, pool_nft_token, 1)
+  let outcome = PoolState {
+    quantity_lp: (pool_script_hash, pool_lp_token, 99),
+    quantity_a: (ada_policy_id, ada_asset_name, 100_000_000),
+    quantity_b: (rberry_policy_id, rberry_token_name, 100_000_000)
+  }
+  let protocol_fees = 2_000_000
+
+  pool_validator.has_expected_pool_value(
+    pool_script_hash,
+    identifier,
+    pool_value,
+    outcome,
+    protocol_fees,
+  )
+}
+
+test has_expected_pool_value_test2() {
+  // other/rberry pool
+  let other_policy_id = #"fafafafafafafafafafafafafafafafafafafafafafafafafafafafa"
+  let other_token_name = #"fafafafa"
+  let rberry_policy_id = #"9a9693a9a37912a5097918f97918d15240c92ab729a0b7c4aa144d77"
+  let rberry_token_name = #"524245525259"
+
+  let pool_script_hash = #"00000000000000000000000000000000000000000000000000000000"
+  let identifier = #"00000000000000000000000000000000000000000000000000000000"
+  let (_, pool_nft_token, pool_lp_token) = shared.pool_token_names(identifier)
+  let pool_value = value.from_lovelace(2_000_000)
+    |> value.add(other_policy_id, other_token_name, 100_000_000)
+    |> value.add(rberry_policy_id, rberry_token_name, 100_000_000)
+    |> value.add(pool_script_hash, pool_nft_token, 1)
+  let outcome = PoolState {
+    quantity_lp: (pool_script_hash, pool_lp_token, 99),
+    quantity_a: (other_policy_id, other_token_name, 100_000_000),
+    quantity_b: (rberry_policy_id, rberry_token_name, 100_000_000)
+  }
+  let protocol_fees = 2_000_000
+
+  pool_validator.has_expected_pool_value(
+    pool_script_hash,
+    identifier,
+    pool_value,
+    outcome,
+    protocol_fees,
+  )
 }