Fix math::sqrt and bump version to v1.0.4

Move.toml

@@ -1,6 +1,6 @@
 [package]
 name = 'swap'
-version = '1.0.3'
+version = '1.0.4'
 license = "Apache-2.0"
 upgrade_policy = "compatible"
 

init/Move.toml

@@ -1,6 +1,6 @@
 [package]
 name = 'init'
-version = '1.0.3'
+version = '1.0.4'
 license = "Apache-2.0"
 upgrade_policy = "compatible"
 

lp/Move.toml

@@ -1,6 +1,6 @@
 [package]
 name = 'lp'
-version = '1.0.3'
+version = '1.0.4'
 license = "Apache-2.0"
 upgrade_policy = "compatible"
 

sources/implements.move

@@ -4,14 +4,15 @@ module swap::implements {
     use std::signer;
     use std::string::{Self, String};
 
-    use aptos_framework::account::{Self, SignerCapability};
-    use aptos_framework::coin::{Self, Coin};
+    use aptos_framework::account::SignerCapability;
+    use aptos_framework::coin;
     use aptos_framework::timestamp;
+
     use lp::lp_coin::LP;
 
     use swap::event;
     use swap::init;
-    use swap::math;
+    use swap::math::{Self, mul_to_u128, mul_div, mul_div_u128, sqrt};
 
     friend swap::interface;
     friend swap::controller;
@@ -126,7 +127,7 @@ module swap::implements {
         account::get_signer_capability_address(&config.fee_cap)
     }
 
-    public fun is_pool_exists<X,Y>():bool acquires Config {
+    public fun is_pool_exists<X, Y>(): bool acquires Config {
         let pool_account = pool_account();
         let pool_address = signer::address_of(&pool_account);
         if (!exists<LiquidityPool<X, Y>>(pool_address)) {
@@ -231,7 +232,7 @@ module swap::implements {
 
         let lp_coins_total = option::extract(&mut coin::supply<LP<X, Y>>());
         let provided_liq = if (0 == lp_coins_total) {
-            let initial_liq = math::sqrt(x_provided_val) * math::sqrt(y_provided_val);
+            let initial_liq = sqrt(mul_to_u128(x_provided_val, y_provided_val));
             assert!(initial_liq > MINIMAL_LIQUIDITY, ERR_LIQUID_NOT_ENOUGH);
             initial_liq - MINIMAL_LIQUIDITY
         } else {
@@ -316,7 +317,7 @@ module swap::implements {
         // Multiply coin_in by the current exchange rate:
         // current_exchange_rate = reserve_out / reserve_in
         // amount_in_after_fees * current_exchange_rate -> amount_out
-        math::mul_div_u128(coin_in_val_after_fees, // scaled to 1000
+        mul_div_u128(coin_in_val_after_fees, // scaled to 1000
             (reserve_out as u128),
             new_reserve_in  // scaled to 1000
         )
@@ -354,7 +355,7 @@ module swap::implements {
         reserve_out: u64,
     ) acquires LiquidityPool, Config {
         let fee_multiplier = FEE_MULTIPLIER / 5; // 20% fee to swap fundation.
-        let fee_value = math::mul_div(coin_in_value, fee_multiplier, FEE_SCALE);
+        let fee_value = mul_div(coin_in_value, fee_multiplier, FEE_SCALE);
 
         let coin_out_value = get_amount_out(coin_in_value, reserve_in, reserve_out);
         assert!(coin_out_value >= coin_out_min_value, ERR_COIN_OUT_NUM_LESS_THAN_EXPECTED_MINIMUM, );
@@ -399,7 +400,7 @@ module swap::implements {
         reserve_out: u64,
     ) acquires LiquidityPool, Config {
         let fee_multiplier = FEE_MULTIPLIER / 5; // 20% fee to swap fundation.
-        let fee_value = math::mul_div(coin_in_value, fee_multiplier, FEE_SCALE);
+        let fee_value = mul_div(coin_in_value, fee_multiplier, FEE_SCALE);
 
         let coin_out_value = get_amount_out(coin_in_value, reserve_in, reserve_out);
         assert!(coin_out_value >= coin_out_min_value, ERR_COIN_OUT_NUM_LESS_THAN_EXPECTED_MINIMUM, );
@@ -483,7 +484,7 @@ module swap::implements {
 
         // exchange_price = reserve_out / reserve_in_size
         // amount_returned = coin_in_val * exchange_price
-        let res = math::mul_div(coin_in, reserve_out, reserve_in);
+        let res = mul_div(coin_in, reserve_out, reserve_in);
         (res as u64)
     }
 

sources/math.move

@@ -6,49 +6,29 @@ module swap::math {
 
     const U64_MAX: u64 = 18446744073709551615;
 
-    /// Get a nearest lower integer Square Root for `x`. Given that this
-    /// function can only operate with integers, it is impossible
-    /// to get perfect (or precise) integer square root for some numbers.
-    ///
-    /// Example:
-    /// ```
-    /// math::sqrt(9) => 3
-    /// math::sqrt(8) => 2 // the nearest lower square root is 4;
-    /// ```
-    ///
-    /// In integer math, one of the possible ways to get results with more
-    /// precision is to use higher values or temporarily multiply the
-    /// value by some bigger number. Ideally if this is a square of 10 or 100.
-    ///
-    /// Example:
-    /// ```
-    /// math::sqrt(8) => 2;
-    /// math::sqrt(8 * 10000) => 282;
-    /// // now we can use this value as if it was 2.82;
-    /// // but to get the actual result, this value needs
-    /// // to be divided by 100 (because sqrt(10000)).
-    ///
-    ///
-    /// math::sqrt(8 * 1000000) => 2828; // same as above, 2828 / 1000 (2.828)
-    /// ```
-    public fun sqrt(
-        x: u64
-    ): u64 {
-        let bit = 1u128 << 64;
-        let res = 0u128;
-        let x = (x as u128);
+    /// Multiple two u64 and get u128, e.g. ((`x` * `y`) as u128).
+    public fun mul_to_u128(x: u64, y: u64): u128 {
+        (x as u128) * (y as u128)
+    }
 
-        while (bit != 0) {
-            if (x >= res + bit) {
-                x = x - (res + bit);
-                res = (res >> 1) + bit;
+    /// Get square root of `y`.
+    /// Babylonian method (https://en.wikipedia.org/wiki/Methods_of_computing_square_roots#Babylonian_method)
+    public fun sqrt(y: u128): u64 {
+        if (y < 4) {
+            if (y == 0) {
+                0u64
             } else {
-                res = res >> 1;
+                1u64
+            }
+        } else {
+            let z = y;
+            let x = y / 2 + 1;
+            while (x < z) {
+                z = x;
+                x = (y / x + x) / 2;
             };
-            bit = bit >> 2;
-        };
-
-        (res as u64)
+            (z as u64)
+        }
     }
 
     /// Implements: `x` * `y` / `z`.

tests/interface_tests.move

@@ -4,16 +4,16 @@ module swap::interface_tests {
     use std::signer;
     use std::string::utf8;
 
-    use aptos_framework::account;
     use aptos_framework::aptos_coin::{Self, AptosCoin};
     use aptos_framework::coin::{Self, MintCapability};
     use aptos_framework::genesis;
 
     use lp::lp_coin::LP;
+
     use swap::implements;
     use swap::init;
     use swap::interface;
-    use swap::math;
+    use swap::math::{sqrt, mul_to_u128};
 
     const MAX_U64: u64 = 18446744073709551615;
 
@@ -112,7 +112,7 @@ module swap::interface_tests {
 
         assert!(
             coin::balance<LP<USDT, XBTC>>(admin_address)
-                == math::sqrt(xbtc_val) * math::sqrt(usdt_val) - 1000,
+                == sqrt(mul_to_u128(xbtc_val, usdt_val)) - 1000,
             coin::balance<LP<USDT, XBTC>>(admin_address)
         );
     }
@@ -313,8 +313,8 @@ module swap::interface_tests {
         assert!(coin::balance<USDT>(user) == usdt_val - usdt_val / 100, 3);
         assert!(coin::balance<XBTC>(user) == xbtc_val - xbtc_val / 100, 4);
         assert!(
-            137840390 == coin::balance<LP<USDT,XBTC>>(user),
-            coin::balance<LP<USDT,XBTC>>(user)
+            137840477 == coin::balance<LP<USDT, XBTC>>(user),
+            coin::balance<LP<USDT, XBTC>>(user)
         )
     }
 
@@ -346,30 +346,28 @@ module swap::interface_tests {
         assert!(coin::balance<USDT>(user) == usdt_val - usdt_val / 100, 3);
         assert!(coin::balance<XBTC>(user) == xbtc_val - xbtc_val / 100, 4);
         assert!(
-            coin::balance<LP<USDT,XBTC>>(user) == 137840390,
-            coin::balance<LP<USDT,XBTC>>(user)
+            coin::balance<LP<USDT, XBTC>>(user) == 137840477,
+            coin::balance<LP<USDT, XBTC>>(user)
         );
 
         interface::remove_liquidity<USDT, XBTC>(
             &user_account,
-            13784039,
+            137840477 / 10,
             1,
             1,
         );
 
         assert!(
-            coin::balance<LP<USDT,XBTC>>(user) == 137840390 - 13784039,
-            coin::balance<LP<USDT,XBTC>>(user)
+            coin::balance<LP<USDT, XBTC>>(user) == 137840477 - 137840477 / 10,
+            coin::balance<LP<USDT, XBTC>>(user)
         );
 
         assert!(
-            coin::balance<USDT>(user)
-                == usdt_val - usdt_val / 100 + usdt_val / 1000,
+            coin::balance<USDT>(user) == 1882899999896,
             coin::balance<USDT>(user)
         );
         assert!(
-            coin::balance<XBTC>(user)
-                == xbtc_val - xbtc_val / 100 + xbtc_val / 1000,
+            coin::balance<XBTC>(user) == 99099999,
             coin::balance<XBTC>(user)
         );
     }
@@ -402,8 +400,8 @@ module swap::interface_tests {
         assert!(coin::balance<USDT>(user) == usdt_val - usdt_val / 100, 3);
         assert!(coin::balance<XBTC>(user) == xbtc_val - xbtc_val / 100, 4);
         assert!(
-            137840390 == coin::balance<LP<USDT,XBTC>>(user),
-            coin::balance<LP<USDT,XBTC>>(user)
+            137840477 == coin::balance<LP<USDT, XBTC>>(user),
+            coin::balance<LP<USDT, XBTC>>(user)
         );
 
         let (reserve_usdt, reserve_xbtc) = implements::get_reserves_size<USDT, XBTC>();