From 91c133026f9149446431a0c1fe5f7a589e3b95d3 Mon Sep 17 00:00:00 2001
From: Pi Lanningham <pi.lanningham@gmail.com>
Date: Mon, 12 Feb 2024 12:53:42 -0500
Subject: [PATCH] Resolve SSW-306
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Since we already know the number of LP tokens being minted, we can check
that it is correct with an is_sqrt function, rather than calculating the
sqrt directly.

We're not overly concerned with the pool creation, but it was simple
enough to implement.

Base:
    │ PASS [mem:  1173694, cpu:  470849067] mint_test
Fixed:
    │ PASS [mem:  1056495, cpu:  395561132] mint_test
---
 lib/shared.ak      | 13 +++++++++++++
 validators/pool.ak | 10 ++++++----
 2 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/lib/shared.ak b/lib/shared.ak
index 82d6bb9..4c770b3 100644
--- a/lib/shared.ak
+++ b/lib/shared.ak
@@ -128,6 +128,19 @@ pub fn count_orders(tx_inputs: List<Input>) -> Int {
   )
 }
 
+// Taken from unmerged PR: https://github.com/aiken-lang/stdlib/pull/73/files
+pub fn is_sqrt(self: Int, x: Int) -> Bool {
+  x * x <= self && ( x + 1 ) * ( x + 1 ) > self
+}
+
+test is_sqrt1() {
+  is_sqrt(44203, 210)
+}
+
+test is_sqrt2() {
+  is_sqrt(975461057789971041, 987654321)
+}
+
 /// Check whether a specific value has *exactly* the right amount of tokens
 /// This is important to, for example, efficiently check the pool output has the correct value
 /// Check where this is used, because it's a very subtle bit of logic for where this is safe to do
diff --git a/validators/pool.ak b/validators/pool.ak
index 0a1eaff..0ed2526 100644
--- a/validators/pool.ak
+++ b/validators/pool.ak
@@ -4,7 +4,6 @@ use aiken/dict
 use aiken/hash
 use aiken/interval
 use aiken/list
-use aiken/math
 use aiken/option
 use aiken/transaction.{
   Datum, InlineDatum, Input, Mint, NoDatum, Output, OutputReference,
@@ -388,11 +387,14 @@ validator(settings_policy_id: PolicyId) {
             coin_a_amt
           }
 
-        // Calculate the initial number of LP tokens; In particular, we adopt Uniswaps convention of
-        // computing this value as the square root of the product of the two values.
+        // Check that the quantity of LP tokens is correct; In particular, we adopt Uniswaps convention of
+        // using the sqrt of the product of the two values for the initial number of LP tokens to mint..
         // This helps minimize precision loss: it gives decent initial liquidity values for a range of
         // sizes of pools, such that an individual LP token is granular enough for depositing and withdrawing for most users.
-        expect Some(initial_lq) = math.sqrt(coin_a_amt_sans_protocol_fees * coin_b_amt)
+        // In particular, though, we don't calculate the sqrt here, which is an expensive function; we instead verify that the
+        // amount minted is valid by checking that it squares to the correct product
+        let initial_lq = pool_output_datum.circulating_lp
+        expect shared.is_sqrt(coin_a_amt_sans_protocol_fees * coin_b_amt, initial_lq)
 
         // And check that we mint the correct tokens, and nothing else.
         let expected_mint =