balancer_adapter.py

"""
balancer_adapter.py — Balancer V2 WeightedPool 的 ProtocolAdapter 实现

将现有 vault_client.py 中的六步流程封装为标准 Adapter 接口，
使 Balancer V2 可与其他协议在统一框架下运行。
"""
from typing import Dict, List, Optional

from web3 import Web3

from .core import (
    AMMType,
    AccuracyReport,
    AMMConstraint,
    PoolMeta,
    ProtocolAdapter,
    RawState,
    SimulationDifficulty,
    SwapParams,
)
from .models import ONE, BalancerPoolMeta, BalancerRawState
from .weighted_math import BalancerWeightedConstraint
from .vault_client import (
    fetch_raw_state,
    validate_and_convert,
    simulate_swap_onchain,
    scan_all_pools,
)


class BalancerV2Adapter(ProtocolAdapter):
    """Balancer V2 WeightedPool 适配器。"""

    @property
    def protocol_name(self) -> str:
        return "BalancerV2Weighted"

    @property
    def amm_type(self) -> AMMType:
        return AMMType.CONSTANT_FUNCTION

    @property
    def simulation_difficulty(self) -> SimulationDifficulty:
        return SimulationDifficulty.TRIVIAL

    def discover_pools(
        self,
        w3: Web3,
        from_block: int,
        to_block: Optional[int] = None,
    ) -> List[PoolMeta]:
        return scan_all_pools(w3, from_block, to_block)

    def fetch_state(
        self,
        pool_meta: PoolMeta,
        w3: Web3,
        block_number: Optional[int] = None,
    ) -> Optional[RawState]:
        if not isinstance(pool_meta, BalancerPoolMeta):
            return None
        return fetch_raw_state(pool_meta, w3, block_number)

    def build_constraint(
        self,
        raw_state: RawState,
        pool_meta: PoolMeta,
    ) -> Optional[AMMConstraint]:
        if not isinstance(raw_state, BalancerRawState):
            return None
        if not isinstance(pool_meta, BalancerPoolMeta):
            return None
        return validate_and_convert(raw_state, pool_meta)

    def quote_onchain(
        self,
        pool_meta: PoolMeta,
        w3: Web3,
        params: SwapParams,
        block_number: int,
    ) -> Optional[int]:
        if not isinstance(pool_meta, BalancerPoolMeta):
            return None
        return simulate_swap_onchain(
            pool_meta, w3,
            params.amount_in,
            params.token_in_index,
            params.token_out_index,
            block_number,
        )

    def _generate_test_amounts(
        self,
        constraint: AMMConstraint,
        token_in_index: int,
        num_points: int,
    ) -> List[int]:
        if not isinstance(constraint, BalancerWeightedConstraint):
            return []
        B_in = constraint.balances[token_in_index]
        low = max(1, B_in // 10_000)
        high = B_in // 33
        if high <= low:
            return []
        return [
            int(low * ((high / low) ** (i / max(1, num_points - 1))))
            for i in range(num_points)
        ]