Rewrite It in Rust: Pool State & Path Solver

Part V: The Good Stuff

This entry builds on the theoretical concepts explored in the Numerical Optimization series. If you didn’t read those, shame on you! But for your convenience, here they are again:

Numerical Optimization — Part I: Introduction

Numerical Optimization — Part II: Derivative-Free Methods

Numerical Optimization — Part III: Derivative-Free Methods (Multivariate)

Numerical Optimization — Part IV: Newtonian Methods

Numerical Optimization — Part V: Closed-Form Solutions

Liquidity Pool Structure

A given liquidity pool can be expressed as a tuple of immutable identity values and mutable state values.

For example, a Uniswap V2 pool consists of an identity:

pub struct V2PoolIdentity {
    pub address: Address,
    pub token0: Address,
    pub token1: Address,
    pub fee_token0: (u64, u64),
    pub fee_token1: (u64, u64),
    pub factory: Address,
    pub deployer: Address,
    pub init_hash: B256,
    pub variant: DexVariant,
    pub stable_swap: bool,
    pub fee_denominator: Option<u64>,
}

and a state:

pub struct V2PoolState {
    pub reserve0: U112,
    pub reserve1: U112,
    pub update_block: u64,
    /// a bounded-length journal of block states to support 
    /// restoring a previous state following a reorg
    pub journal: ReorgJournal<V2BlockDelta>,
}

Each pool is tracked by a Bot instance (see Part II), and the state is updated whenever the appropriate event listener observes an event that modifies this mutable state, and instructs the bot to apply the state transition function to the pool.

Pool Calculations

When we want to perform calculations against this state, we will ask the bot to evaluate the request against that pool’s state. Under the hood, the state is collapsed into a simplified data structure that arranges the reserves based on the input and output direction:

pub struct IntHopState {
    pub reserve_in: U256,
    pub reserve_out: U256,

    /// Gamma numerator: the retained fraction (e.g. 997 for 0.3% fee).
    pub gamma_numer: U256,
    
    /// Fee denominator (e.g. 1000 for 0.3% fee).
    pub fee_denom: U256,
}

This struct has an impl block for a swap method that replicates the Uniswap V2 formula:

pub fn swap(&self, x: U256) -> U256 {
    let amount_in_with_fee = x
        .checked_mul(self.gamma_numer)
        .ok_or(HopSwapError::Overflow)?;
    let numerator = amount_in_with_fee
        .checked_mul(self.reserve_out)
        .ok_or(HopSwapError::Overflow)?;
    let denom = self
        .reserve_in
        .checked_mul(self.fee_denom)
        .ok_or(HopSwapError::Overflow)?
        .checked_add(amount_in_with_fee)
        .ok_or(HopSwapError::Overflow)?;

    if denom.is_zero() {
        return Ok(U256::ZERO);
    }

    // Floor division — EVM `DIV` semantics.
    Ok(numerator / denom)
}

The user-facing bot wraps this up with a generalized function that can perform a swap calculation against different pool types defined in the PoolEntry enum:

pub enum PoolEntry {
    V2(V2PoolIdentity, V2PoolState),
    V3(V3PoolIdentity, V3PoolState),
    V4(V4PoolIdentity, V4PoolState),
    Curve(CurvePoolIdentity, CurvePoolState),
    BalancerWeighted(BalancerWeightedPoolIdentity, BalancerWeightedPoolState),
    BalancerStable(BalancerStablePoolIdentity, BalancerStablePoolState),
    AerodromeV2(AerodromeV2PoolIdentity, AerodromeV2PoolState),
}

pub fn simulate_swap(
    entry: &PoolEntry,
    zero_for_one: bool,
    amount_in: U256,
) -> Result<U256, SimulateSwapError> {
    match entry {
        PoolEntry::V2(identity, state) => {
            if amount_in.is_zero() {
                return Ok(U256::ZERO);
            }
            let (reserve_in, reserve_out, gamma_numer, fee_denom) = if zero_for_one {
                (
                    state.reserve0.to::<U256>(),
                    state.reserve1.to::<U256>(),
                    identity.fee_token0.0,
                    identity.fee_token0.1,
                )
            } else {
                (
                    state.reserve1.to::<U256>(),
                    state.reserve0.to::<U256>(),
                    identity.fee_token1.0,
                    identity.fee_token1.1,
                )
            };
            let hop = IntHopState::new(reserve_in, reserve_out, gamma_numer, fee_denom);
            hop.swap(amount_in)
                .map_err(|_| SimulateSwapError::NotComputable)
        }
        // V3 concentrated-liquidity math.
        PoolEntry::V3(identity, state) => {
             ...
        }
        // V4 concentrated-liquidity math.
        PoolEntry::V4(identity, state) => {
            ...
        }
        // not yet ported
        PoolEntry::Curve(..)
        | PoolEntry::BalancerWeighted(..)
        | PoolEntry::BalancerStable(..)
        | PoolEntry::AerodromeV2(..) => Ok(U256::ZERO),
    }
}

Path Solver

With this type-erased structure, swap calculations can be chained together as part of an arbitrage path and optimized by a dedicated solver. The solver does not care about the identity of the pool, its tokens, what position they are in, or the contents of the rollback journal. It only cares about the reserves associated with the token going in and out.

I have created a solver called UniswapEngine that handles Uniswap-style paths using these intermediate hops. It holds a reference to a Bot, which in turn holds the individual pools.

It implements several solver functions that operate on different kinds of paths. For example, a V2-only solver based on the closed-form Möbius solution is very simple. It takes an IntHopState array generated from the pools, then evaluates its 2x2 matrix: