stable_pool: math

contracts/pool_stable/src/contract.rs

@@ -345,8 +345,8 @@ impl StableLiquidityPoolTrait for StableLiquidityPool {
             &env,
             amp as u128,
             &[
-                scale_value(new_balance_a, token_a_decimals, DECIMAL_PRECISION),
-                scale_value(new_balance_b, token_b_decimals, DECIMAL_PRECISION),
+                scale_value(&env, new_balance_a, token_a_decimals, DECIMAL_PRECISION),
+                scale_value(&env, new_balance_b, token_b_decimals, DECIMAL_PRECISION),
             ],
         );
 
@@ -373,11 +373,13 @@ impl StableLiquidityPoolTrait for StableLiquidityPool {
                 amp as u128,
                 &[
                     scale_value(
+                        &env,
                         convert_i128_to_u128(old_balance_a),
                         token_a_decimals,
                         DECIMAL_PRECISION,
                     ),
                     scale_value(
+                        &env,
                         convert_i128_to_u128(old_balance_b),
                         token_b_decimals,
                         DECIMAL_PRECISION,
@@ -1029,13 +1031,14 @@ pub fn compute_swap(
         env,
         amp as u128,
         scale_value(
+            env,
             offer_pool + offer_amount,
             greatest_precision,
             DECIMAL_PRECISION,
         ),
         &[
-            scale_value(offer_pool, offer_pool_precision, DECIMAL_PRECISION),
-            scale_value(ask_pool, ask_pool_precision, DECIMAL_PRECISION),
+            scale_value(env, offer_pool, offer_pool_precision, DECIMAL_PRECISION),
+            scale_value(env, ask_pool, ask_pool_precision, DECIMAL_PRECISION),
         ],
         greatest_precision,
     );
@@ -1084,13 +1087,14 @@ pub fn compute_offer_amount(
         env,
         amp as u128,
         scale_value(
+            env,
             ask_pool - convert_i128_to_u128(before_commission),
             greatest_precision,
             DECIMAL_PRECISION,
         ),
         &[
-            scale_value(offer_pool, offer_pool_precision, DECIMAL_PRECISION),
-            scale_value(ask_pool, ask_pool_precision, DECIMAL_PRECISION),
+            scale_value(env, offer_pool, offer_pool_precision, DECIMAL_PRECISION),
+            scale_value(env, ask_pool, ask_pool_precision, DECIMAL_PRECISION),
         ],
         greatest_precision,
     );

contracts/pool_stable/src/math.rs

@@ -20,18 +20,28 @@ const DECIMAL_FRACTIONAL: u128 = 1_000_000_000_000_000_000;
 /// 1e-6
 const TOL: u128 = 1000000000000;
 
-pub fn scale_value(atomics: u128, decimal_places: u32, target_decimal_places: u32) -> u128 {
-    const TEN: u128 = 10;
-
-    if decimal_places < target_decimal_places {
-        let factor = TEN.pow(target_decimal_places - decimal_places);
-        atomics
-            .checked_mul(factor)
-            .expect("Multiplication overflow")
+pub fn scale_value(
+    env: &Env,
+    atomics: u128,
+    decimal_places: u32,
+    target_decimal_places: u32,
+) -> u128 {
+    let ten = U256::from_u128(env, 10);
+    let atomics = U256::from_u128(env, atomics);
+
+    let scaled_value = if decimal_places < target_decimal_places {
+        let power = target_decimal_places - decimal_places;
+        let factor = ten.pow(power);
+        atomics.mul(&factor)
     } else {
-        let factor = TEN.pow(decimal_places - target_decimal_places);
-        atomics.checked_div(factor).expect("Division overflow")
-    }
+        let power = decimal_places - target_decimal_places;
+        let factor = ten.pow(power);
+        atomics.div(&factor)
+    };
+
+    scaled_value
+        .to_u128()
+        .expect("Value doesn't fit into u128!")
 }
 
 fn abs_diff(a: &U256, b: &U256) -> U256 {
@@ -74,9 +84,8 @@ pub(crate) fn compute_current_amp(env: &Env, amp_params: &AmplifierParameters) -
 /// A * sum(x_i) * n**n + D = A * D * n**n + D**(n+1) / (n**n * prod(x_i))
 pub fn compute_d(env: &Env, amp: u128, pools: &[u128]) -> U256 {
     let leverage = U256::from_u128(env, (amp / AMP_PRECISION as u128) * N_COINS_PRECISION);
-    let amount_a_times_coins = pools[0] * N_COINS;
-    let amount_b_times_coins = pools[1] * N_COINS;
-
+    let amount_a_times_coins = U256::from_u128(env, pools[0]).mul(&U256::from_u128(env, N_COINS));
+    let amount_b_times_coins = U256::from_u128(env, pools[1]).mul(&U256::from_u128(env, N_COINS));
     let sum_x = U256::from_u128(env, pools[0] + pools[1]); // sum(x_i), a.k.a S
     let zero = U256::from_u128(env, 0u128);
     if sum_x == zero {
@@ -88,10 +97,8 @@ pub fn compute_d(env: &Env, amp: u128, pools: &[u128]) -> U256 {
 
     // Newton's method to approximate D
     for _ in 0..ITERATIONS {
-        let d_product = d.pow(3).div(&U256::from_u128(
-            env,
-            amount_a_times_coins * amount_b_times_coins,
-        ));
+        let a_times_b_product = amount_a_times_coins.mul(&amount_b_times_coins);
+        let d_product = d.pow(3).div(&a_times_b_product);
         d_previous = d.clone();
         d = calculate_step(env, &d, &leverage, &sum_x, &d_product);
         // Equality with the precision of 1e-6
@@ -137,50 +144,207 @@ fn calculate_step(
     l_val.div(&r_val)
 }
 
-/// Compute the swap amount `y` in proportion to `x`.
+/// Compute the swap amount `y` in proportion to `x` using a partially-reordered formula.
 ///
-/// * **Solve for y**
+/// Original stable-swap equation:
+/// ```text
+/// y² + y * (sum' - (A*n^n - 1) * D / (A * n^n)) = D^(n+1) / (n^(2n) * prod' * A)
 ///
-/// y**2 + y * (sum' - (A*n**n - 1) * D / (A * n**n)) = D ** (n + 1) / (n ** (2 * n) * prod' * A)
-///
-/// y**2 + b*y = c
+/// => y² + b·y = c
+/// ```
+/// We chunk up multiplications/divisions to avoid overflow when computing `d³ * amp_prec`.
 pub(crate) fn calc_y(
     env: &Env,
     amp: u128,
-    new_amount: u128,
+    new_amount_u128: u128,
     xp: &[u128],
     target_precision: u32,
 ) -> u128 {
-    let n_coins = U256::from_u128(env, N_COINS);
-    let new_amount = U256::from_u128(env, new_amount);
+    // number of coins in the pool, e.g. 2 for a two-coin stableswap.
+    let coins_count = U256::from_u128(env, N_COINS);
+
+    // convert `new_amount_u128` to U256 for big math.
+    let new_u256_amount = U256::from_u128(env, new_amount_u128);
+
+    // compute the stableswap invariant D.
+    let invariant_d = compute_d(env, amp, xp);
+
+    let amp_precision_factor = U256::from_u128(env, (AMP_PRECISION as u128) * DECIMAL_FRACTIONAL);
+
+    // compute "leverage" = amp * DECIMAL_FRACTIONAL * n_coins.
+    let leverage = U256::from_u128(env, amp)
+        .mul(&U256::from_u128(env, DECIMAL_FRACTIONAL))
+        .mul(&coins_count);
 
-    let d = compute_d(env, amp, xp);
-    let leverage = U256::from_u128(env, amp * DECIMAL_FRACTIONAL * N_COINS);
-    let amp_prec = U256::from_u128(env, AMP_PRECISION as u128 * DECIMAL_FRACTIONAL);
+    // ------------------------------------------------------------------
+    // Now we compute:
+    //   c = (D^3 * amp_precision_factor) / (new_amount * n_coins^2 * leverage)
+    // but we do it in multiple steps to prevent overflow.
+    // ------------------------------------------------------------------
 
-    let c = d
-        .pow(3)
-        .mul(&amp_prec)
-        .div(&new_amount.mul(&n_coins.mul(&n_coins)).mul(&leverage));
+    // Step A: D²
+    let invariant_sq = invariant_d.mul(&invariant_d);
 
-    let b = new_amount.add(&d.mul(&amp_prec).div(&leverage));
+    // Step B: multiply coins_count by itself => n_coins^2, then times new_u256_amount.
+    // denominator_chunk1 = n_coins^2 * new_amount
+    let coins_count_sq = coins_count.mul(&coins_count);
+    let denominator_chunk1 = coins_count_sq.mul(&new_u256_amount);
 
-    // Solve for y by approximating: y**2 + b*y = c
-    let mut y_prev;
-    let mut y = d.clone();
+    // Step C: partial factor => (D² / (n_coins^2 * new_amount))
+    let temp_factor1 = invariant_sq.div(&denominator_chunk1);
+
+    // Step D: multiply by D => (D³ / (n_coins^2 * new_amount))
+    let temp_factor2 = temp_factor1.mul(&invariant_d);
+
+    // Step E: multiply by amp_precision_factor => (D³ * amp_prec) / (n_coins^2 * new_amount)
+    let temp_factor3 = temp_factor2.mul(&amp_precision_factor);
+
+    // Step F: finally divide by leverage =>
+    //   c = (D³ * amp_prec) / (n_coins^2 * new_amount * leverage)
+    let constant_c = temp_factor3.div(&leverage);
+
+    // ------------------------------------------------------------------
+    // b = new_amount + (D * amp_precision_factor / leverage)
+    // ------------------------------------------------------------------
+    let coefficient_b = {
+        let scaled_d = invariant_d.mul(&amp_precision_factor).div(&leverage);
+        new_u256_amount.add(&scaled_d)
+    };
+
+    // ------------------------------------------------------------------
+    // Solve for y in the equation:
+    //   y² + b·y = c  ==>  y = (y² + c) / (n_coins·y + b - D)
+    // Using iteration (Newton-like).
+    // ------------------------------------------------------------------
+    let mut y_guess = invariant_d.clone();
     for _ in 0..ITERATIONS {
-        y_prev = y.clone();
-        y = (y.pow(2).add(&c)).div(&(y.mul(&n_coins).add(&b).sub(&d)));
-        if abs_diff(&y, &y_prev) <= U256::from_u128(env, TOL) {
+        let y_prev = y_guess.clone();
+
+        // Numerator = y² + c
+        let numerator = y_guess.pow(2).add(&constant_c);
+
+        // Denominator = n_coins·y + b - D
+        let denominator = coins_count
+            .mul(&y_guess)
+            .add(&coefficient_b)
+            .sub(&invariant_d);
+
+        // Next approximation for y
+        y_guess = numerator.div(&denominator);
+
+        // Check convergence
+        if abs_diff(&y_guess, &y_prev) <= U256::from_u128(env, TOL) {
+            // Scale down from DECIMAL_PRECISION to `target_precision`.
             let divisor = 10u128.pow(DECIMAL_PRECISION - target_precision);
-            return y
+            return y_guess
                 .to_u128()
-                .expect("Pool stable: calc_y: conversion to u128 failed")
+                .expect("calc_y: final y doesn't fit in u128!")
                 / divisor;
         }
     }
 
-    // Should definitely converge in 64 iterations.
-    log!(&env, "Pool Stable: calc_y: y is not converging");
-    panic_with_error!(&env, ContractError::CalcYErr);
+    // If not converged in 64 iterations, we treat that as an error.
+    log!(env, "calc_y: not converging in 64 iterations!");
+    panic_with_error!(env, ContractError::CalcYErr);
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use soroban_sdk::Env;
+
+    #[test]
+    fn test_scale_value_up() {
+        let env = Env::default();
+        // example: 123 (with 3 decimals) => want 6 decimals
+        // 10^(6-3) = 10^3 = 1000, result => 123 * 1000 = 123000
+        let val = scale_value(&env, 123, 3, 6);
+        assert_eq!(val, 123_000);
+    }
+
+    #[test]
+    fn test_scale_value_down() {
+        let env = Env::default();
+        // example: 123_000 (with 6 decimals) => want 3 decimals
+        // 10^(6-3) = 10^3 = 1000, result => 123000 / 1000 = 123
+        let val = scale_value(&env, 123_000, 6, 3);
+        assert_eq!(val, 123);
+    }
+
+    #[test]
+    fn test_scale_value_no_change() {
+        let env = Env::default();
+        // if decimal_places == target_decimal_places, value is unchanged
+        let val = scale_value(&env, 999_999, 5, 5);
+        assert_eq!(val, 999_999);
+    }
+
+    #[test]
+    fn test_scale_value_big_numbers() {
+        let env = Env::default();
+        // something bigger, e.g. 1_234_567 with decimal_places=2 => target=6
+        // 10^(6-2) = 10^4 = 10000, result => 1234567 * 10000 = 12345670000
+        let val = scale_value(&env, 1_234_567, 2, 6);
+        assert_eq!(val, 12_345_670_000);
+    }
+
+    #[test]
+    fn test_compute_d_zero_sum() {
+        let env = Env::default();
+        // if sum_x=0 => function returns zero
+        let d = compute_d(&env, 100, &[0, 0]);
+        assert_eq!(d, U256::from_u128(&env, 0));
+    }
+
+    #[test]
+    fn test_compute_d_basic() {
+        let env = Env::default();
+        // With amp=100, each of the two tokens having a balance of 1000,
+        // the stableswap invariant D converges to 2000 in the current formula.
+        let amp = 100;
+        let pools = [1000u128, 1000u128];
+
+        let d = compute_d(&env, amp, &pools);
+
+        // Check that we get exactly 2000, which is the expected stable-swap invariant
+        assert_eq!(d, U256::from_u128(&env, 2000));
+    }
+
+    #[test]
+    #[should_panic(expected = "attempt to add with overflow")]
+    fn test_compute_d_non_convergence() {
+        let env = Env::default();
+        // forcing a scenario that should cause no convergence, eg.
+        // unbalanced pools or something that triggers the iteration to never meet TOL.
+        let amp = 1_000_000_000;
+        let pools = [u128::MAX, u128::MAX];
+        compute_d(&env, amp, &pools);
+    }
+
+    #[test]
+    fn test_calc_y_simple() {
+        let env = Env::default();
+        let amp = 100;
+        let xp = [1000u128, 1000u128];
+        let new_amount = 500u128;
+        let target_precision = 6;
+
+        // the math above shows final y == 0 after the scaling.
+        let result = calc_y(&env, amp, new_amount, &xp, target_precision);
+        assert_eq!(result, 0);
+    }
+
+    #[test]
+    #[should_panic(expected = "attempt to add with overflow")]
+    fn test_calc_y_extreme_overflow() {
+        let env = Env::default();
+        // using very large `xp` or `new_amount` to see if we’d attempt a final `y` that can't fit u128
+        calc_y(
+            &env,
+            1_000_000_000_000_000_000, // big `AMP`
+            u128::MAX,                 // big `new_amount`
+            &[u128::MAX, u128::MAX],
+            18,
+        );
+    }
 }