Balancer precision calculation test example

balancer_example.py

from wake.testing import *
from wake.testing.fuzzing import *
from pytypes.contracts.ComposableStablePoolFactory import ComposableStablePoolFactory
from pytypes.contracts.ComposableStablePool import ComposableStablePool
from pytypes.wake.interfaces.IERC20 import IERC20

from pytypes.__.__.node_modules.balancerlabs.v2interfaces.contracts.poolutils.IPoolVersion import IPoolVersion
from pytypes.__.__.node_modules.balancerlabs.v2interfaces.contracts.poolutils.IRateProvider import IRateProvider
from pytypes.__.__.node_modules.balancerlabs.v2interfaces.contracts.solidityutils.helpers.IVersion import IVersion
from pytypes.__.__.node_modules.balancerlabs.v2interfaces.contracts.solidityutils.openzeppelin.IERC20 import IERC20
from pytypes.__.__.node_modules.balancerlabs.v2interfaces.contracts.standaloneutils.IProtocolFeePercentagesProvider import IProtocolFeePercentagesProvider
from pytypes.__.__.node_modules.balancerlabs.v2interfaces.contracts.vault.IVault import IVault
from pytypes.__.__.node_modules.balancerlabs.v2poolutils.contracts.test.MockVault import MockVault
from pytypes.tests.MockAuthorizor import MockAuthorizor
from pytypes.__.__.node_modules.balancerlabs.v2poolutils.contracts.test.MockRateProvider import MockRateProvider
from pytypes.tests.MockProtocolFeePercentagesProvider import MockProtocolFeePercentagesProvider
from pytypes.__.__.node_modules.balancerlabs.v2interfaces.contracts.vault.IPoolSwapStructs import IPoolSwapStructs
from decimal import Decimal, getcontext

import logging

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)


# ============================================================================
# StableMath Python Implementation - Matches Solidity exactly
# ============================================================================

AMP_PRECISION = 1000


class BalanceGraph:
    """
    Lightweight ASCII balance graph for quick inspection during fuzzing.
    """

    def __init__(self, width: int = 40):
        self._width = width

    @staticmethod
    def _label(token: IERC20) -> str:
        if hasattr(token, "label") and token.label:
            return token.label
        try:
            return str(token.address)
        except Exception:
            return "token"

    def render(
        self,
        tokens: list[IERC20],
        balances: dict[IERC20, int],
        invariant: int,
        iteration: int,
    ) -> None:
        if not balances:
            return

        max_value = max(balances[token] for token in tokens if token in balances)
        max_value = max_value if max_value > 0 else 1

        lines = [
            "",
            f"Balances after flow #{iteration} (Invariant={invariant})",
        ]

        for token in tokens:
            if token not in balances:
                continue
            value = balances[token]
            bar_length = max(1, value * self._width // max_value)
            bar = "#" * bar_length
            lines.append(f"{self._label(token):<10} | {bar:<{self._width}} {value}")

        print("\n".join(lines))


def calc_invariant(amplification_parameter: int, balances: list[int]) -> int:
    """
    Calculate the StableSwap invariant using Newton-Raphson approximation.

    Matches: StableMath._calculateInvariant

    Formula:
        D = invariant
        A = amplification coefficient
        A * n^n * S + D = A * D * n^n + D^(n+1) / (n^n * P)
    """
    total = sum(balances)
    if total == 0:
        return 0

    num_tokens = len(balances)
    prev_invariant = 0
    invariant = total
    amp_times_total = amplification_parameter * num_tokens

    # Newton-Raphson iteration (max 255 iterations)
    for _ in range(255):
        D_P = invariant

        # Calculate D_P = D^(n+1) / (n^n * P)
        for balance in balances:
            D_P = (D_P * invariant) // (balance * num_tokens)

        prev_invariant = invariant

        # numerator = ((ampTimesTotal * sum) / AMP_PRECISION + D_P * numTokens) * invariant
        numerator_part1 = (amp_times_total * total) // AMP_PRECISION
        numerator_part2 = D_P * num_tokens
        numerator = (numerator_part1 + numerator_part2) * invariant

        # denominator = ((ampTimesTotal - AMP_PRECISION) * invariant) / AMP_PRECISION + (numTokens + 1) * D_P
        denominator_part1 = ((amp_times_total - AMP_PRECISION) * invariant) // AMP_PRECISION
        denominator_part2 = (num_tokens + 1) * D_P
        denominator = denominator_part1 + denominator_part2

        invariant = numerator // denominator
        # Check for convergence (difference <= 1)
        if invariant > prev_invariant:
            if invariant - prev_invariant <= 1:
                return invariant
        elif prev_invariant - invariant <= 1:
            return invariant

    raise Exception("STABLE_INVARIANT_DIDNT_CONVERGE")

def calc_invariant_pure_math(amplification_parameter: int, balances: list[int]) -> float:
    """
    Calculate StableSwap invariant using the pure mathematical formula.

    Starting from: A * n^n * S + D = A * D * n^n + D^(n+1) / (n^n * P)

    Newton-Raphson iteration:
    D_new = (A * n^n * S + D_P * n) * D / ((A * n^n - 1) * D + (n+1) * D_P)

    where D_P = D^(n+1) / (n^n * P)
    """
    from decimal import Decimal, getcontext
    getcontext().prec = 1000

    S = sum(balances)
    if S == 0:
        return 0.0

    n = len(balances)
    A = Decimal(amplification_parameter) / Decimal(AMP_PRECISION)  # Actual A value
    S_dec = Decimal(S)
    n_dec = Decimal(n)
    balances_dec = [Decimal(b) for b in balances]

    # Initial guess
    D = Decimal(S)
    Ann = A * (n_dec)  # A * n^n

    for iteration in range(255):
        # Calculate D_P = D^(n+1) / (n^n * P)
        D_P = D
        for balance in balances_dec:
            D_P = D_P * D / (balance * n_dec)

        D_prev = D

        # Newton-Raphson formula:
        # numerator = (A * n^n * S + D_P * n) * D
        numerator = (Ann * S_dec + D_P * n_dec) * D

        # denominator = (A * n^n - 1) * D + (n+1) * D_P
        denominator = (Ann - 1) * D + (n_dec + 1) * D_P

        D = numerator / denominator

        delta = abs(D - D_prev)

        if delta <= 1:
            return float(D)

    raise Exception("STABLE_INVARIANT_DIDNT_CONVERGE")

def get_token_balance_pure_math_quiet(
    amplification_parameter: int,
    balances: list[int],
    invariant: float,
    token_index: int
) -> float:
    """
    Same as get_token_balance_pure_math but without print statements.

    Solves: y^2 + y * b - c = 0
    Where:
        b = S' + D / (A * n^n) - D
        c = D^(n+1) / (A * n^2n * P')
    """
    from decimal import Decimal, getcontext
    getcontext().prec = 1000

    n = len(balances)

    # Convert to Decimal
    A = Decimal(amplification_parameter) / Decimal(AMP_PRECISION)
    D = Decimal(invariant)
    n_dec = Decimal(n)
    balances_dec = [Decimal(b) for b in balances]

    # A * n^n
    Ann = A * (n_dec)

    # S' = sum excluding target
    S_prime = sum(balances_dec[j] for j in range(n) if j != token_index)

    # Calculate c = D^(n+1) / (A * n^2n * P')
    c = D
    for _ in range(n):
        c = c * D / n_dec

    for j in range(n):
        if j != token_index:
            c = c / balances_dec[j]

    c = c / Ann

    # Calculate b
    b = S_prime + D / Ann - D

    # Solve quadratic: y = (-b + sqrt(b^2 + 4c)) / 2
    discriminant = b * b + 4 * c
    y = (-b + discriminant.sqrt()) / 2

    return float(y)

def get_token_balance_given_invariant_and_all_other_balances(
    amplification_parameter: int,
    balances: list[int],
    invariant: int,
    token_index: int
) -> int:
    """
    Calculate the balance of a specific token given the invariant and all other balances.

    Matches: StableMath._getTokenBalanceGivenInvariantAndAllOtherBalances
    """
    num_tokens = len(balances)
    amp_times_total = amplification_parameter * num_tokens

    # Calculate S' (sum excluding target) and P_D (product term)
    total = balances[0]
    P_D = balances[0] * num_tokens

    for j in range(1, num_tokens):
        # Accumulate product: (P_D * balance_j * n) / D
        P_D = (P_D * balances[j] * num_tokens) // invariant
        total += balances[j]

    # Remove the target token from sum to get S'
    total = total - balances[token_index]

    # print("sum", total)

    # Calculate coefficient 'c' in the quadratic equation
    inv2 = invariant * invariant
    # print(f"inv2 {inv2}")

    # IMPORTANT: Must match Solidity's order of operations!
    # Solidity: divUp(inv2, ampTimesTotal * P_D) * AMP_PRECISION * balances[token_index]
    c_denominator = amp_times_total * P_D
    temp = (inv2 + c_denominator - 1) // c_denominator  # divUp first
    c = temp * AMP_PRECISION * balances[token_index]  # multiply after
    # print(f"c: {c}")

    # Calculate coefficient 'b' in the quadratic equation
    # IMPORTANT: Must match Solidity's order of operations!
    # Solidity: sum + divDown(invariant, ampTimesTotal) * AMP_PRECISION
    b = total + ((invariant // amp_times_total) * AMP_PRECISION)
    # print(f"b: {b}")

    # Solve using Newton-Raphson iteration
    # Initial approximation: y_0 = divUp(D^2 + c, D + b)
    numerator = inv2 + c
    denominator = invariant + b
    token_balance = (numerator + denominator - 1) // denominator
    # print("first balance", token_balance)

    # Iterate until convergence (max 255 iterations)
    for _ in range(255):
        prev_token_balance = token_balance

        # Newton-Raphson formula: y_(k+1) = divUp(y_k^2 + c, 2*y_k + b - D)
        numerator = token_balance * token_balance + c
        denominator = token_balance * 2 + b - invariant
        token_balance = (numerator + denominator - 1) // denominator

        # Check for convergence (within 1 wei)
        if token_balance > prev_token_balance:
            if token_balance - prev_token_balance <= 1:
                return token_balance
        elif prev_token_balance - token_balance <= 1:
            return token_balance

    raise Exception("STABLE_GET_BALANCE_DIDNT_CONVERGE")


def calc_out_given_in(
    amplification_parameter: int,
    balances: list[int],
    token_index_in: int,
    token_index_out: int,
    token_amount_in: int,
    invariant: int
) -> int:
    """
    Calculate how many tokens can be received for a given amount of tokens sent.

    Matches: StableMath._calcOutGivenIn
    """
    # Add the incoming tokens to the balance
    balances[token_index_in] = balances[token_index_in] + token_amount_in

    # Calculate what the new balance of tokenOut should be
    final_balance_out = get_token_balance_given_invariant_and_all_other_balances(
        amplification_parameter,
        balances,
        invariant,
        token_index_out
    )
    # print("final_balance_out", final_balance_out)
    pure_math_final_balance_out = get_token_balance_pure_math_quiet(
        amplification_parameter,
        balances,
        invariant,
        token_index_out
    )


    diff = final_balance_out - pure_math_final_balance_out
    assert -2 <= diff # vault does not lose value as rounding
    rel_error = diff / pure_math_final_balance_out if pure_math_final_balance_out != 0 else 0.0
    assert abs(rel_error) < 1e-9, f"Token balance mismatch too high: abs={diff}, rel={rel_error}"

    # Restore the original balance
    balances[token_index_in] = balances[token_index_in] - token_amount_in

    # Amount out is the difference (minus 1 for rounding)
    amount_out = balances[token_index_out] - final_balance_out - 1
    # print("amount_out calculated", amount_out)

    return amount_out

def calc_in_given_out(
    amplification_parameter: int,
    balances: list[int],
    token_index_in: int,
    token_index_out: int,
    token_amount_out: int,
    invariant: int
) -> int:
    """
    Calculate how many tokens must be sent to receive a given amount of tokens.

    Matches: StableMath._calcInGivenOut
    """
    # Subtract the outgoing tokens from the balance
    balances[token_index_out] = balances[token_index_out] - token_amount_out

    final_balance_in = get_token_balance_given_invariant_and_all_other_balances(
        amplification_parameter,
        balances,
        invariant,
        token_index_in
    )
    # print("final_balance_in", final_balance_in)

    pure_math_final_balance_out = get_token_balance_pure_math_quiet(
        amplification_parameter,
        balances,
        invariant,
        token_index_in
    )

    diff = final_balance_in - pure_math_final_balance_out # diff should be at least bigger
    assert -2 <= diff # vault does not lose value as rounding
    rel_error = diff / pure_math_final_balance_out if pure_math_final_balance_out != 0 else 0.0
    assert abs(rel_error) < 1e-9, f"Token balance mismatch too high: abs={diff}, rel={rel_error}"

    # Restore the original balance
    balances[token_index_out] = balances[token_index_out] + token_amount_out

    # Amount in is the difference (plus 1 for rounding)
    amount_in = final_balance_in - balances[token_index_in] + 1
    # print("amount_in calculated", amount_in)

    return amount_in


# ============================================================================
# Test Helper Functions
# ============================================================================

# Print failing tx call trace
def revert_handler(e: RevertError):
    if e.tx is not None:
        print(e.tx.call_trace)


class FuzzSwap(FuzzTest):

    balances_dict: dict[IERC20, int]

    def get_balance_array(self) -> list[int]:
        return [self.balances_dict[token] for token in self.token_list]


    def pre_sequence(self):
        self.something = 0
        self.weth = IERC20(0xC02AAa39b223FE8D0A0e5C4F27eAD9083C756Cc2)
        self.weth.label = "weth"
        self.oseth = IERC20(0xf1C9ACDc66974DFB6dECB12AA385b9cd01190e38)
        self.oseth.label = "oseth"

        authorizor = MockAuthorizor.deploy()
        self.mock_vault = MockVault.deploy(authorizor)

        mock_protocol_fee_percentages_provider = MockProtocolFeePercentagesProvider.deploy()
        protocol_fee_percentages_provider = IProtocolFeePercentagesProvider(
            mock_protocol_fee_percentages_provider
        )
        min_value = max_value = 10**15
        mock_protocol_fee_percentages_provider.registerFeeType(0, "something", min_value, max_value)
        mock_protocol_fee_percentages_provider.registerFeeType(2, "something", min_value, max_value)
        mock_protocol_fee_percentages_provider.registerFeeType(3, "something", min_value, max_value)

        mock_rate_provider = MockRateProvider.deploy()
        second_mock_rate_provider = MockRateProvider.deploy()
        second_mock_rate_provider.mockRate(123_4567_8999_9999_9999)

        self.swap_fee_percentage = 10**12
        self.amplification_parameter = 10 * 10**3
        amplification_parameter_for_deployment_argument = self.amplification_parameter // 10**3

        self.pool_factory = ComposableStablePoolFactory.deploy(
            vault=self.mock_vault,
            protocolFeeProvider=protocol_fee_percentages_provider,
            factoryVersion="something",
            poolVersion="something",
            initialPauseWindowDuration=0,
            bufferPeriodDuration=0,
        )

        tx = self.pool_factory.create(
            "something",
            "SOMETHING",
            tokens=[self.weth, self.oseth],
            amplificationParameter=amplification_parameter_for_deployment_argument,
            rateProviders=[mock_rate_provider, second_mock_rate_provider],
            tokenRateCacheDurations=[100, 1000],
            exemptFromYieldProtocolFeeFlags=[False, False],
            swapFeePercentage=self.swap_fee_percentage,
            owner=chain.accounts[0],
            salt=random_bytes(32),
        )

        self.pool = tx.return_value
        self.pool.label = "bpt"
        self.pool_id = self.pool.getPoolId()

        token_list, _ = self.mock_vault.getPoolTokens(self.pool_id)
        self.token_list = list(token_list)
        self.bpt_index = self.token_list.index(self.pool)

        self.scaling_factor_dict: dict[IERC20, int] = {
            self.pool: 10**18,
            self.weth: 10**18,
            self.oseth: 123_4567_8999_9999_9999,
        }

        self.balances_dict = {
            self.weth: 0,
            self.oseth: 0,
            self.pool: 0,
        }

        self.pool.onJoinPool(
            poolId=self.pool_id,
            sender=chain.accounts[3],
            balances=[0, 0, 0],
            lastChangeBlock=chain.blocks["latest"].number,
            protocolSwapFeePercentage=2,
            userData=abi.encode(uint8(0), [uint256(10**20), uint256(10**20), uint256(10**20)]),
            recipient=chain.accounts[3],
            from_=self.mock_vault,
        )

        self.balances_dict = {
            self.weth: 10**20,
            self.oseth: 10**20,
            self.pool: 10**20,
        }
        self.balance_graph = BalanceGraph()
        self._flow_counter = 0
        self._last_invariant_value = 0

    def _token_index_without_bpt(self, token: IERC20) -> int:
        index = self.token_list.index(token)
        return index - 1 if self.bpt_index < index else index

    @staticmethod
    def _fixed_point_div_up(amount: int, divisor: int) -> int:
        if amount == 0:
            return 0
        inflated = amount * 10**18
        return ((inflated - 1) // divisor) + 1

    @flow()
    def flow_swap(self):
        swap_amount = random_int(0, 10**18)
        swap_kind = random.choice([IVault.SwapKind.GIVEN_IN, IVault.SwapKind.GIVEN_OUT])
        token_in = random.choice([self.weth, self.oseth])
        token_out = self.oseth if token_in == self.weth else self.weth

        swap_request = IPoolSwapStructs.SwapRequest(
            kind=swap_kind,
            tokenIn=token_in,
            tokenOut=token_out,
            amount=swap_amount,
            poolId=self.pool_id,
            lastChangeBlock=chain.blocks["latest"].number,
            from_=chain.accounts[3],
            to=chain.accounts[3],
            userData=b"",
        )

        balances = self.get_balance_array()

        with may_revert() as e:
            tx = self.mock_vault.callGeneralPoolSwap(
                pool=self.pool,
                request=swap_request,
                balances=balances,
                indexIn=self.token_list.index(swap_request.tokenIn),
                indexOut=self.token_list.index(swap_request.tokenOut),
            )
        if e.value is not None:
            if e.value == Error("BAL#001"):
                return "balance limit"
            if e.value == Error("BAL#321"):
                return "invariant did not converge"
            else:
                raise e.value

        swap_event = next(event for event in tx.events if isinstance(event, MockVault.Swap))

        effective_amount = swap_amount
        if swap_kind == IVault.SwapKind.GIVEN_IN:
            assert swap_amount == swap_event.amountIn
            fee = (swap_amount * self.swap_fee_percentage + 10**18 - 1) // 10**18
            effective_amount = swap_amount - fee
            self.balances_dict[swap_request.tokenIn] += swap_event.amountIn
            self.balances_dict[swap_request.tokenOut] -= swap_event.amountOut
        else:
            assert swap_amount == swap_event.amountOut
            self.balances_dict[swap_request.tokenIn] += swap_event.amountIn
            self.balances_dict[swap_request.tokenOut] -= swap_event.amountOut

        upscaled_balances = [
            self.scaling_factor_dict[token] * balances[i] // 10**18
            for i, token in enumerate(self.token_list)
        ]
        scaling_factor_for_given = (
            self.scaling_factor_dict[swap_request.tokenIn]
            if swap_kind == IVault.SwapKind.GIVEN_IN
            else self.scaling_factor_dict[swap_request.tokenOut]
        )
        scaled_swap_amount = scaling_factor_for_given * effective_amount // 10**18

        balances_for_calc = [
            upscaled_balances[i] for i, token in enumerate(self.token_list) if token != self.pool
        ]

        index_in_for_calc = self._token_index_without_bpt(swap_request.tokenIn)
        index_out_for_calc = self._token_index_without_bpt(swap_request.tokenOut)

        invariant = calc_invariant(self.amplification_parameter, balances_for_calc)
        float_invariant = calc_invariant_pure_math(self.amplification_parameter, balances_for_calc)
        # invariant should be smaller than float value by roudning
        diff = float_invariant - invariant # should be smaller invariant.
        rel_error = diff / invariant if invariant != 0 else 0.0
        assert rel_error < 1e-9, f"Invariant mismatch too high: abs={diff}, rel={rel_error}"

        if swap_kind == IVault.SwapKind.GIVEN_IN:
            amount_out = calc_out_given_in(
                amplification_parameter=self.amplification_parameter,
                balances=balances_for_calc.copy(),
                token_index_in=index_in_for_calc,
                token_index_out=index_out_for_calc,
                token_amount_in=scaled_swap_amount,
                invariant=invariant,
            )

            # down scale down
            amount_out = amount_out * 10**18 // self.scaling_factor_dict[swap_request.tokenOut]
            assert amount_out == swap_event.amountOut, f"Mismatch: Python={amount_out}, Solidity={swap_event.amountOut}"


        else:
            amount_in = calc_in_given_out(
                amplification_parameter=self.amplification_parameter,
                balances=balances_for_calc.copy(),
                token_index_in=index_in_for_calc,
                token_index_out=index_out_for_calc,
                token_amount_out=scaled_swap_amount,
                invariant=invariant,
            )
            amount_in = self._fixed_point_div_up(amount_in, self.scaling_factor_dict[swap_request.tokenIn])
            complement_percentage = 10**18 - self.swap_fee_percentage
            amount_in = self._fixed_point_div_up(amount_in, complement_percentage)
            assert amount_in == swap_event.amountIn, f"Mismatch: Python={amount_in}, Solidity={swap_event.amountIn}"

        self._last_invariant_value = invariant
        self._flow_counter += 1

        logger.info(
            f"Flow #{self._flow_counter}: Swap {swap_kind.name} {swap_amount} "
            f"{swap_request.tokenIn.label} -> {swap_request.tokenOut.label}"
        )

    @invariant()
    def invariant_draw_balance_graph(self):
        if not hasattr(self, "balance_graph"):
            return
        if self._flow_counter == 0:
            return
        self.balance_graph.render(
            tokens=self.token_list,
            balances=self.balances_dict,
            invariant=self._last_invariant_value,
            iteration=self._flow_counter,
        )



@chain.connect(fork="http://localhost:8685@23717398")
@on_revert(revert_handler)
def test_fuzz():
    FuzzSwap.run(sequences_count=1, flows_count=100000)