Aqua Pool Integration Guide

Introduction

The Aqua Pool is a decentralized liquidity pool that enables efficient token swaps. Integrating with the Aqua Pool involves obtaining specific pool data, performing calculations to determine the expected output amount for a given input, and executing the swap transaction. This guide will walk you through the necessary steps and provide code examples to facilitate the integration process.

Step 1: Obtaining Aqua Pool Data

To interact with the Aqua Pool, you first need to obtain the pool's data. This can be achieved by either:

• Calling the smart contract's functions directly.

• Using the getAquaPoolData function.

Required Parameters

The following parameters are required to obtain the pool data:

• sender: The address of the user initiating the transaction.

• tokenIn: The address of the input token.

• tokenOut: The address of the output token.

• amountIn: The amount of the input token that the user wants to swap.

Data Freshness

If you already have the necessary pool data, you may proceed to Step 2. However, please note that using outdated data may lead to inaccurate results due to changes in the pool's state. It is recommended to always fetch the latest data before performing calculations.

Step 2: Calculating Output Amount

Once you have obtained the pool data, you can calculate the expected output amount by calling the getAmountOut function.

Function Usage

To use the getAmountOut function, provide the following arguments:

• poolData: The data obtained in Step 1.

• tokenIn: The address of the input token.

• amountIn: The amount of the input token.

• sender: The address of the user initiating the transaction.

• shouldEstimateTweakPrice

The function will return:

• amountOut: The expected amount of the output token that the user will receive.

Step 3: Performing the Swap

After estimating the output amount, you can proceed to perform the actual swap transaction by calling the swap function of the Aqua Pool contract.

Swap Function Parameters

To perform the swap, you need to call the swap function with the following parameters:

• _data: Encoded data containing:

  • tokenIn: Address of the input token.

  • to: Address where the output tokens will be sent.

  • withdrawMode: A uint8 value indicating the withdrawal mode (explained below).

• _sender: Address of the user initiating the swap.

• _callback: Address of a callback contract (can be set to zero address if not used).

• _callbackData: Additional data for the callback (can be empty if not used).

Withdraw Mode

The withdrawMode parameter determines how the output tokens are handled, particularly for tokens like ETH that can exist in native or wrapped form.

Withdraw modes:

• 0 - Vault Internal Transfer: The output tokens remain within the vault system and are transferred internally.

• 1 - Withdraw and Unwrap to Native ETH: The output tokens are withdrawn from the vault and unwrapped to native ETH. Use this mode if you want to receive ETH directly.

• 2 - Withdraw and Wrap to wETH: The output tokens are withdrawn from the vault and wrapped into wETH (Wrapped Ether). Use this mode if you prefer wETH tokens.

Example Usage

Encoding Swap Data

Before calling the swap function, you need to encode the _data parameter using ABI encoding:

Examples and Code Snippets

Below are examples demonstrating how to implement the above steps in code. Please refer to the Appendix for detailed function definitions and additional code.

Example

import { BigNumber } from "ethers";
import { JsonRpcProvider } from "@ethersproject/providers";

const sender = '0x0000000000000000000000000000000000000000';
const poolAddress = '0x40B768De8b2E4ed83d982804Cb2FcC53D2529bE9'; // aqua pool address
const tokenOut = '0x5a7d6b2f92c77fad6ccabd7ee0624e64907eaf3e'; // zk 
const tokenIn = '0x3355df6d4c9c3035724fd0e3914de96a5a83aaf4'; // usdc.e

const amountIn = BigNumber.from(1000).mul(DECIMALS_6); // 1000 zk
const provider = new JsonRpcProvider('https://zksync-mainnet.g.alchemy.com/v2/fwRki-oc766ZWMFK79Qy4ChKzkdwfyx7/');

const pool: Crypto_pool = Crypto_pool__factory.connect(poolAddress, provider)
// const pool = new ethers.Contract(poolAddress, Crypto_pool__factory.abi, provider);


const poolData = await getAquaPoolData({
    pool: pool,
    sender,
    tokenIn,
    tokenOut,
});

const result = getAmountOut(poolData, tokenIn, amountIn, sender, true);

console.log(`Expected output amount: ${ethers.utils.formatUnits(result, tokenOutDecimals)}`);

const swapData = utils.defaultAbiCoder.encode(
    ['address', 'address', 'uint8'],
    [tokenIn, sender, 0]
);

const signer = provider.getSigner();

// Create a contract instance with the signer
const poolWithSigner: Crypto_pool = Crypto_pool__factory.connect(poolAddress, signer)

// Call the swap function
const tx = await poolWithSigner.swap(
  swapData,
  sender,
  ethers.constants.AddressZero, // No callback
  '0x' // Empty callback data
);

// Wait for the transaction to be mined
const receipt = await tx.wait();

console.log('Swap transaction successful:', receipt.transactionHash);

Appendix

1. Get a aqua pool info

Params

token0

Returns the address of the first token in the pool.

token1

Returns the address of the second token in the pool.

getReserves

Returns the current reserves of the tokens in the pool.

priceScale

Returns the current price scale, which determines the price band around which liquidity is concentrated.

token0PrecisionMultiplie, token1PrecisionMultiplie

Multipliers for each pooled token's precision to get to the pool precision decimals.

which is agnostic to the pool, but usually is 18.

For example, TBTC has 18 decimals, so the multiplier should be 10 ** (18 - 18) = 1. WBTC has 8, so the multiplier should be 10 ** (18 - 8) => 10 ** 10.

The value is only for crypto pools, and has no effects on non-crypto pools.

invariantLast

Invariant of the pool as of immediately after the most recent liquidity event.

a, futureA, gamma

Ramp params

totalSupply

virtualPrice

virtual price.

PoolType

Pool type 3 for aqua pools.

Code

import { Crypto_pool } from "@/abis";

async function fetchPoolData<T extends keyof Crypto_pool>(
    pool: Crypto_pool,
    calls: [T, Parameters<Extract<Crypto_pool[T], (...args: any) => any>>?][]
): Promise<{ [K in T]: Awaited<ReturnType<Extract<Crypto_pool[K], (...args: any) => any>>> }> {
    const callPromises = calls.map(([method, args]) => {
        const poolMethod = pool[method];
        if (typeof poolMethod === "function") {
            return (poolMethod as (this: Crypto_pool, ...args: any[]) => any).apply(pool, args || []); // Call with args if provided
        } else {
            throw new Error(`Property ${String(method)} is not a callable function`);
        }
    });

    const results = await Promise.all(callPromises);

    return Object.fromEntries(calls.map(([method], index) => [method, results[index]])) as {
        [K in T]: Awaited<ReturnType<Extract<Crypto_pool[K], (...args: any) => any>>>;
    };
}


export async function getAquaPoolData(pool: Crypto_pool) {
    const poolData = await fetchPoolData(pool, [
        ["token0"], // pool.token0()
        ["token1"],
        ["reserve0"],
        ["reserve1"],
        ["token0PrecisionMultiplier"],
        ["token1PrecisionMultiplier"],
        ["invariantLast"],
        ["getParams"],
        ["getReserves"],
        ["totalSupply"],
        ["virtualPrice"],
        ["priceScale"],
        ['poolType'],
    ]);
    return poolData;
}

2. Quote amount out

typescript

import { BigNumber, utils } from "ethers";
import { JsonRpcProvider } from "@ethersproject/providers";
import {
    DECIMALS_18,
    DECIMALS_6,
    ONE,
    ZERO,
    ZERO_ADDRESS,
    computeD, estimateTweakPriceAqua, getY
} from "./utils/math";
import { getAquaPoolData } from "./utils";

import { Crypto_pool } from "./abis/Crypto_pool";
import { Crypto_pool__factory } from "./abis/Crypto_pool__factory";



interface GetAmountOutArgs {
    reserve0: BigNumber;
    reserve1: BigNumber;
    a: BigNumber;
    gamma: BigNumber;
    futureTime: BigNumber;
    priceScale: BigNumber;
    invariant: BigNumber;
    token0PrecisionMultiplier: BigNumber;
    token1PrecisionMultiplier: BigNumber;
    swapFee: BigNumber;
    virtualPrice: BigNumber;
    totalSupply: BigNumber;
}

interface GetAmountOutArgs2 {
    xp0: BigNumber;
    xp1: BigNumber;
    y: BigNumber;
    xp1PriceScale: BigNumber;
}

export function getAmountOut(
    poolData: any,
    tokenIn: string,
    amountIn: BigNumber,
    sender: string,
    shouldEstimateTweakPrice?: boolean
): BigNumber {
    const it: GetAmountOutArgs = {} as GetAmountOutArgs;
    [it.reserve0, it.reserve1] = poolData.getReserves;
    [it.a, it.gamma, it.futureTime] = poolData.getParams;
    it.swapFee = BigNumber.from(poolData.getSwapFee);

    it.priceScale = poolData.priceScale;
    it.token0PrecisionMultiplier = poolData.token0PrecisionMultiplier;
    it.token1PrecisionMultiplier = poolData.token1PrecisionMultiplier;
    it.totalSupply = poolData.totalSupply;
    it.virtualPrice = poolData.virtualPrice;

    const currentTime = BigNumber.from(Math.floor(Date.now() / 1000));

    if (it.futureTime.gt(currentTime)) {
        const xp0 = it.reserve0.mul(it.token0PrecisionMultiplier);
        const xp1 = it.token1PrecisionMultiplier.mul(it.reserve1).mul(it.priceScale).div(DECIMALS_18);

        it.invariant = computeD(it.a, it.gamma, xp0, xp1);
    } else {
        it.invariant = poolData.invariantLast;
    }

    const token0 = poolData.token0;
    const token1 = poolData.token1;


    const [_amountOut, _] = _getAmountOut(
        token0,
        token1,
        sender,
        amountIn,
        tokenIn.toLowerCase() === token0.toLowerCase(),
        it,
        shouldEstimateTweakPrice
    );

    return _amountOut;
}


function _getAmountOut(
    _token0: string,
    _token1: string,
    _sender: string,
    _amountIn: BigNumber,
    _token0In: boolean,
    args: GetAmountOutArgs,
    shouldEstimateTweakPrice?: boolean
): [BigNumber, BigNumber] {
    let _amountOut = BigNumber.from(0);
    let _feeIn = BigNumber.from(0);
    let newXpOut; // used for estimateTweakPriceAqua

    if (!_amountIn.isZero()) {
        if (_token0In) {
            args.reserve0 = args.reserve0.add(_amountIn);
        } else {
            args.reserve1 = args.reserve1.add(_amountIn);
        }
        const it: GetAmountOutArgs2 = {} as GetAmountOutArgs2;
        it.xp0 = args.reserve0.mul(args.token0PrecisionMultiplier);
        it.xp1PriceScale = args.priceScale.mul(args.token1PrecisionMultiplier);
        it.xp1 = args.reserve1.mul(it.xp1PriceScale).div(DECIMALS_18);

        if (_token0In) {

            it.y = getY(args.a, args.gamma, it.xp0, it.xp1, args.invariant, 1);
            _amountOut = it.xp1.sub(it.y);
            it.xp1 = it.y;

            newXpOut = it.xp1.sub(it.y);

            _amountOut = _amountOut.sub(ONE);
            _amountOut = _amountOut.mul(DECIMALS_18).div(it.xp1PriceScale);

        } else {

            it.y = getY(args.a, args.gamma, it.xp0, it.xp1, args.invariant, 0);
            _amountOut = it.xp0.sub(it.y);
            it.xp0 = it.y;

            newXpOut = it.xp0.sub(it.y);

            _amountOut = _amountOut.sub(ONE);
            _amountOut = _amountOut.div(args.token0PrecisionMultiplier);

        }
        _feeIn = args.swapFee.mul(_amountOut).div(BigNumber.from('100000'));
        _amountOut = _amountOut.sub(_feeIn);

        if (shouldEstimateTweakPrice) {
            if (_token0In) {
                newXpOut = args.reserve1.sub(_amountOut).mul(it.xp1PriceScale.mul(args.token1PrecisionMultiplier)).div(DECIMALS_18);
                estimateTweakPriceAqua(
                    args.a, args.gamma, it.xp0, newXpOut, ZERO, args.virtualPrice, args.priceScale, args.totalSupply, args.futureTime
                );
            } else {
                newXpOut = args.reserve0.sub(_amountOut).mul(args.token0PrecisionMultiplier);
                estimateTweakPriceAqua(
                    args.a, args.gamma, newXpOut, it.xp0, ZERO, args.virtualPrice, args.priceScale, args.totalSupply, args.futureTime
                );
            }
        }
    }

    return [_amountOut, _feeIn];
}

solidity

struct GetAmountOutArgs {
    uint xp0;
    uint xp1;
    uint invariant;
    uint xp1PriceScale;
    uint y;
}

struct GetAmountOutInputParams {
    address sender;
    uint a;
    uint gamma;
    uint balance0;
    uint balance1;
    uint reserve0;
    uint reserve1;
    bool token0In;
    uint futureTime;
}

function _getAmountOut(
    GetAmountOutInputParams memory p,
    ICryptoMath _MATH
)
    private
    view
    returns (
        uint _amountOut,
        uint24 _swapFee,
        uint _feeAmountOut,
        uint _newXp0,
        uint _newXp1,
        uint _k0_prev
    )
{
    GetAmountOutArgs memory it;
    it.xp0 = p.balance0 * token0PrecisionMultiplier;
    it.xp1PriceScale = priceScale * token1PrecisionMultiplier;
    it.xp1 = Math.mulDivUnsafeLast(p.balance1, it.xp1PriceScale, 1e18);

    if (p.futureTime > block.timestamp) {
        if (p.token0In) {
            // Resue `it.y` as `xp0` here
            it.y = unsafe_mul(p.reserve0, token0PrecisionMultiplier); // <---- already safe_mul above for `xp0`

            it.invariant = _MATH.computeD(p.a, p.gamma, it.y, it.xp1, 0);
        } else {
            // Resue `it.y` as `xp1` here
            it.y = unsafe_div(unsafe_mul(p.reserve1, it.xp1PriceScale), 1e18); // <---- already safe_mul above for `xp1`

            it.invariant = _MATH.computeD(p.a, p.gamma, it.xp0, it.y, 0);
        }
    } else {
        it.invariant = invariantLast;
    }

    if (p.token0In) {
        (it.y, _k0_prev) = _MATH.getY(
            p.a,
            p.gamma,
            it.xp0,
            it.xp1,
            it.invariant,
            1
        );
        _amountOut = it.xp1 - it.y;
        it.xp1 -= _amountOut;

        _amountOut -= 1;
        _amountOut = Math.mulDivUnsafeFirst(1e18, _amountOut, it.xp1PriceScale);

        _swapFee = _getFee(p.sender, token0, token1, it.xp0, it.xp1);
        _feeAmountOut = Math.mulDivUnsafeLast(_swapFee, _amountOut, 1e5);
        _amountOut -= _feeAmountOut;

        _newXp0 = it.xp0;
        _newXp1 = unsafe_div(
            unsafe_mul(p.reserve1 - _amountOut, it.xp1PriceScale),
            1e18
        ); // <---- already safe_mul above for `xp1`
    } else {
        (it.y, _k0_prev) = _MATH.getY(
            p.a,
            p.gamma,
            it.xp0,
            it.xp1,
            it.invariant,
            0
        );
        _amountOut = it.xp0 - it.y;
        it.xp0 -= _amountOut;

        _amountOut -= 1;
        _amountOut = Math.divUnsafeLast(_amountOut, token0PrecisionMultiplier);

        _swapFee = _getFee(p.sender, token1, token0, it.xp0, it.xp1);
        _feeAmountOut = Math.mulDivUnsafeLast(_swapFee, _amountOut, 1e5);
        _amountOut -= _feeAmountOut;

        _newXp0 = unsafe_mul(
            p.reserve0 - _amountOut,
            token0PrecisionMultiplier
        );
        _newXp1 = it.xp1;
    }
}

3. Math functions details

3.1 computedD

const N_COINS = BigNumber.from(2);
const A_MULTIPLIER = BigNumber.from(10000);

function computedD(ANN: BigNumber, gamma: BigNumber, x0: BigNumber, x1: BigNumber): BigNumber {

    if (x0.isZero() || x1.isZero()) {
        return ZERO;
    }

    // Initial value of invariant D is that for constant-product invariant
    let x: [BigNumber, BigNumber];//[x0, x1].sort((a, b) => b.sub(a).isNegative() ? -1 : 1);
    if (x0.lt(x1)) {
        x = [x1, x0];
    } else {
        x = [x0, x1];
    }

    if (!(x[0].gte(DECIMALS_9) && x[0].lte(DECIMALS_33))) {
        //console.warn("computedD: unsafe values x[0]", x[0].toString());
        return ZERO;
    }
    if (!(x[1].mul(DECIMALS_18).div(x[0]).gte(DECIMALS_14))) {
        //console.warn("computedD: unsafe values x[i] (input)", x[1].toString());
        return ZERO;
    }

    let D = N_COINS.mul(geometricMean(x0, x1));
    const S = x0.add(x1);
    const __g1k0 = gamma.add(DECIMALS_18);

    for (let i = 0; i < 255; i++) {
        const D_prev = D;
        //invariant(D.gt(ZERO), "Unsafe value D");
        if (D.lte(ZERO)) {
            //console.warn("Unsafe value D", D.toString());
            return ZERO;
        }

        const K0 = ((DECIMALS_18.mul(N_COINS.pow(2))).mul(x[0]).div(D)).mul(x[1]).div(D);

        let _g1k0 = __g1k0;
        if (_g1k0.gt(K0)) {
            _g1k0 = _g1k0.sub(K0).add(1);
        } else {
            _g1k0 = K0.sub(_g1k0).add(1);
        }

        const mul1 = ((DECIMALS_18.mul(D).div(gamma)).mul(_g1k0).div(gamma)).mul(_g1k0).mul(A_MULTIPLIER).div(ANN);
        const mul2 = ((TWO.mul(DECIMALS_18)).mul(N_COINS)).mul(K0).div(_g1k0);

        const neg_fprime = (S.add(S.mul(mul2).div(DECIMALS_18))).add(mul1.mul(N_COINS).div(K0)).sub(mul2.mul(D).div(DECIMALS_18));

        const D_plus = D.mul(neg_fprime.add(S)).div(neg_fprime);
        let D_minus = D.mul(D).div(neg_fprime);

        if (DECIMALS_18.gt(K0)) {
            D_minus = D_minus.add(D.mul(mul1.div(neg_fprime)).div(DECIMALS_18).mul(DECIMALS_18.sub(K0)).div(K0));
        } else {
            D_minus = D_minus.sub(D.mul(mul1.div(neg_fprime)).div(DECIMALS_18).mul(K0.sub(DECIMALS_18)).div(K0));
        }

        if (D_plus.gt(D_minus)) {
            D = D_plus.sub(D_minus);
        } else {
            D = D_minus.sub(D_plus).div(2);
        }

        const diff = D.gt(D_prev) ? D.sub(D_prev) : D_prev.sub(D);

        if (diff.mul(DECIMALS_14).lt(max(DECIMALS_16, D))) {
            const frac = x0.mul(DECIMALS_18).div(D);
            if (!(frac.gte(DECIMALS_16) && frac.lte(DECIMALS_20))) {
                //console.warn("computedD: unsafe values x[i]-1", frac.toString());
                return ZERO;
            }
            //nvariant(frac.gte(DECIMALS_16) && frac.lte(DECIMALS_20), "computedD: unsafe values x[i]-1");

            const frac1 = x1.mul(DECIMALS_18).div(D);
            if (!(frac1.gte(DECIMALS_16) && frac1.lte(DECIMALS_20))) {
                //console.warn("computedD: unsafe values x[i]-2", frac1.toString());
                return ZERO;
            }
            //invariant(frac1.gte(DECIMALS_16) && frac1.lte(DECIMALS_20), "computedD: unsafe values x[i]-2");

            return D;
        }
    }

    //throw new Error("Did not converge");
    console.warn("Did not converge");
    return ZERO;
}

Vyper

Math library

3.2 getY

Typescript

function getY(ANN: BigNumber, gamma: BigNumber, x0: BigNumber, x1: BigNumber, D: BigNumber, i: number): BigNumber {
    if (!(D.gte(DECIMALS_17) && D.lte(DECIMALS_33))) {
        //console.warn("getY: unsafe values D", D.toString());
        return ZERO;
    }

    let x_j = x0;
    if (i === 0) {
        x_j = x1;
    }

    let y = D.pow(TWO).div(x_j.mul(N_COINS.pow(TWO)));
    const K0_i = ETHER.mul(N_COINS).mul(x_j).div(D);

    if (!(K0_i.gte(DECIMALS_16.mul(N_COINS)) && K0_i.lte(DECIMALS_20.mul(N_COINS)))) {
        //console.warn("getY: unsafe values x[i]", K0_i.toString());
        return ZERO;
    }

    const convergence_limit = max(max(x_j.div(DECIMALS_14), D.div(DECIMALS_14)), DECIMALS_2);
    const __g1k0 = gamma.add(ETHER);

    for (let j = 0; j < 255; j++) {
        const y_prev = y;

        const K0 = K0_i.mul(y).mul(N_COINS).div(D);
        const S = x_j.add(y);

        let _g1k0 = __g1k0;
        if (_g1k0.gt(K0)) {
            _g1k0 = _g1k0.sub(K0).add(1);
        } else {
            _g1k0 = K0.sub(_g1k0).add(1);
        }

        const mul1 = (DECIMALS_18.mul(D).div(gamma).mul(_g1k0).div(gamma).mul(_g1k0).mul(A_MULTIPLIER)).div(ANN);
        const mul2 = DECIMALS_18.add((TWO.mul(DECIMALS_18)).mul(K0)).div(_g1k0);

        let yfprime = DECIMALS_18.mul(y).add(S.mul(mul2)).add(mul1);
        const _dyfprime = D.mul(mul2);

        if (yfprime.lt(_dyfprime)) {
            y = y_prev.div(TWO);
            continue;
        } else {
            yfprime = yfprime.sub(_dyfprime);
        }

        const fprime = yfprime.div(y);

        const y_minus_temp = mul1.div(fprime);
        const y_plus = (yfprime.add(ETHER.mul(D))).div(fprime).add(y_minus_temp.mul(ETHER).div(K0));
        const y_minus = y_minus_temp.add(DECIMALS_18.mul(S).div(fprime));

        if (y_plus.lt(y_minus)) {
            y = y_prev.div(2);
        } else {
            y = y_plus.sub(y_minus);
        }

        let diff;
        if (y.gt(y_prev)) {
            diff = y.sub(y_prev);
        } else {
            diff = y_prev.sub(y);
        }

        if (diff.lt(max(convergence_limit, y.div(DECIMALS_14)))) {
            const frac = y.mul(DECIMALS_18).div(D);
            if (!(frac.gte(DECIMALS_16) && frac.lte(DECIMALS_20))) {
                //console.warn("getY: unsafe values for y", frac.toString());
                return ZERO;
            }

            return y;
        }
    }

    //throw new Error("Did not converge");
    console.warn("Did not converge");
    return ZERO;
}

Vyper

Math library

3.3 estimateTweakPriceAqua

typescript

function estimateTweakPriceAqua(
    a: BigNumber, gamma: BigNumber, xp0: BigNumber, xp1: BigNumber, newInvariant: BigNumber,
    oldVirtualPrice: BigNumber, priceScale: BigNumber, totalSupply: BigNumber, futureParamsTime: BigNumber,
) {
    let unadjustedInvariant;
    if (newInvariant.isZero()) {
        unadjustedInvariant = computeD(a, gamma, xp0, xp1);
    } else {
        unadjustedInvariant = newInvariant;
    }

    if (!oldVirtualPrice.isZero()) {
        const _xp0 = unadjustedInvariant.div(2);
        const _xp1 = ETHER.mul(unadjustedInvariant).div(priceScale.mul(2));

        const xcp = geometricMean(_xp0, _xp1);
        const newVirtualPrice = ETHER.mul(xcp).div(totalSupply);
        if (futureParamsTime.lt(BigNumber.from(Math.floor(Date.now() / 1000)))) {
            if (newVirtualPrice.lt(oldVirtualPrice)) {
                console.warn('estimateTweakPriceAqua: loss, newVirtualPrice', newVirtualPrice.toString(), 'oldVirtualPrice', oldVirtualPrice.toString(), 'xcp', xcp.toString(), 'newXp0', _xp0.toString(), 'newXp1', _xp1.toString(), 'priceScale', priceScale.toString());
                throw Error('loss');
            }
        }
    }
}

3.4 _tweakPrice

solidity

struct TweakPriceArgs {
    uint priceScale;
    uint priceOracle;
    uint lastPrices;
    uint lastPricesTimestamp;
    uint totalSupply;
    uint xcpProfit;
    uint virtualPrice;
    uint oldVirtualPrice;
    uint newInvariant;
    uint norm;
}

/// @notice Tweaks `priceOracle`, `lastPrices` and conditionally adjusts `priceScale`.
/// This is called whenever there is an unbalanced liquidity operation: swap, mint, or burnSingle.
/// @dev Contains main liquidity rebalancing logic, by tweaking `priceScale`.
function _tweakPrice(
    uint a,
    uint gamma,
    uint xp0,
    uint xp1,
    uint invariantUnadjusted,
    uint _futureTime,
    uint _k0_prev,
    ICryptoMath _MATH
) private returns (uint, uint) {
    // returns (newInvariant, newPriceScale)
    TweakPriceArgs memory it;

    // ---------------------------- Read storage ------------------------------

    RebalancingParams memory _rebalancingParams = rebalancingParams;

    it.priceOracle = cachedPriceOracle;
    it.priceScale = priceScale;
    it.lastPricesTimestamp = lastPricesTimestamp;

    it.totalSupply = totalSupply;
    it.oldVirtualPrice = virtualPrice;

    // ----------------------- Update MA if needed ----------------------------

    if (it.lastPricesTimestamp < block.timestamp) {
        // The moving average price oracle is calculated using the `lastPrices`
        // of the trade at the previous block, and the price oracle logged
        // before that trade. This can happen only once per block.

        // ------------------ Calculate moving average params -----------------

        // ---------------------------------------------- Update price oracles.

        // ----------------- We cap state price that goes into the EMA with
        //                                                  2 x priceScale.

        it.priceOracle = _MATH.getPriceOracle(
            it.lastPricesTimestamp,
            _rebalancingParams.maTime,
            lastPrices,
            it.priceScale,
            it.priceOracle
        );

        cachedPriceOracle = it.priceOracle;
        lastPricesTimestamp = block.timestamp;
    }

    // `priceOracle` is used further on to calculate its vector
    // distance from `priceScale`. This distance is used to calculate
    // the amount of adjustment to be done to the `priceScale`.

    //  ---------------- If `invariantUnadjusted` is 0, calculate it -----------------

    if (invariantUnadjusted == 0) {
        invariantUnadjusted = _MATH.computeD(a, gamma, xp0, xp1, _k0_prev);
    }

    // ----------------- Calculate and update `lastPrices` --------------------

    lastPrices = _MATH.getLastPrices(
        a,
        gamma,
        xp0,
        xp1,
        invariantUnadjusted,
        it.priceScale
    );

    // ------------------------- Update `xcpProfit` ---------------------------

    uint _xp0;
    uint _xp1;
    if (it.oldVirtualPrice != 0) {
        _xp0 = Math.divUnsafeLast(invariantUnadjusted, 2);
        _xp1 = Math.mulDivUnsafeFirst(
            1e18,
            invariantUnadjusted,
            Math.mulUnsafeFirst(2, it.priceScale)
        );

        it.virtualPrice = Math.mulDivUnsafeFirst(
            1e18,
            Math.geometricMean(_xp0, _xp1),
            it.totalSupply
        );
        it.xcpProfit = Math.divUnsafeLast(
            xcpProfit * it.virtualPrice,
            it.oldVirtualPrice
        );

        // If A and gamma are NOT undergoing ramps (t < block.timestamp),
        // ensure new `virtualPrice` is not less than `oldVirtualPrice`,
        // else the pool suffers a loss.
        if (_futureTime < block.timestamp) {
            if (it.virtualPrice <= it.oldVirtualPrice) {
                revert Loss();
            }
        }
    } else {
        it.xcpProfit = 1e18;
        it.virtualPrice = 1e18;
    }

    xcpProfit = it.xcpProfit;

    // ------------ Rebalance liquidity if there's enough profits to adjust it:
    if (
        Math.mulUnsafeFirst(2, it.virtualPrice) - 1e18 >
        it.xcpProfit + unsafe_mul(_rebalancingParams.allowedExtraProfit, 2)
    ) {
        // ------------------- Get adjustment step ----------------------------

        // Calculate the vector distance between `priceScale` and `priceOracle`.
        it.norm = Math.mulDivUnsafeFirst(1e18, it.priceOracle, it.priceScale);
        if (it.norm > 1e18) {
            it.norm = unsafe_sub(it.norm, 1e18);
        } else {
            it.norm = unsafe_sub(1e18, it.norm);
        }

        // Reuse `_futureTime` as `_adjustmentStep` here to avoid stake too deep errors.
        _futureTime = Math.max(
            _rebalancingParams.adjustmentStep,
            Math.divUnsafeLast(it.norm, 5)
        );

        // We only adjust prices if the vector distance between `priceOracle`
        // and `priceScale` is large enough.
        // This check ensures that no rebalancing occurs if the distance is low
        // i.e. the pool prices are pegged to the oracle prices.
        if (it.norm > _futureTime) {
            // ------------------------------------- Calculate new price scale.

            // Reuse `_futureTime` as `newPriceScale` to avoid stake too deep errors.
            _futureTime = Math.divUnsafeLast(
                it.priceScale *
                    unsafe_sub(it.norm, _futureTime) +
                    (_futureTime * it.priceOracle),
                it.norm
            ); // <- norm is non-zero and gt adjustment_step; unsafe = safe

            // --------------- Update stale xp (using price_scale) with p_new.

            _xp0 = xp0;
            _xp1 = Math.mulDivUnsafeLast(xp1, _futureTime, it.priceScale);
            // unsafe_div because we did safediv before -----^

            // ------------------------------------------ Update D with new xp.

            // Calculate "extended constant product" invariant xCP and virtual price.
            it.newInvariant = _MATH.computeD(a, gamma, _xp0, _xp1, 0);

            // ------------------------------------- Convert xp to real prices.
            _xp0 = Math.divUnsafeLast(it.newInvariant, 2);
            _xp1 = Math.mulDivUnsafeFirst(
                1e18,
                it.newInvariant,
                Math.mulUnsafeFirst(2, _futureTime)
            );

            // ---------- Calculate new virtual price using new xp and D.
            //            Reuse `oldVirtualPrice` (but it has new virtual price).
            it.oldVirtualPrice = Math.mulDivUnsafeFirst(
                1e18,
                Math.geometricMean(_xp0, _xp1),
                it.totalSupply
            );

            // ---------------------------- Proceed if we've got enough profit.
            if (it.oldVirtualPrice > 1e18) {
                if (
                    Math.mulUnsafeFirst(2, it.oldVirtualPrice) - 1e18 >
                    it.xcpProfit
                ) {
                    priceScale = _futureTime; // `newPriceScale`
                    emit UpdatePriceScale(_futureTime);

                    invariantLast = it.newInvariant;
                    emit UpdateInvariant(it.newInvariant);

                    virtualPrice = it.oldVirtualPrice;
                    emit UpdatePoolProfit(it.xcpProfit, it.oldVirtualPrice);

                    return (it.newInvariant, _futureTime);
                }
            }
        }
    }

    // --------- `priceScale` was not adjusted. Update the profit counter and D.

    // If we are here, the `priceScale` adjustment did not happen,
    // Still need to update the profit counter and D.
    invariantLast = invariantUnadjusted;
    emit UpdateInvariant(invariantUnadjusted);

    virtualPrice = it.virtualPrice;
    emit UpdatePoolProfit(it.xcpProfit, it.virtualPrice);

    return (invariantUnadjusted, it.priceScale);
}

3.5 Math library

AquaMath.vy

# pragma version 0.3.10
# pragma evm-version paris

N_COINS: constant(uint256) = 2
A_MULTIPLIER: constant(uint256) = 10000

# MIN_GAMMA: constant(uint256) = 10**10
# MAX_GAMMA: constant(uint256) = 2 * 10**16

# MIN_A: constant(uint256) = N_COINS**N_COINS * A_MULTIPLIER / 10
# MAX_A: constant(uint256) = N_COINS**N_COINS * A_MULTIPLIER * 100000

PRECISION: constant(uint256) = 10**18
EXP_PRECISION: constant(uint256) = 10**10

@external
@pure
def getY(ANN: uint256, gamma: uint256, x0: uint256, x1: uint256, D: uint256, i: uint256) -> uint256[2]:
    """
    Calculating x[i] given other balances x[0..N_COINS-1] and invariant D
    ANN = A * N**N
    """
    # Safety checks
    # assert ANN > MIN_A - 1 and ANN < MAX_A + 1  # dev: unsafe values A
    # assert gamma > MIN_GAMMA - 1 and gamma < MAX_GAMMA + 1  # dev: unsafe values gamma
    assert D > 10**17 - 1 and D < 10**15 * 10**18 + 1 # dev: unsafe values D

    x_j: uint256 = x0
    if i == 0:
        x_j = x1

    y: uint256 = D**2 / (x_j * N_COINS**2)
    K0_i: uint256 = unsafe_div((10**18 * N_COINS) * x_j, D)
    # S_i = x_j

    # frac = x_j * 1e18 / D => frac = K0_i / N_COINS
    assert (K0_i > 10**16*N_COINS - 1) and (K0_i < 10**20*N_COINS + 1)  # dev: unsafe values x[i]

    # x_sorted: uint256[N_COINS] = x
    # x_sorted[i] = 0
    # x_sorted = self.sort(x_sorted)  # From high to low
    # x[not i] instead of x_sorted since x_soted has only 1 element

    convergence_limit: uint256 = max(max(unsafe_div(x_j, 10**14), unsafe_div(D, 10**14)), 100)

    __g1k0: uint256 = unsafe_add(gamma, 10**18)

    for j in range(255):
        y_prev: uint256 = y

        K0: uint256 = unsafe_div(K0_i * y * N_COINS, D)
        S: uint256 = x_j + y

        _g1k0: uint256 = __g1k0
        if _g1k0 > K0:
            _g1k0 = unsafe_sub(_g1k0, K0) + 1
        else:
            _g1k0 = unsafe_sub(K0, _g1k0) + 1

        # D / (A * N**N) * _g1k0**2 / gamma**2
        mul1: uint256 = unsafe_div(unsafe_div(unsafe_div(10**18 * D, gamma) * _g1k0, gamma) * _g1k0 * A_MULTIPLIER, ANN)

        # 2*K0 / _g1k0
        mul2: uint256 = unsafe_div(10**18 + (2 * 10**18) * K0, _g1k0)

        yfprime: uint256 = 10**18 * y + S * mul2 + mul1
        _dyfprime: uint256 = D * mul2
        if yfprime < _dyfprime:
            y = unsafe_div(y_prev, 2)
            continue
        else:
            yfprime = unsafe_sub(yfprime, _dyfprime)
        fprime: uint256 = yfprime / y

        # y -= f / f_prime;  y = (y * fprime - f) / fprime
        # y = (yfprime + 10**18 * D - 10**18 * S) // fprime + mul1 // fprime * (10**18 - K0) // K0
        y_minus: uint256 = mul1 / fprime
        y_plus: uint256 = (yfprime + 10**18 * D) / fprime + y_minus * 10**18 / K0
        y_minus += 10**18 * S / fprime

        if y_plus < y_minus:
            y = unsafe_div(y_prev, 2)
        else:
            y = unsafe_sub(y_plus, y_minus)

        # dev: reused `y_minus` as `diff`
        if y > y_prev:
            y_minus = unsafe_sub(y, y_prev)
        else:
            y_minus = unsafe_sub(y_prev, y)

        if y_minus < max(convergence_limit, unsafe_div(y, 10**14)):
            y_minus = unsafe_div(y * 10**18, D)  # dev: reused as `frac`
            assert (y_minus > 10**16 - 1) and (y_minus < 10**20 + 1)  # dev: unsafe value for y
            return [y, 0]

    raise "Did not converge"


@external
@view
def computeD(ANN: uint256, gamma: uint256, x0: uint256, x1: uint256, K0_prev: uint256 = 0) -> uint256:
    """
    Finding the invariant using Newton method.
    ANN is higher by the factor A_MULTIPLIER
    ANN is already A * N**N
    Currently uses 60k gas
    """
    # Safety checks
    # assert ANN > MIN_A - 1 and ANN < MAX_A + 1  # dev: unsafe values A
    # assert gamma > MIN_GAMMA - 1 and gamma < MAX_GAMMA + 1  # dev: unsafe values gamma

    # Initial value of invariant D is that for constant-product invariant
    x: uint256[N_COINS] = [x0, x1]
    if x[0] < x[1]:
        x = [x1, x0]

    assert x[0] > 10**9 - 1 and x[0] < 10**15 * 10**18 + 1  # dev: unsafe values x[0]
    assert unsafe_div(x[1] * 10**18, x[0]) > 10**14-1  # dev: unsafe values x[i] (input)

    #D: uint256 = N_COINS * self.geometric_mean(x0, x1)
    S: uint256 = unsafe_add(x[0], x[1])  # can unsafe add here because we checked x[0] bounds

    D: uint256 = 0
    if K0_prev == 0:
        D = N_COINS * isqrt(unsafe_mul(x[0], x[1]))
    else:
        # D = isqrt(x[0] * x[1] * 4 / K0_prev * 10**18)
        D = isqrt(unsafe_mul(unsafe_div(unsafe_mul(unsafe_mul(4, x[0]), x[1]), K0_prev), 10**18))
        if S < D:
            D = S

    __g1k0: uint256 = unsafe_add(gamma, 10**18)

    for i in range(255):
        D_prev: uint256 = D
        assert D > 0
        # Unsafe ivision by D is now safe

        # K0: uint256 = 10**18
        # for _x in x:
        #     K0 = K0 * _x * N_COINS / D
        # collapsed for 2 coins
        K0: uint256 = unsafe_div(unsafe_div((10**18 * N_COINS**2) * x[0], D) * x[1], D)

        _g1k0: uint256 = __g1k0
        if _g1k0 > K0:
            _g1k0 = unsafe_sub(_g1k0, K0) + 1  # > 0
        else:
            _g1k0 = unsafe_sub(K0, _g1k0) + 1  # > 0

        # D / (A * N**N) * _g1k0**2 / gamma**2
        mul1: uint256 = unsafe_div(unsafe_div(unsafe_div(10**18 * D, gamma) * _g1k0, gamma) * _g1k0 * A_MULTIPLIER, ANN)

        # 2*N*K0 / _g1k0
        mul2: uint256 = unsafe_div(((2 * 10**18) * N_COINS) * K0, _g1k0)

        neg_fprime: uint256 = (S + unsafe_div(S * mul2, 10**18)) + mul1 * N_COINS / K0 - unsafe_div(mul2 * D, 10**18)

        # D -= f / fprime; neg_fprime safediv being validated
        D_plus: uint256 = D * (neg_fprime + S) / neg_fprime
        D_minus: uint256 = unsafe_div(D * D,  neg_fprime)
        if 10**18 > K0:
            D_minus += unsafe_div(unsafe_div(D * unsafe_div(mul1, neg_fprime), 10**18) * unsafe_sub(10**18, K0), K0)
        else:
            D_minus -= unsafe_div(unsafe_div(D * unsafe_div(mul1, neg_fprime), 10**18) * unsafe_sub(K0, 10**18), K0)

        if D_plus > D_minus:
            D = unsafe_sub(D_plus, D_minus)
        else:
            D = unsafe_div(unsafe_sub(D_minus, D_plus), 2)

        # dev: reused `D_minus` as `diff`
        if D > D_prev:
            D_minus = unsafe_sub(D, D_prev)
        else:
            D_minus = unsafe_sub(D_prev, D)

        if D_minus * 10**14 < max(10**16, D):  # Could reduce precision for gas efficiency here
            # Test that we are safe with the next newton_y
            D_minus = unsafe_div(x0 * 10**18, D)  # dev: reused as `frac`
            assert (D_minus > 10**16 - 1) and (D_minus < 10**20 + 1)  # dev: unsafe values x[i]

            D_minus = unsafe_div(x1 * 10**18, D)
            assert (D_minus > 10**16 - 1) and (D_minus < 10**20 + 1)  # dev: unsafe values x[i]
            return D

    raise "Did not converge"


@internal
@pure
def geometric_mean(a: uint256, b: uint256) -> uint256:
    """
    (x[0] * x[1] * ...) ** (1/N)
    """
    return isqrt(a * b)


@external
@view
def getPriceOracle(last_timestamp: uint256, ma_time: uint256, last_prices: uint256, price_scale: uint256, price_oracle: uint256) -> uint256:
    alpha: uint256 = self.halfpow(
        unsafe_div(
            unsafe_mul(block.timestamp - last_timestamp, 10**18),
            ma_time
        )
    )

    return unsafe_div(
        min(last_prices, 2 * price_scale) * (10**18 - alpha) +
        (price_oracle * alpha),  # ^-------- Cap spot price into EMA.
        10**18
    )


@internal
@pure
def halfpow(power: uint256) -> uint256:
    """
    1e18 * 0.5 ** (power/1e18)

    Inspired by: https://github.com/balancer-labs/balancer-core/blob/master/contracts/BNum.sol#L128
    """
    intpow: uint256 = unsafe_div(power, 10**18)
    if intpow > 59:
        return 0
    otherpow: uint256 = unsafe_sub(power, unsafe_mul(intpow, 10**18))  # < 10**18
    # result: uint256 = unsafe_div(10**18, pow_mod256(2, intpow))
    result: uint256 = pow_mod256(2, intpow)
    result = unsafe_div(10**18, result)
    if otherpow == 0:
        return result

    term: uint256 = 10**18
    S: uint256 = 10**18
    neg: bool = False

    for i in range(1, 256):
        K: uint256 = unsafe_mul(i, 10**18)  # <= 255 * 10**18; >= 10**18
        c: uint256 = unsafe_sub(K, 10**18)  # <= 254 * 10**18; < K
        if otherpow > c:  # c < otherpow < 10**18 <= K -> c < K
            c = unsafe_sub(otherpow, c)
            neg = not neg
        else:
            c = unsafe_sub(c, otherpow)  # c < K
        # c <= 254 * 10**18, < K -> (c/2) / K < 1 -> term * c/2 / K <= 10**18
        term = unsafe_div(unsafe_mul(term, unsafe_div(c, 2)), K)
        if neg:
            S -= term
        else:
            S += term
        if term < EXP_PRECISION:
            return unsafe_div(result * S, 10**18)

    raise "Did not converge"


@external
@view
def getLastPrices(
    _A: uint256, _gamma: uint256, _xp0: uint256, _xp1: uint256, _D: uint256, _price_scale: uint256
) -> uint256:
    return unsafe_div(
        self.getPrice(_A, _gamma, _xp0, _xp1, _D) * _price_scale,
        10**18
    )


@internal
@view
def getPrice(
    _A: uint256, _gamma: uint256, _xp0: uint256, _xp1: uint256, _D: uint256
) -> uint256:
    """
    @notice Calculates dx/dy.
    @dev Output needs to be multiplied with price_scale to get the actual value.
    @param _A Amplification coefficient of the pool.
    @param _gamma Gamma of the pool.
    @param _xp0 Balances of the pool.
    @param _xp1 Balances of the pool.
    @param _D Current value of D.
    """

    assert _D > 10**17 - 1 and _D < 10**15 * 10**18 + 1  # dev: unsafe D values

    # K0 = P * N**N / D**N.
    # K0 is dimensionless and has 10**36 precision:
    K0: uint256 = unsafe_div(
        unsafe_div(4 * _xp0 * _xp1, _D) * 10**36,
        _D
    )

    # GK0 is in 10**36 precision and is dimensionless.
    # GK0 = (
    #     2 * _K0 * _K0 / 10**36 * _K0 / 10**36
    #     + (gamma + 10**18)**2
    #     - (_K0 * _K0 / 10**36 * (2 * gamma + 3 * 10**18) / 10**18)
    # )
    # GK0 is always positive. So the following should never revert:
    GK0: uint256 = (
        unsafe_div(unsafe_div(2 * K0 * K0, 10**36) * K0, 10**36)
        + pow_mod256(unsafe_add(_gamma, 10**18), 2)
        - unsafe_div(
            unsafe_div(pow_mod256(K0, 2), 10**36) * unsafe_add(unsafe_mul(2, _gamma), 3 * 10**18),
            10**18
        )
    )

    # NNAG2 = N**N * A * gamma**2
    NNAG2: uint256 = unsafe_div(unsafe_mul(_A, pow_mod256(_gamma, 2)), A_MULTIPLIER)

    # denominator = (GK0 + NNAG2 * x / D * _K0 / 10**36)
    denominator: uint256 = (GK0 + unsafe_div(unsafe_div(NNAG2 * _xp0, _D) * K0, 10**36) )

    # p_xy = x * (GK0 + NNAG2 * y / D * K0 / 10**36) / y * 10**18 / denominator
    # p is in 10**18 precision.
    return unsafe_div(
        _xp0 * ( GK0 + unsafe_div(unsafe_div(NNAG2 * _xp1, _D) * K0, 10**36) ) / _xp1 * 10**18,
        denominator
    )

4. Swap function details

function swap(
    bytes memory _data,
    address _sender,
    address _callback,
    bytes memory _callbackData
) public override nonReentrant returns (TokenAmount memory _tokenAmount) {
    ICallback.BaseSwapCallbackParams memory params;

    (params.tokenIn, params.to, params.withdrawMode) = abi.decode(
        _data,
        (address, address, uint8)
    );
    (params.reserve0, params.reserve1) = getReserves();
    (params.balance0, params.balance1) = _balances();

    // Gets swap fee for the sender.
    _sender = _getVerifiedSender(_sender);

    // ----------- Update invariant if A, gamma are undergoing ramps ---------

    // In case ramp is happening
    (uint a, uint gamma, uint futureTime) = _params();

    // ----------------------- Calculate dy and fees --------------------------

    uint xp0;
    uint xp1;
    uint _k0_prev;
    uint _feeAmountOut;
    ICryptoMath _MATH = MATH;

    if (params.tokenIn == token0) {
        (
            params.amountOut,
            params.swapFee,
            _feeAmountOut,
            xp0,
            xp1,
            _k0_prev
        ) = _getAmountOut(
            GetAmountOutInputParams({
                sender: _sender,
                a: a,
                gamma: gamma,
                balance0: params.balance0,
                balance1: params.balance1,
                reserve0: params.reserve0,
                reserve1: params.reserve1,
                token0In: true,
                futureTime: futureTime
            }),
            _MATH
        );

        params.tokenOut = token1;
        params.balance1 -= params.amountOut;
        params.amountIn = params.balance0 - params.reserve0;

        emit Swap(
            msg.sender,
            params.amountIn,
            0,
            0,
            params.amountOut,
            params.to
        );
        emit Fee(0, _feeAmountOut);
    } else {
        require(params.tokenIn == token1);
        (
            params.amountOut,
            params.swapFee,
            _feeAmountOut,
            xp0,
            xp1,
            _k0_prev
        ) = _getAmountOut(
            GetAmountOutInputParams({
                sender: _sender,
                a: a,
                gamma: gamma,
                balance0: params.balance0,
                balance1: params.balance1,
                reserve0: params.reserve0,
                reserve1: params.reserve1,
                token0In: false,
                futureTime: futureTime
            }),
            _MATH
        );

        params.tokenOut = token0;
        params.balance0 -= params.amountOut;
        params.amountIn = params.balance1 - params.reserve1;

        emit Swap(
            msg.sender,
            0,
            params.amountIn,
            params.amountOut,
            0,
            params.to
        );
        emit Fee(_feeAmountOut, 0);
    }

    require(params.amountIn != 0 && params.amountOut != 0);

    // ------------------------ Tweak `priceScale` ----------------------------

    _tweakPrice(a, gamma, xp0, xp1, 0, futureTime, _k0_prev, _MATH);

    // Updates reserves with up-to-date balances (updated above).
    _updateReserves(params.balance0, params.balance1);

    // Transfers output tokens.
    _transferTokens(
        params.tokenOut,
        params.to,
        params.amountOut,
        params.withdrawMode
    );

    // Calls callback with data.
    if (_callback != address(0)) {
        // Fills additional values for callback params.
        params.sender = _sender;
        params.callbackData = _callbackData;

        // Calculates fee amount for callback.
        params.feeIn = Math.mulDivUnsafeFirstLast(
            params.amountIn,
            params.swapFee,
            1e5
        );

        ICallback(_callback).syncSwapBaseSwapCallback(params);
    }

    _tokenAmount.token = params.tokenOut;
    _tokenAmount.amount = params.amountOut;

    emit Swapped(
        msg.sender,
        _sender,
        params.tokenOut,
        params.amountIn,
        params.amountOut,
        params.swapFee,
        params.to
    );
}