balancer_maths_rust/pools/stable/

stable_math.rs

use crate::common::errors::PoolError;
use crate::common::maths::div_up;
use alloy_primitives::U256;

/// Amplification precision
pub const AMP_PRECISION: u64 = 1000;

/// Invariant growth limit: non-proportional add cannot cause the invariant to increase by more than this ratio.
pub const _MIN_INVARIANT_RATIO: u64 = 60e16 as u64; // 60%
/// Invariant shrink limit: non-proportional remove cannot cause the invariant to decrease by less than this ratio.
pub const _MAX_INVARIANT_RATIO: u64 = 500e16 as u64; // 500%

/// Calculate the invariant for the stable swap curve.
///
/// Returns `Ok(0)` when every balance is zero (empty/uninitialized pool), matching the
/// Solidity/TS/Python reference. A single zero balance among non-zero others makes the
/// invariant undefined; those inputs return `Err(PoolError::StableZeroBalance)` instead of
/// panicking on divide-by-zero (Solidity would revert with `Panic(0x12)`).
pub fn compute_invariant(
    amplification_parameter: &U256,
    balances: &[U256],
) -> Result<U256, PoolError> {
    // Calculate the sum of balances
    let total_balance: U256 = balances.iter().sum();
    let num_tokens = balances.len() as u64;

    if total_balance == U256::ZERO {
        return Ok(U256::ZERO);
    }

    // Any individual zero balance would divide by zero in the D_P product below.
    if balances.iter().any(|b| b.is_zero()) {
        return Err(PoolError::StableZeroBalance);
    }

    // Initial invariant and amplification
    let mut invariant = total_balance;
    let amp_times_total = amplification_parameter * U256::from(num_tokens);

    // Iteratively compute the invariant
    for _ in 0..255 {
        let mut d_p = invariant;

        for balance in balances {
            d_p = (d_p * invariant) / (balance * U256::from(num_tokens));
        }

        let prev_invariant = invariant;

        let numerator = (amp_times_total * total_balance) / U256::from(AMP_PRECISION);
        let numerator = numerator + (d_p * U256::from(num_tokens));
        let numerator = numerator * invariant;

        let amp_minus_precision = amp_times_total - U256::from(AMP_PRECISION);
        let denominator = (amp_minus_precision * invariant) / U256::from(AMP_PRECISION);
        let denominator = denominator + (U256::from(num_tokens + 1) * d_p);

        invariant = numerator / denominator;

        // Check for convergence
        if invariant > prev_invariant {
            if invariant - prev_invariant <= U256::ONE {
                return Ok(invariant);
            }
        } else if prev_invariant - invariant <= U256::ONE {
            return Ok(invariant);
        }
    }

    Err(PoolError::StableInvariantDidntConverge)
}

/// Compute how many tokens can be taken out of a pool if `token_amount_in` are sent
pub fn compute_out_given_exact_in(
    amplification_parameter: &U256,
    balances: &[U256],
    token_index_in: usize,
    token_index_out: usize,
    token_amount_in: &U256,
    invariant: &U256,
) -> Result<U256, PoolError> {
    let mut balances_copy = balances.to_vec();

    // Add the token amount to the input balance
    balances_copy[token_index_in] += token_amount_in;

    // Calculate the final balance out
    let final_balance_out = compute_balance(
        amplification_parameter,
        &balances_copy,
        invariant,
        token_index_out,
    )?;

    // Calculate and return the amount of tokens out, rounding down
    Ok(balances_copy[token_index_out] - final_balance_out - U256::ONE)
}

/// Compute how many tokens must be sent to a pool to take out `token_amount_out`
pub fn compute_in_given_exact_out(
    amplification_parameter: &U256,
    balances: &[U256],
    token_index_in: usize,
    token_index_out: usize,
    token_amount_out: &U256,
    invariant: &U256,
) -> Result<U256, PoolError> {
    if &balances[token_index_out] <= token_amount_out {
        return Err(PoolError::TokenAmountOutIsGreaterThanBalance);
    }

    let mut balances_copy = balances.to_vec();

    // Subtract the token amount from the output balance
    balances_copy[token_index_out] -= token_amount_out;

    // Calculate the final balance in
    let final_balance_in = compute_balance(
        amplification_parameter,
        &balances_copy,
        invariant,
        token_index_in,
    )?;

    // Calculate and return the amount of tokens in, rounding up
    Ok(final_balance_in - balances_copy[token_index_in] + U256::ONE)
}

/// Compute the balance of a token given the invariant.
///
/// Zero balances (or a zero invariant) make the stable balance equation undefined; return
/// a typed error rather than panicking on divide-by-zero.
pub fn compute_balance(
    amplification_parameter: &U256,
    balances: &[U256],
    invariant: &U256,
    token_index: usize,
) -> Result<U256, PoolError> {
    if invariant.is_zero() {
        return Err(PoolError::ZeroInvariant);
    }
    if balances.iter().any(|b| b.is_zero()) {
        return Err(PoolError::StableZeroBalance);
    }

    let num_tokens = balances.len() as u64;
    let amp_times_total = amplification_parameter * U256::from(num_tokens);

    // Calculate sum and P_D
    let mut sum = balances[0];
    let mut p_d = balances[0] * U256::from(num_tokens);

    for balance in balances.iter().skip(1) {
        p_d = (p_d * balance * U256::from(num_tokens)) / invariant;
        sum += balance;
    }

    sum -= &balances[token_index];

    // Calculate inv2 and c
    let inv2 = invariant * invariant;
    let c = div_up(
        &(inv2 * U256::from(AMP_PRECISION)),
        &(amp_times_total * p_d),
    )? * balances[token_index];

    let b = sum + (invariant * U256::from(AMP_PRECISION)) / amp_times_total;

    // Initial approximation
    let mut token_balance = div_up(&(inv2 + c), &(invariant + b))?;

    // Iteratively solve for tokenBalance
    for _i in 0..255 {
        let prev_token_balance = token_balance;
        token_balance = div_up(
            &(token_balance * token_balance + c),
            &(token_balance * U256::from(2) + b - invariant),
        )?;

        // Check for convergence
        if token_balance > prev_token_balance {
            if token_balance - prev_token_balance <= U256::ONE {
                return Ok(token_balance);
            }
        } else if prev_token_balance - token_balance <= U256::ONE {
            return Ok(token_balance);
        }
    }

    Err(PoolError::StableInvariantDidntConverge)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn compute_invariant_all_zero_balances_returns_zero() {
        let result = compute_invariant(&U256::from(1000u64), &[U256::ZERO, U256::ZERO]);
        assert_eq!(result, Ok(U256::ZERO));
    }

    #[test]
    fn compute_invariant_partial_zero_balance_errors() {
        // 3-token pool with one drained balance — previously panic'd on / 0.
        let balances = [
            U256::from(1_000_000u64),
            U256::ZERO,
            U256::from(1_000_000u64),
        ];
        let result = compute_invariant(&U256::from(1000u64), &balances);
        assert!(matches!(result, Err(PoolError::StableZeroBalance)));
    }

    #[test]
    fn compute_balance_zero_balance_errors() {
        let balances = [
            U256::from(1_000_000u64),
            U256::ZERO,
            U256::from(1_000_000u64),
        ];
        let result = compute_balance(
            &U256::from(1000u64),
            &balances,
            &U256::from(2_000_000u64),
            0,
        );
        assert!(matches!(result, Err(PoolError::StableZeroBalance)));
    }

    #[test]
    fn compute_balance_zero_invariant_errors() {
        let balances = [U256::from(1_000_000u64), U256::from(1_000_000u64)];
        let result = compute_balance(&U256::from(1000u64), &balances, &U256::ZERO, 0);
        assert!(matches!(result, Err(PoolError::ZeroInvariant)));
    }
}