tycho-execution 0.383.0

use num_bigint::BigUint;

use super::{
    constants::{PRICE_LEVEL_STREAM_PREFIX, PROPAMM_FALLBACK_PREFIX},
    group_swaps::group_swaps,
};
use crate::encoding::models::{Solution, Strategy, UserTransferType};

/// Default gas usage for an ERC-20 `transferFrom` or `transfer`. Used as fallback when the token
/// has no measured gas usage.
pub const DEFAULT_TOKEN_TRANSFER_GAS: u64 = 40_000;

/// Gas usage for an ERC-20 `approve`.
pub const TOKEN_APPROVAL_GAS: u64 = 25_000;

/// Gas usage for a permit2 transfer from the user
pub const USER_PERMIT2_TRANSFER: u64 = 80_000;

/// Callback-based protocols: the input `transferFrom` happens inside the callback, which is
/// inside `swap()`. The gas is thus already included in `get_amount_out`.
pub const PROTOCOLS_CALLBACK: &[&str] = &[
    "uniswap_v3",
    "pancakeswap_v3",
    "uniswap_v4",
    "uniswap_v4_hooks",
    "ekubo_v2",
    "ekubo_v3",
    "aerodrome_slipstreams",
    "velodrome_slipstreams",
    "vm:balancer_v3",
];

/// Exact-name protocols where the router-to-pool input transfer is skipped (funds are sent
/// directly without an intermediate router hop). The input transfer cost is therefore not
/// double-counted — it is only charged once, on the first hop or when the optimized path applies.
///
/// This list is incomplete on its own: whole protocol families can qualify too. Use
/// [`optimizable_transfer_in`] for the full classification.
pub const PROTOCOLS_OPTIMIZABLE_TRANSFER_IN: &[&str] =
    &["erc4626", "maverick_v2", "uniswap_v2", "sushiswap_v2", "pancakeswap_v2", "quickswap_v2"];

/// Whether the router-to-pool input transfer is skipped for `protocol_system` (see
/// [`PROTOCOLS_OPTIMIZABLE_TRANSFER_IN`]). Price-level-stream pAMMs are push-payment venues whose
/// `fundsExpectedAddress` is the venue itself, so the whole `pricelevelstream:` family qualifies.
pub fn optimizable_transfer_in(protocol_system: &str) -> bool {
    PROTOCOLS_OPTIMIZABLE_TRANSFER_IN.contains(&protocol_system) ||
        protocol_system.starts_with(PRICE_LEVEL_STREAM_PREFIX)
}

/// Exact-name protocols where the router must `approve(protocol)` before swapping
/// (ProtocolWillDebit). The protocol's `transferFrom` is inside `swap()` and already in the gas
/// computation of `get_amount_out`, but the approval is not.
///
/// Incomplete on its own — use `needs_approval` for the full classification.
pub const PROTOCOLS_NEEDING_APPROVAL: &[&str] = &[
    "vm:balancer_v2",
    "vm:curve",
    "rfq:bebop",
    "rfq:hashflow",
    "rfq:liquorice",
    "rfq:metric",
    "erc4626",
    "ring_swap_v2",
    "sky",
];

/// Whether the router must approve the protocol before swapping (see
/// `PROTOCOLS_NEEDING_APPROVAL`). The PropAMMRouter pulls `tokenIn` with `transferFrom`, so the
/// whole `propammfallback:` family needs the approval.
pub fn needs_approval(protocol_system: &str) -> bool {
    PROTOCOLS_NEEDING_APPROVAL.contains(&protocol_system) ||
        protocol_system.starts_with(PROPAMM_FALLBACK_PREFIX)
}

/// Extra gas the PropAMMRouter adds around the venue swap: the `transferFrom` that pulls `tokenIn`
/// out of the TychoRouter, the venue whitelist read, and its own accounting. Measured on mainnet
/// forks by the router authors at +48,972 (FermiSwap) and +69,952 (Kipseli) over a direct venue
/// call; the higher figure is used so the estimate is not optimistic.
///
/// The pull is inside this figure, which is why `estimate_transfer_overhead` charges no separate
/// input transfer for the family.
///
/// The Uniswap V3 retry is not included, and it costs far more than this: the whole router call on
/// the retry path measures 249,895 gas against 106,784 for the same Uniswap V3 swap called
/// directly (`PropAMMFallbackExecutorTest.testFallbackPathGas` and `testDirectUniswapV3Gas`, block
/// 25682938). A leg that actually falls back therefore costs about 180k more than estimated here.
/// The estimate prices the venue path because that is the path a fresh quote takes.
pub const PROPAMM_FALLBACK_OVERHEAD_GAS: u64 = 70_000;

/// `outputToRouter = true`: the pool sends output to the router, which then does an extra
/// `_transferOut` to the receiver.
pub const PROTOCOLS_OUTPUT_TO_ROUTER: &[&str] =
    &["vm:curve", "rocketpool", "fluid_v1", "native_wrapper"];

pub const ROUTER_FEES_ACTIVE: bool = true;

/// Estimates the total gas cost for executing a `Solution`.
///
/// Groups consecutive swaps on the same groupable protocol (UniswapV4, BalancerV3, Ekubo) via
/// `group_swaps`, which discounts intermediate token transfers saved by flash accounting. For each
/// group (or ungrouped swap), sums:
///
/// - **Pool gas** (`group.estimated_gas`): the simulation-reported cost with intermediate transfer
///   savings already subtracted for multi-swap groups (see `compute_group_gas`).
/// - **Transfer overhead**: the input transfer, approval, and output transfer gas that is NOT
///   captured by `get_amount_out`, computed via `estimate_transfer_overhead` once per group.
/// - **Router overhead**: the user input transfer and — when `ROUTER_FEES_ACTIVE` — the extra
///   output transfer from the fee path. For non-split swaps, callback protocols already include a
///   regular `transferFrom` in their pool gas, so only the Permit2 delta (`USER_PERMIT2_TRANSFER -
///   DEFAULT_TOKEN_TRANSFER_GAS`) is added for Permit2 users; regular transferFrom and vault users
///   pay nothing extra. Non-callback protocols add the full user transfer cost. For split swaps,
///   funds always go through the router first, so the full user input transfer is always added
///   regardless of protocol type.
pub fn estimate_gas_usage(solution: &Solution, strategy: Strategy) -> BigUint {
    let mut total_gas = BigUint::ZERO;
    let swap_groups = group_swaps(solution.swaps());
    for group in &swap_groups {
        let first_swap = &group.swaps[0];
        let last_swap = &group.swaps[group.swaps.len() - 1];
        let group_transfer_overhead = estimate_transfer_overhead(
            &group.protocol_system,
            first_swap.token_in(),
            last_swap.token_out(),
            &strategy,
        );
        total_gas += group_transfer_overhead + &group.estimated_gas;
    }

    // Add user transfer overhead
    if strategy != Strategy::Split {
        if let Some(first_group) = swap_groups.first() {
            let protocol = first_group.protocol_system.as_str();
            // if the solution is not a split swap, the protocol gas usage for callback protocols
            // already includes a regular transfer gas usage:
            // - for the permit2 case we need to deduct the transfer usage already included
            // - for the regular transfer from and vault, do nothing, everything is already
            //   accounted
            if PROTOCOLS_CALLBACK.contains(&protocol) {
                if *solution.user_transfer_type() == UserTransferType::TransferFromPermit2 {
                    total_gas += BigUint::from(USER_PERMIT2_TRANSFER - DEFAULT_TOKEN_TRANSFER_GAS);
                }
            } else {
                total_gas += BigUint::from(match *solution.user_transfer_type() {
                    UserTransferType::TransferFromPermit2 => USER_PERMIT2_TRANSFER,
                    _ => DEFAULT_TOKEN_TRANSFER_GAS, /* TransferFrom and UseVaultsFunds have
                                                      * similar overheads */
                });
            }
        }
    } else {
        // for split swaps the funds always go through the router first
        total_gas += BigUint::from(match *solution.user_transfer_type() {
            UserTransferType::TransferFromPermit2 => USER_PERMIT2_TRANSFER,
            _ => DEFAULT_TOKEN_TRANSFER_GAS, /* TransferFrom and UseVaultsFunds have
                                              * similar overheads */
        });
    }

    // Add fees overhead: when fees are active and the last swap's protocol does not
    // already route output through the router, the fee path adds an extra transfer.
    if let Some(last_group) = swap_groups.last() {
        let last_swap = &last_group.swaps[last_group.swaps.len() - 1];
        let protocol: &str = &last_swap.component().protocol_system;
        let final_swap_output_to_router = (ROUTER_FEES_ACTIVE || strategy == Strategy::Split) &&
            !PROTOCOLS_OUTPUT_TO_ROUTER.contains(&protocol);
        if final_swap_output_to_router {
            total_gas += transfer_token_gas(last_swap.token_out());
        }
    }

    total_gas
}

fn transfer_token_gas(token: &tycho_common::models::token::Token) -> BigUint {
    let measured = token.gas_usage();
    if measured == BigUint::ZERO {
        BigUint::from(DEFAULT_TOKEN_TRANSFER_GAS)
    } else {
        measured
    }
}

/// Gas overhead for token transfers NOT captured by `get_amount_out`.
///
/// `get_amount_out` includes pool computation gas plus any transfers the pool performs during
/// `swap()`. The input transfer is included in `get_amount_out` for:
///
/// - **Callback protocols**: the user pays the pool directly inside the callback.
/// - **ProtocolWillDebit** (`PROTOCOLS_NEEDING_APPROVAL`): the vault pulls from the router via
///   `transferFrom` inside `swap()`.
/// - **Optimizable-transfer-in protocols** (`PROTOCOLS_OPTIMIZABLE_TRANSFER_IN`): the
///   router-to-pool hop is skipped so there is no extra transfer to account for.
/// - **Split strategy**: funds always flow through the router first; the caller accounts for the
///   user input transfer separately.
///
/// For all other protocols (e.g. UniswapV2), the Dispatcher transfers tokens to the pool
/// *before* calling `executor.swap()`, so the cost is NOT included in `get_amount_out`.
///
/// The output transfer (pool → receiver/router) is always included in `get_amount_out`.
///
/// This function adds what's missing:
///
/// - **Input**: for non-callback, non-optimizable, non-split protocols, the `transfer` that happens
///   before `executor.swap()`.
/// - **Approval**: the router's `approve(vault)` for ProtocolWillDebit protocols (Balancer V2,
///   Curve, etc.), regardless of other exclusions.
/// - **Output**: the extra `_transferOut(router, receiver)` for protocols with `outputToRouter =
///   true`, or any protocol when the strategy is Split (output always routes through the router).
fn estimate_transfer_overhead(
    protocol_system: &str,
    token_in: &tycho_common::models::token::Token,
    token_out: &tycho_common::models::token::Token,
    strategy: &Strategy,
) -> BigUint {
    let mut overhead = BigUint::ZERO;

    // Input transfer: only needed when it happens outside swap().
    // - Callback protocols handle it inside the callback (part of swap gas).
    // - Protocols that can have an optimizable transfer in should not be included here either
    //   because the extra transfer is skipped but only if the strategy is not Split
    // - The PropAMMRouter pulls tokenIn itself and PROPAMM_FALLBACK_OVERHEAD_GAS already prices
    //   that pull, so charging a transfer here would count it twice.
    if !PROTOCOLS_CALLBACK.contains(&protocol_system) &&
        !protocol_system.starts_with(PROPAMM_FALLBACK_PREFIX) &&
        (!optimizable_transfer_in(protocol_system) || *strategy == Strategy::Split)
    {
        overhead += transfer_token_gas(token_in);
    }

    if needs_approval(protocol_system) {
        overhead += BigUint::from(TOKEN_APPROVAL_GAS);
    }

    // The venue swap gas from `get_amount_out` prices a direct call to an already-funded venue,
    // not one wrapped by the PropAMMRouter.
    if protocol_system.starts_with(PROPAMM_FALLBACK_PREFIX) {
        overhead += BigUint::from(PROPAMM_FALLBACK_OVERHEAD_GAS);
    }

    // Output transfer: router -> receiver/next pool (only when outputToRouter).
    if PROTOCOLS_OUTPUT_TO_ROUTER.contains(&protocol_system) {
        overhead += transfer_token_gas(token_out);
    }

    overhead
}

#[cfg(test)]
mod tests {
    use num_bigint::BigUint;
    use tycho_common::{models::protocol::ProtocolComponent, Bytes};

    use super::*;
    use crate::encoding::models::{default_token, Solution, Strategy, Swap, UserTransferType};

    fn make_swap(protocol: &str) -> Swap {
        Swap::new(
            ProtocolComponent { protocol_system: protocol.to_string(), ..Default::default() },
            default_token(Bytes::from(vec![0x01u8; 20])),
            default_token(Bytes::from(vec![0x02u8; 20])),
            BigUint::from(100_000u64), // pool gas
        )
    }

    fn make_solution(swaps: Vec<Swap>) -> Solution {
        Solution::new(
            Bytes::from(vec![0x00u8; 20]),
            Bytes::from(vec![0x00u8; 20]),
            Bytes::from(vec![0x01u8; 20]),
            Bytes::from(vec![0x02u8; 20]),
            BigUint::from(1u64),
            BigUint::from(0u64),
            BigUint::from(0u64),
            swaps,
        )
    }

    fn make_swap_with_tokens(protocol: &str, token_in: Bytes, token_out: Bytes, gas: u64) -> Swap {
        Swap::new(
            ProtocolComponent { protocol_system: protocol.to_string(), ..Default::default() },
            default_token(token_in),
            default_token(token_out),
            BigUint::from(gas),
        )
    }

    fn token_a() -> Bytes {
        Bytes::from(vec![0x01u8; 20])
    }
    fn token_b() -> Bytes {
        Bytes::from(vec![0x02u8; 20])
    }
    fn token_c() -> Bytes {
        Bytes::from(vec![0x03u8; 20])
    }

    #[test]
    fn test_single_optimizable_transfer_in() {
        // uniswap_v2 is in PROTOCOLS_OPTIMIZABLE_TRANSFER_IN: the router-to-pool input transfer is
        // skipped (tokens are sent directly). Only the user transfer and the fee path transfer
        // remain.
        let solution = make_solution(vec![make_swap("uniswap_v2")]);
        let gas = estimate_gas_usage(&solution, Strategy::Single);

        // user transfer (TransferFrom)         40_000  ← DEFAULT_TOKEN_TRANSFER_GAS
        // input transfer                            0  ← optimizable, skip router→pool hop
        // pool gas                            100_000
        // fee output transfer                  60_000  ← not in OUTPUT_TO_ROUTER, fee path adds it
        assert_eq!(gas, BigUint::from(200_000u64));
    }

    #[test]
    fn test_single_price_level_stream_pamm() {
        // Push-payment pAMMs are funded directly (fundsExpectedAddress is the venue), so like
        // uniswap_v2 the router-to-pool input transfer is skipped — for every venue under the
        // prefix.
        let solution = make_solution(vec![make_swap("pricelevelstream:fermiswap")]);
        let gas = estimate_gas_usage(&solution, Strategy::Single);

        // user transfer (TransferFrom)         40_000  ← DEFAULT_TOKEN_TRANSFER_GAS
        // input transfer                            0  ← push-payment, funds sent directly
        // pool gas                            100_000
        // fee output transfer                  60_000  ← not in OUTPUT_TO_ROUTER
        assert_eq!(gas, BigUint::from(200_000u64));
    }

    #[test]
    fn test_single_propamm_fallback() {
        // Routing the same venue through the PropAMMRouter switches the leg to ProtocolWillDebit:
        // the router approves, then the PropAMMRouter pulls tokenIn itself.
        let solution = make_solution(vec![make_swap("propammfallback:fermiswap")]);
        let gas = estimate_gas_usage(&solution, Strategy::Single);

        // user transfer (TransferFrom)         40_000  ← DEFAULT_TOKEN_TRANSFER_GAS
        // input transfer                            0  ← the PropAMMRouter pulls, see below
        // approval                             25_000  ← TOKEN_APPROVAL_GAS
        // PropAMMRouter overhead               70_000  ← PROPAMM_FALLBACK_OVERHEAD_GAS, which
        //                                                prices the pull
        // pool gas                            100_000
        // fee output transfer                  60_000  ← not in OUTPUT_TO_ROUTER
        assert_eq!(gas, BigUint::from(295_000u64));
    }

    #[test]
    fn test_split_propamm_fallback() {
        // Split does not reintroduce the input transfer: the leg is still ProtocolWillDebit, so
        // the PropAMMRouter pulls from the router either way.
        let solution = make_solution(vec![
            make_swap("propammfallback:fermiswap").with_split(0.5),
            make_swap("propammfallback:kipseli").with_split(0.5),
        ]);
        let gas = estimate_gas_usage(&solution, Strategy::Split);

        // user transfer (TransferFrom)         40_000  ← DEFAULT_TOKEN_TRANSFER_GAS
        // leg1 approval + router overhead      95_000
        // leg1 pool gas                       100_000
        // leg2 approval + router overhead      95_000
        // leg2 pool gas                       100_000
        // extra output transfer (→ router)     60_000  ← TOKEN_GAS
        assert_eq!(gas, BigUint::from(490_000u64));
    }

    #[test]
    fn test_single_transfer_from_with_approval() {
        let solution = make_solution(vec![make_swap("ring_swap_v2")]);
        let gas = estimate_gas_usage(&solution, Strategy::Single);

        // user transfer (TransferFrom)         40_000  ← DEFAULT_TOKEN_TRANSFER_GAS
        // input transfer (router → wrapper)    60_000  ← measured token gas
        // approval (ProtocolWillDebit)         25_000  ← TOKEN_APPROVAL_GAS
        // pool gas                            100_000
        // fee output transfer                  60_000  ← measured token gas
        assert_eq!(gas, BigUint::from(285_000u64));
    }

    #[test]
    fn test_single_callback_protocol() {
        // uniswap_v3 is a callback protocol: the pool pulls funds inside the callback, so that
        // cost is already in pool gas. The only extra is the fee path's output transfer.
        let solution = make_solution(vec![make_swap("uniswap_v3")]);
        let gas = estimate_gas_usage(&solution, Strategy::Single);

        // user transfer                             0  ← callback + TransferFrom, nothing extra
        // input transfer                            0  ← callback, in pool gas
        // pool gas                            100_000
        // fee output transfer                  60_000  ← TOKEN_GAS, fee path adds _transferOut
        assert_eq!(gas, BigUint::from(160_000u64));
    }

    #[test]
    fn test_single_permit2_with_approval() {
        // rfq:bebop is not callback, not optimizable, and needs an approval (ProtocolWillDebit).
        // With Permit2 the user transfer costs 80k instead of 40k.
        let solution = make_solution(vec![make_swap("rfq:bebop")])
            .with_user_transfer_type(UserTransferType::TransferFromPermit2);
        let gas = estimate_gas_usage(&solution, Strategy::Single);

        // user permit2 transfer                80_000  ← USER_PERMIT2_TRANSFER
        // approval (ProtocolWillDebit)         25_000  ← TOKEN_APPROVAL_GAS
        // input transfer (router → pool)       60_000  ← TOKEN_GAS
        // pool gas                            100_000
        // fee output transfer                  60_000  ← TOKEN_GAS
        assert_eq!(gas, BigUint::from(325_000u64));
    }

    #[test]
    fn test_sequential_two_hops() {
        // Sequential uniswap_v2 → uniswap_v3: uniswap_v2 is in PROTOCOLS_OPTIMIZABLE_TRANSFER_IN
        // so its input transfer is skipped; uniswap_v3 is callback so its input is also
        // skipped. User transfer once (first hop, non-callback). Fee transfer once (last
        // hop, uniswap_v3).
        let solution = make_solution(vec![make_swap("uniswap_v2"), make_swap("uniswap_v3")]);
        let gas = estimate_gas_usage(&solution, Strategy::Sequential);

        // user transfer (first hop, TransferFrom) 40_000  ← DEFAULT_TOKEN_TRANSFER_GAS
        // hop1 input transfer                       0  ← uniswap_v2 is optimizable
        // hop1 pool gas                        100_000
        // hop2 input transfer                       0  ← callback
        // hop2 pool gas                        100_000
        // fee output transfer (last hop)        60_000  ← not in OUTPUT_TO_ROUTER
        assert_eq!(gas, BigUint::from(300_000u64));
    }

    #[test]
    fn test_sequential_two_hops_output_to_router() {
        // Sequential curve → uniswap_v2:
        // vm:curve outputs to router so a new transfer needs to be accounted for at the end of the
        // first hop
        let solution = make_solution(vec![make_swap("vm:curve"), make_swap("uniswap_v2")]);
        let gas = estimate_gas_usage(&solution, Strategy::Sequential);

        // user transfer (first hop, TransferFrom) 40_000  ← DEFAULT_TOKEN_TRANSFER_GAS
        // hop1 pool gas                        100_000
        // hop1 input transfer (router → pool)   60_000  ← vm:curve is not optimizable/callback
        // hop1 approval (ProtocolWillDebit)     25_000  ← TOKEN_APPROVAL_GAS
        // hop1 output to router                 60_000  ← vm:curve in OUTPUT_TO_ROUTER
        // hop2 pool gas                        100_000
        // hop2 input transfer                       0  ← uniswap_v2 is optimizable
        // fee output transfer (last hop)        60_000  ← uniswap_v2 not in OUTPUT_TO_ROUTER
        assert_eq!(gas, BigUint::from(445_000u64));
    }

    #[test]
    fn test_single_grouped_two_hops_usv4() {
        // Two consecutive USV4 hops get grouped: intermediate transfer is saved.
        // A→B→C on uniswap_v4 becomes one group with flash accounting.
        let solution = make_solution(vec![
            make_swap_with_tokens("uniswap_v4", token_a(), token_b(), 100_000),
            make_swap_with_tokens("uniswap_v4", token_b(), token_c(), 100_000),
        ]);
        let gas = estimate_gas_usage(&solution, Strategy::Single);

        // group estimated_gas: 200_000 - token_b.gas(60k) - token_b.gas(60k) = 80_000
        //   (first.token_out=60k + last.token_in=60k subtracted by compute_group_gas)
        // transfer overhead: 0 (callback protocol)
        // user transfer: 0 (callback + TransferFrom)
        // fee output transfer: 60_000 (not in OUTPUT_TO_ROUTER)
        assert_eq!(gas, BigUint::from(140_000u64));
    }

    #[test]
    fn test_sequential_mixed_grouped_and_ungrouped() {
        // A→B→C on uniswap_v4 (grouped), then C→D on uniswap_v2 (separate).
        let token_d = Bytes::from(vec![0x04u8; 20]);
        let solution = make_solution(vec![
            make_swap_with_tokens("uniswap_v4", token_a(), token_b(), 100_000),
            make_swap_with_tokens("uniswap_v4", token_b(), token_c(), 100_000),
            make_swap_with_tokens("uniswap_v2", token_c(), token_d.clone(), 100_000),
        ]);
        let gas = estimate_gas_usage(&solution, Strategy::Sequential);

        // group1 (usv4): 200_000 - 60k - 60k = 80_000 (intermediate B transfers saved)
        // group1 transfer overhead: 0 (callback)
        // group2 (usv2): 100_000
        // group2 transfer overhead: 0 (optimizable transfer in)
        // user transfer: 0 (first group is callback + TransferFrom)
        // fee output transfer: 60_000 (last swap is uniswap_v2, not in OUTPUT_TO_ROUTER)
        assert_eq!(gas, BigUint::from(240_000u64));
    }

    #[test]
    fn test_split_two_legs() {
        // Split swap with two equal legs. User transfer once. Extra router transfer at the end
        // because it is a split swap
        let solution = make_solution(vec![
            make_swap("uniswap_v2").with_split(0.5),
            make_swap("uniswap_v3").with_split(0.5),
        ]);
        let gas = estimate_gas_usage(&solution, Strategy::Split);

        // user transfer (TransferFrom)          40_000  ← DEFAULT_TOKEN_TRANSFER_GAS
        // transfer for leg1 (not optimized)     60_000  ← TOKEN_GAS
        // leg1 pool gas                        100_000
        // leg2 pool gas                        100_000
        // extra output transfer (→ router)       60_000  ← TOKEN_GAS
        assert_eq!(gas, BigUint::from(360_000u64));
    }
}