balancer-maths-rust 0.5.0

use crate::common::constants::WAD;
use crate::common::errors::PoolError;
use crate::common::maths::{
    complement_fixed, div_down_fixed, div_up_fixed, mul_down_fixed, mul_up_fixed, pow_down_fixed,
    pow_up_fixed,
};
use alloy_primitives::{uint, U256};

// A minimum normalized weight imposes a maximum weight ratio. We need this due to limitations in the
// implementation of the power function, as these ratios are often exponents.
pub const MIN_WEIGHT: U256 = uint!(10000000000000000_U256); // 0.01e18

// Pool limits that arise from limitations in the fixed point power function (and the imposed 1:100 maximum weight
// ratio).

// Swap limits: amounts swapped may not be larger than this percentage of the total balance.
pub const MAX_IN_RATIO: U256 = uint!(300000000000000000_U256); // 0.3e18
pub const MAX_OUT_RATIO: U256 = uint!(300000000000000000_U256); // 0.3e18

// Invariant growth limit: non-proportional joins cannot cause the invariant to increase by more than this ratio.
pub const MAX_INVARIANT_RATIO: U256 = uint!(3000000000000000000_U256); // 3e18
                                                                       // Invariant shrink limit: non-proportional exits cannot cause the invariant to decrease by less than this ratio.
pub const MIN_INVARIANT_RATIO: U256 = uint!(700000000000000000_U256); // 0.7e18

/// Compute the invariant, rounding down.
///
/// The invariant functions are called by the Vault during various liquidity operations, and require a specific
/// rounding direction in order to ensure safety (i.e., that the final result is always rounded in favor of the
/// protocol. The invariant (i.e., all token balances) must always be greater than 0, or it will revert.
///
/// invariant               _____
/// wi = weight index i      | |      wi
/// bi = balance index i     | |  bi ^   = i
/// i = invariant
pub fn compute_invariant_down(
    normalized_weights: &[U256],
    balances: &[U256],
) -> Result<U256, PoolError> {
    let mut invariant = WAD;

    for i in 0..normalized_weights.len() {
        let pow_result = pow_down_fixed(&balances[i], &normalized_weights[i])?;
        invariant = mul_down_fixed(&invariant, &pow_result)?;
    }

    if invariant == U256::ZERO {
        return Err(PoolError::ZeroInvariant);
    }

    Ok(invariant)
}

/// Compute the invariant, rounding up.
///
/// The invariant functions are called by the Vault during various liquidity operations, and require a specific
/// rounding direction in order to ensure safety (i.e., that the final result is always rounded in favor of the
/// protocol. The invariant (i.e., all token balances) must always be greater than 0, or it will revert.
///
/// invariant               _____
/// wi = weight index i      | |      wi
/// bi = balance index i     | |  bi ^   = i
/// i = invariant
pub fn compute_invariant_up(
    normalized_weights: &[U256],
    balances: &[U256],
) -> Result<U256, PoolError> {
    let mut invariant = WAD;

    for i in 0..normalized_weights.len() {
        invariant = mul_up_fixed(
            &invariant,
            &pow_up_fixed(&balances[i], &normalized_weights[i])?,
        )?;
    }

    if invariant == U256::ZERO {
        return Err(PoolError::ZeroInvariant);
    }

    Ok(invariant)
}

/// Computes how many tokens can be taken out of a pool if `amount_in` are sent, given the
/// current balances and weights.
///
/// outGivenExactIn
/// aO = amountOut
/// bO = balanceOut
/// bI = balanceIn              /      /            bI             \    (wI / wO) \
/// aI = amountIn    aO = bO * |  1 - | --------------------------  | ^            |
/// wI = weightIn               \      \       ( bI + aI )         /              /
/// wO = weightOut
pub fn compute_out_given_exact_in(
    balance_in: &U256,
    weight_in: &U256,
    balance_out: &U256,
    weight_out: &U256,
    amount_in: &U256,
) -> Result<U256, PoolError> {
    if amount_in > &mul_down_fixed(balance_in, &MAX_IN_RATIO)? {
        return Err(PoolError::MaxInRatioExceeded);
    }

    let denominator = balance_in + amount_in;
    let base = div_up_fixed(balance_in, &denominator)?;
    let exponent = div_down_fixed(weight_in, weight_out)?;
    let power = pow_up_fixed(&base, &exponent)?;

    // Because of rounding up, power can be greater than one. Using complement prevents reverts.
    mul_down_fixed(balance_out, &complement_fixed(&power)?)
}

/// Computes how many tokens must be sent to a pool in order to take `amount_out`, given the
/// current balances and weights.
///
/// inGivenExactOut
/// aO = amountOut
/// bO = balanceOut
/// bI = balanceIn              /  /            bO             \    (wO / wI)      \
/// aI = amountIn    aI = bI * |  | --------------------------  | ^            - 1  |
/// wI = weightIn               \  \       ( bO - aO )         /                   /
/// wO = weightOut
pub fn compute_in_given_exact_out(
    balance_in: &U256,
    weight_in: &U256,
    balance_out: &U256,
    weight_out: &U256,
    amount_out: &U256,
) -> Result<U256, PoolError> {
    if amount_out > &mul_down_fixed(balance_out, &MAX_OUT_RATIO)? {
        return Err(PoolError::MaxOutRatioExceeded);
    }

    let base = div_up_fixed(balance_out, &(balance_out - amount_out))?;
    let exponent = div_up_fixed(weight_out, weight_in)?;
    let power = pow_up_fixed(&base, &exponent)?;

    // Because the base is larger than one (and the power rounds up), the power should always be larger than one, so
    // the following subtraction should never revert.
    let ratio = power - WAD;

    mul_up_fixed(balance_in, &ratio)
}

/// Calculate balance out given invariant
///
/// calculateBalanceGivenInvariant
/// o = balanceOut
/// b = balanceIn                      (1 / w)
/// w = weight              o = b * i ^
/// i = invariantRatio
pub fn compute_balance_out_given_invariant(
    current_balance: &U256,
    weight: &U256,
    invariant_ratio: &U256,
) -> Result<U256, PoolError> {
    // Rounds result up overall.
    // Calculate by how much the token balance has to increase to match the invariantRatio.
    let balance_ratio = pow_up_fixed(invariant_ratio, &div_up_fixed(&WAD, weight)?)?;

    mul_up_fixed(current_balance, &balance_ratio)
}