AMM Contracts

┌─────────────────────────────────────────────────────────────┐
│                      AMM System                              │
├─────────────────────────────────────────────────────────────┤
│                                                              │
│  ┌─────────────┐     ┌─────────────┐     ┌─────────────┐   │
│  │   Router    │────►│   Factory   │────►│    Pairs    │   │
│  │             │     │             │     │             │   │
│  │ • swap      │     │ • createPair│     │ • mint      │   │
│  │ • addLiq    │     │ • getPair   │     │ • burn      │   │
│  │ • removeLiq │     │ • feeTo     │     │ • swap      │   │
│  └─────────────┘     └─────────────┘     └─────────────┘   │
│                                                              │
└─────────────────────────────────────────────────────────────┘

// Initialize pair
nexus_fn! {
    fn init_pair(t0: Address, t1: Address, factory_addr: Address) {
        TOKEN0.set(&b"val".as_slice(), t0);
        TOKEN1.set(&b"val".as_slice(), t1);
        FACTORY.set(&b"val".as_slice(), factory_addr);
        // ...
    }
}

// Mint LP tokens
nexus_fn! {
    fn mint(to: Address) {
        let t0 = TOKEN0.get(&b"val".as_slice());
        let t1 = TOKEN1.get(&b"val".as_slice());
        let r0 = RES0.get(&b"val".as_slice());
        let r1 = RES1.get(&b"val".as_slice());
        
        let balance0 = get_token_balance(&t0, &this_addr);
        let balance1 = get_token_balance(&t1, &this_addr);
        
        let amount0 = balance0.sub(r0);
        let amount1 = balance1.sub(r1);
        
        let liquidity = if supply == U256::zero() {
            // First liquidity: sqrt(amount0 * amount1) - MINIMUM_LIQUIDITY
            sqrt_u256(amount0.mul(amount1)).sub(U256::from(MINIMUM_LIQUIDITY))
        } else {
            // Subsequent: min(amount0/reserve0, amount1/reserve1) * supply
            core::cmp::min(
                amount0.mul(supply).div(r0),
                amount1.mul(supply).div(r1)
            )
        };
        
        LP_BALANCES.set(&to, bal.add(liquidity));
        _update(balance0, balance1, r0, r1);
    }
}

// Swap tokens
nexus_fn! {
    fn swap(amount0_out: U256, amount1_out: U256, to: Address, data_len: u64) {
        require!(amount0_out > U256::zero() || amount1_out > U256::zero());
        require!(amount0_out < r0 && amount1_out < r1);
        
        // Transfer output tokens
        if amount0_out > U256::zero() {
            xcc::call(&t0.0, "transfer", &json_encode_addr_u256(&to, &amount0_out));
        }
        if amount1_out > U256::zero() {
            xcc::call(&t1.0, "transfer", &json_encode_addr_u256(&to, &amount1_out));
        }
        
        // Flash swap callback (if data provided)
        if data_len > 0 {
            xcc::call(&to.0, "satsAmmV2Call", &args);
        }
        
        // Verify K invariant (with 0.3% fee)
        let balance0_adj = balance0.mul(1000).sub(amount0_in.mul(3));
        let balance1_adj = balance1.mul(1000).sub(amount1_in.mul(3));
        require!(balance0_adj.mul(balance1_adj) >= r0.mul(r1).mul(1000000));
        
        _update(balance0, balance1, r0, r1);
    }
}

fn _update(balance0: U256, balance1: U256, reserve0: U256, reserve1: U256) {
    let time_elapsed = block_timestamp - last_timestamp;
    
    if time_elapsed > 0 && reserve0 > 0 && reserve1 > 0 {
        // Accumulate price * time for TWAP
        let price0 = reserve1.mul(2^112).div(reserve0);
        let price1 = reserve0.mul(2^112).div(reserve1);
        
        PRICE0_CUMULATIVE_LAST += price0 * time_elapsed;
        PRICE1_CUMULATIVE_LAST += price1 * time_elapsed;
    }
    
    RES0.set(balance0);
    RES1.set(balance1);
    BLOCK_TIMESTAMP_LAST.set(block_timestamp);
}

// Embedded pair bytecode (like Solidity's type(Pair).creationCode)
const PAIR_CODE: &[u8] = include_bytes!("../../sats_amm_v2.wasm");

nexus_fn! {
    fn createPair(tokenA: Address, tokenB: Address) {
        // Sort tokens
        let (token0, token1) = if tokenA.0 < tokenB.0 {
            (tokenA, tokenB)
        } else {
            (tokenB, tokenA)
        };
        
        require!(!token0.is_zero());
        require!(PAIRS.get(&key).is_zero(), "PAIR_EXISTS");
        
        // Deploy with CREATE2 (deterministic address)
        let mut salt = Vec::new();
        salt.extend_from_slice(&token0.0);
        salt.extend_from_slice(&token1.0);
        
        let pair_addr = xcc::deploy(PAIR_CODE, &[], &salt);
        
        // Initialize pair
        xcc::call(&pair_addr.0, "init_pair", &[token0, token1, factory]);
        
        PAIRS.set(&key, pair_addr);
        emit("PairCreated", &[token0, token1, pair_addr]);
    }
}

nexus_fn! {
    fn swapExactTokensForTokens(
        amount_in: U256,
        amount_out_min: U256,
        path: Vec<Address>,
        to: Address,
        deadline: u64
    ) {
        require!(block.timestamp <= deadline, "EXPIRED");
        
        // Calculate amounts through path
        let mut amounts = vec![amount_in];
        for i in 0..(path.len() - 1) {
            let pair = get_pair(&path[i], &path[i + 1]);
            let (reserve_in, reserve_out) = get_reserves(&pair, &path[i], &path[i + 1]);
            amounts.push(get_amount_out(amounts[i], reserve_in, reserve_out));
        }
        
        require!(amounts.last() >= amount_out_min, "INSUFFICIENT_OUTPUT");
        
        // Transfer input to first pair
        xcc::call(&path[0].0, "transferFrom", &[sender, first_pair, amount_in]);
        
        // Execute swaps through path
        _swap(&amounts, &path, &to);
    }
}

fn get_amount_out(amount_in: U256, reserve_in: U256, reserve_out: U256) -> U256 {
    // 0.3% fee: multiply by 997/1000
    let amount_in_with_fee = amount_in.mul(997);
    let numerator = amount_in_with_fee.mul(reserve_out);
    let denominator = reserve_in.mul(1000).add(amount_in_with_fee);
    numerator.div(denominator)
}

nexus_fn! {
    fn addLiquidity(
        token_a: Address,
        token_b: Address,
        amount_a_desired: U256,
        amount_b_desired: U256,
        amount_a_min: U256,
        amount_b_min: U256,
        to: Address,
        deadline: u64
    ) {
        require!(block.timestamp <= deadline, "EXPIRED");
        
        // Calculate optimal amounts to maintain ratio
        let (amount_a, amount_b) = _add_liquidity(
            token_a, token_b,
            amount_a_desired, amount_b_desired,
            amount_a_min, amount_b_min
        );
        
        // Transfer tokens to pair
        xcc::call(&token_a.0, "transferFrom", &[sender, pair, amount_a]);
        xcc::call(&token_b.0, "transferFrom", &[sender, pair, amount_b]);
        
        // Mint LP tokens
        let liquidity = xcc::call(&pair.0, "mint", &[to]);
        
        ret::bytes(&[amount_a, amount_b, liquidity])
    }
}

// Swap tokens
await router.swapExactTokensForTokens(
  amountIn,
  minAmountOut,
  [tokenA, tokenB],
  recipient,
  deadline
);

// Add liquidity
await router.addLiquidity(
  tokenA, tokenB,
  amountADesired, amountBDesired,
  amountAMin, amountBMin,
  recipient,
  deadline
);

// Remove liquidity
await router.removeLiquidity(
  tokenA, tokenB,
  liquidity,
  amountAMin, amountBMin,
  recipient,
  deadline
);