Skip to main content

tempo_precompiles/tip_fee_manager/
amm.rs

1use crate::{
2    error::{Result, TempoPrecompileError},
3    storage::Handler,
4    tip_fee_manager::{ITIPFeeAMM, TIPFeeAMMError, TIPFeeAMMEvent, TipFeeManager},
5    tip20::{ITIP20, TIP20Token, validate_usd_currency},
6};
7use alloy::{
8    primitives::{Address, B256, U256, keccak256, uint},
9    sol_types::SolValue,
10};
11use tempo_precompiles_macros::Storable;
12
13/// Fee multiplier for fee swaps: 0.9970 scaled by 10000 (30 bps fee).
14pub const M: U256 = uint!(9970_U256);
15/// Fee multiplier for rebalance swaps: 0.9985 scaled by 10000.
16pub const N: U256 = uint!(9985_U256);
17/// Scale factor for fixed-point AMM arithmetic (10000).
18pub const SCALE: U256 = uint!(10000_U256);
19/// Minimum liquidity locked permanently when initializing a pool.
20pub const MIN_LIQUIDITY: U256 = uint!(1000_U256);
21
22/// Computes the output amount for a fee swap: `amount_in * M / SCALE`.
23///
24/// # Errors
25/// - `UnderOverflow` — multiplication of `amount_in * M` overflows
26#[inline]
27pub fn compute_amount_out(amount_in: U256) -> Result<U256> {
28    amount_in
29        .checked_mul(M)
30        .map(|product| product / SCALE)
31        .ok_or(TempoPrecompileError::under_overflow())
32}
33
34/// AMM pool reserves for a user-token / validator-token pair.
35#[derive(Debug, Clone, Default, Storable)]
36pub struct Pool {
37    /// Reserve of the user's fee token.
38    pub reserve_user_token: u128,
39    /// Reserve of the validator's fee token.
40    pub reserve_validator_token: u128,
41}
42
43impl From<Pool> for ITIPFeeAMM::Pool {
44    fn from(value: Pool) -> Self {
45        Self {
46            reserveUserToken: value.reserve_user_token,
47            reserveValidatorToken: value.reserve_validator_token,
48        }
49    }
50}
51
52/// Identifies a directional token pair in the fee AMM.
53#[derive(Debug, Clone, PartialEq, Eq, Hash, Storable)]
54pub struct PoolKey {
55    /// The fee token chosen by the user (transaction sender).
56    pub user_token: Address,
57    /// The fee token chosen by the validator (block producer).
58    pub validator_token: Address,
59}
60
61impl Pool {
62    /// Decodes a [`Pool`] from a raw EVM storage slot value (needed from changeset diffs).
63    pub fn decode_from_slot(slot_value: U256) -> Self {
64        use crate::storage::{LayoutCtx, Storable, packing::PackedSlot};
65
66        // NOTE: fine to expect, as `StorageOps` on `PackedSlot` are infallible
67        Self::load(&PackedSlot(slot_value), U256::ZERO, LayoutCtx::FULL)
68            .expect("unable to decode Pool from slot")
69    }
70}
71
72impl PoolKey {
73    /// Creates a new pool key from user and validator token addresses.
74    /// This key uniquely identifies a trading pair in the AMM.
75    pub fn new(user_token: Address, validator_token: Address) -> Self {
76        Self {
77            user_token,
78            validator_token,
79        }
80    }
81
82    /// Generates a unique pool ID by hashing the token pair addresses.
83    /// Uses keccak256 to create a deterministic identifier for this pool.
84    pub fn get_id(&self) -> B256 {
85        keccak256((self.user_token, self.validator_token).abi_encode())
86    }
87}
88
89/// AMM path [`TipFeeManager`] will take to swap `user_token` into `validator_token` for fee collection.
90#[derive(Debug, Clone, Copy, PartialEq, Eq)]
91pub enum FeeRoute {
92    /// User and validator share the same fee token; no swap is performed.
93    SameToken,
94    /// Direct pool `(user_token, validator_token)` swap.
95    Direct,
96    /// Two-hop swap (T5+): routes through `intermediate = userToken.quoteToken()`.
97    /// Each hop applies the standard `M = 9970/10000` rate sequentially.
98    TwoHop(Address),
99}
100
101/// Pools read during planning, paired with their observed validator-token reserve.
102pub type PoolData = ((Address, Address), u128);
103
104impl TipFeeManager {
105    /// Returns the deterministic pool ID for a directional token pair. Note that the pool id is
106    /// order-dependent: `(A, B)` produces a different ID than `(B, A)`.
107    pub fn pool_id(&self, user_token: Address, validator_token: Address) -> B256 {
108        PoolKey::new(user_token, validator_token).get_id()
109    }
110
111    /// Returns the [`Pool`] reserves for the given user/validator token pair.
112    pub fn get_pool(&self, call: ITIPFeeAMM::getPoolCall) -> Result<Pool> {
113        let pool_id = self.pool_id(call.userToken, call.validatorToken);
114        self.pools[pool_id].read()
115    }
116
117    /// Reserves pool liquidity in transient storage for a pending fee swap.
118    #[inline]
119    pub fn reserve_pool_liquidity(&mut self, pool_id: B256, amount: u128) -> Result<()> {
120        self.pending_fee_swap_reservation[pool_id].t_write(amount)
121    }
122
123    /// Executes a rebalance swap: sells `amount_out` of user-token from the pool in exchange for
124    /// validator-token at the rebalance rate (`N / SCALE`). Used by arbitrageurs to rebalance reserves.
125    ///
126    /// # Errors
127    /// - `InvalidAmount` — `amount_out` is zero or exceeds `u128`
128    /// - `InsufficientReserves` — adding `amount_in` overflows the validator reserve
129    /// - `InsufficientLiquidity` — remaining reserve would violate the pending reservation (T1C+)
130    /// - `UnderOverflow` — arithmetic overflow computing `amount_in`
131    pub fn rebalance_swap(
132        &mut self,
133        msg_sender: Address,
134        user_token: Address,
135        validator_token: Address,
136        amount_out: U256,
137        to: Address,
138    ) -> Result<U256> {
139        if amount_out.is_zero() {
140            return Err(TIPFeeAMMError::invalid_amount().into());
141        }
142
143        let pool_id = self.pool_id(user_token, validator_token);
144        let mut pool = self.pools[pool_id].read()?;
145
146        // Rebalancing swaps are always from validatorToken to userToken
147        // Calculate input and update reserves
148        let amount_in = amount_out
149            .checked_mul(N)
150            .and_then(|product| product.checked_div(SCALE))
151            .and_then(|result| result.checked_add(U256::ONE))
152            .ok_or(TempoPrecompileError::under_overflow())?;
153
154        let amount_in: u128 = amount_in
155            .try_into()
156            .map_err(|_| TIPFeeAMMError::invalid_amount())?;
157        let amount_out: u128 = amount_out
158            .try_into()
159            .map_err(|_| TIPFeeAMMError::invalid_amount())?;
160
161        pool.reserve_validator_token = pool
162            .reserve_validator_token
163            .checked_add(amount_in)
164            .ok_or(TIPFeeAMMError::insufficient_reserves())?;
165
166        pool.reserve_user_token = pool
167            .reserve_user_token
168            .checked_sub(amount_out)
169            .ok_or(TIPFeeAMMError::invalid_amount())?;
170
171        if self.storage.spec().is_t1c() {
172            let reserved = self.pending_fee_swap_reservation[pool_id].t_read()?;
173            if pool.reserve_validator_token < reserved {
174                return Err(TIPFeeAMMError::insufficient_liquidity().into());
175            }
176        }
177
178        self.pools[pool_id].write(pool)?;
179
180        let amount_in = U256::from(amount_in);
181        let amount_out = U256::from(amount_out);
182        TIP20Token::from_address(validator_token)?.system_transfer_from(
183            self.address,
184            msg_sender,
185            amount_in,
186        )?;
187
188        TIP20Token::from_address(user_token)?.transfer(
189            self.address,
190            ITIP20::transferCall {
191                to,
192                amount: amount_out,
193            },
194        )?;
195
196        self.emit_event(TIPFeeAMMEvent::rebalance_swap(
197            user_token,
198            validator_token,
199            msg_sender,
200            amount_in,
201            amount_out,
202        ))?;
203
204        Ok(amount_in)
205    }
206
207    /// Mints LP tokens by depositing validator-token into a pool.
208    ///
209    /// On first deposit the pool is initialized with equal reserves and [`MIN_LIQUIDITY`] is
210    /// permanently locked. Subsequent deposits mint pro-rata to existing supply. Both tokens
211    /// must be distinct, USD-denominated TIP-20s.
212    ///
213    /// NOTE: Validators who also provide liquidity have an information advantage over non-validator
214    /// LPs. Because validators choose their preferred fee token and control transaction inclusion
215    /// order as block producers, a validator-LP can predict which pool will receive fee-swap
216    /// revenue and position liquidity accordingly.
217    ///
218    /// # Errors
219    /// - `IdenticalAddresses` — `user_token` equals `validator_token`
220    /// - `InvalidAmount` — `amount_validator_token` is zero or exceeds `u128`
221    /// - `InvalidCurrency` — either token is not USD-denominated
222    /// - `InsufficientLiquidity` — initial deposit ≤ `MIN_LIQUIDITY`, or zero liquidity minted
223    /// - `InvalidSwapCalculation` — pro-rata arithmetic fails
224    /// - `UnderOverflow` — supply or balance overflow
225    pub fn mint(
226        &mut self,
227        msg_sender: Address,
228        user_token: Address,
229        validator_token: Address,
230        amount_validator_token: U256,
231        to: Address,
232    ) -> Result<U256> {
233        if user_token == validator_token {
234            return Err(TIPFeeAMMError::identical_addresses().into());
235        }
236
237        if amount_validator_token.is_zero() {
238            return Err(TIPFeeAMMError::invalid_amount().into());
239        }
240
241        // Validate both tokens are USD currency
242        validate_usd_currency(user_token)?;
243        validate_usd_currency(validator_token)?;
244
245        let pool_id = self.pool_id(user_token, validator_token);
246        let mut pool = self.pools[pool_id].read()?;
247        let mut total_supply = self.get_total_supply(pool_id)?;
248
249        let liquidity = if pool.reserve_user_token == 0 && pool.reserve_validator_token == 0 {
250            let half_amount = amount_validator_token
251                .checked_div(uint!(2_U256))
252                .ok_or(TempoPrecompileError::under_overflow())?;
253
254            if half_amount <= MIN_LIQUIDITY {
255                return Err(TIPFeeAMMError::insufficient_liquidity().into());
256            }
257
258            total_supply = total_supply
259                .checked_add(MIN_LIQUIDITY)
260                .ok_or(TempoPrecompileError::under_overflow())?;
261            self.set_total_supply(pool_id, total_supply)?;
262
263            half_amount
264                .checked_sub(MIN_LIQUIDITY)
265                .ok_or(TIPFeeAMMError::insufficient_liquidity())?
266        } else {
267            // Subsequent deposits: mint as if user called rebalanceSwap then minted with both
268            // liquidity = amountValidatorToken * _totalSupply / (V + n * U), with n = N / SCALE
269            let product = N
270                .checked_mul(U256::from(pool.reserve_user_token))
271                .and_then(|product| product.checked_div(SCALE))
272                .ok_or(TIPFeeAMMError::invalid_swap_calculation())?;
273
274            let denom = U256::from(pool.reserve_validator_token)
275                .checked_add(product)
276                .ok_or(TIPFeeAMMError::invalid_amount())?;
277
278            if denom.is_zero() {
279                return Err(TIPFeeAMMError::division_by_zero().into());
280            }
281
282            amount_validator_token
283                .checked_mul(total_supply)
284                .and_then(|numerator| numerator.checked_div(denom))
285                .ok_or(TIPFeeAMMError::invalid_swap_calculation())?
286        };
287
288        if liquidity.is_zero() {
289            return Err(TIPFeeAMMError::insufficient_liquidity().into());
290        }
291
292        // Transfer validator tokens from user
293        let _ = TIP20Token::from_address(validator_token)?.system_transfer_from(
294            self.address,
295            msg_sender,
296            amount_validator_token,
297        )?;
298
299        // Update reserves
300        let validator_amount: u128 = amount_validator_token
301            .try_into()
302            .map_err(|_| TIPFeeAMMError::invalid_amount())?;
303
304        pool.reserve_validator_token = pool
305            .reserve_validator_token
306            .checked_add(validator_amount)
307            .ok_or(TIPFeeAMMError::invalid_amount())?;
308
309        self.pools[pool_id].write(pool)?;
310
311        // Mint LP tokens
312        self.set_total_supply(
313            pool_id,
314            total_supply
315                .checked_add(liquidity)
316                .ok_or(TempoPrecompileError::under_overflow())?,
317        )?;
318
319        let balance = self.get_liquidity_balances(pool_id, to)?;
320        self.set_liquidity_balances(
321            pool_id,
322            to,
323            balance
324                .checked_add(liquidity)
325                .ok_or(TempoPrecompileError::under_overflow())?,
326        )?;
327
328        // Emit Mint event
329        self.emit_event(TIPFeeAMMEvent::mint(
330            msg_sender,
331            to,
332            user_token,
333            validator_token,
334            amount_validator_token,
335            liquidity,
336        ))?;
337
338        Ok(liquidity)
339    }
340
341    /// Burns LP tokens and returns the pro-rata share of both pool tokens to `to`.
342    ///
343    /// On T1C+ the burn is rejected if the remaining validator-token reserve would fall below
344    /// the pending fee-swap reservation set by [`TipFeeManager::reserve_pool_liquidity`].
345    ///
346    /// # Errors
347    /// - `IdenticalAddresses` — `user_token` equals `validator_token`
348    /// - `InvalidAmount` — `liquidity` is zero or amounts exceed `u128`
349    /// - `InvalidCurrency` — either token is not USD-denominated
350    /// - `InsufficientLiquidity` — caller's balance < `liquidity`, or remaining reserve would
351    ///   violate the pending reservation (T1C+)
352    /// - `InsufficientReserves` — pool reserves underflow after withdrawal
353    /// - `UnderOverflow` — supply or balance arithmetic overflows
354    pub fn burn(
355        &mut self,
356        msg_sender: Address,
357        user_token: Address,
358        validator_token: Address,
359        liquidity: U256,
360        to: Address,
361    ) -> Result<(U256, U256)> {
362        if user_token == validator_token {
363            return Err(TIPFeeAMMError::identical_addresses().into());
364        }
365
366        if liquidity.is_zero() {
367            return Err(TIPFeeAMMError::invalid_amount().into());
368        }
369
370        // Validate both tokens are USD currency
371        validate_usd_currency(user_token)?;
372        validate_usd_currency(validator_token)?;
373
374        let pool_id = self.pool_id(user_token, validator_token);
375        // Check user has sufficient liquidity
376        let balance = self.get_liquidity_balances(pool_id, msg_sender)?;
377        if balance < liquidity {
378            return Err(TIPFeeAMMError::insufficient_liquidity().into());
379        }
380
381        let mut pool = self.pools[pool_id].read()?;
382        // Calculate amounts to return
383        let (amount_user_token, amount_validator_token) =
384            self.calculate_burn_amounts(&pool, pool_id, liquidity)?;
385
386        // T1C+: Check that burn leaves enough liquidity for pending fee swaps
387        // Reservation is set by reserve_pool_liquidity() in collect_fee_pre_tx
388        let validator_amount: u128 = amount_validator_token
389            .try_into()
390            .map_err(|_| TIPFeeAMMError::invalid_amount())?;
391        let available_after_burn = pool
392            .reserve_validator_token
393            .checked_sub(validator_amount)
394            .ok_or(TIPFeeAMMError::insufficient_reserves())?;
395        if self.storage.spec().is_t1c() {
396            let reserved = self.pending_fee_swap_reservation[pool_id].t_read()?;
397            if available_after_burn < reserved {
398                return Err(TIPFeeAMMError::insufficient_liquidity().into());
399            }
400        }
401
402        // Burn LP tokens
403        self.set_liquidity_balances(
404            pool_id,
405            msg_sender,
406            balance
407                .checked_sub(liquidity)
408                .ok_or(TempoPrecompileError::under_overflow())?,
409        )?;
410        let total_supply = self.get_total_supply(pool_id)?;
411        self.set_total_supply(
412            pool_id,
413            total_supply
414                .checked_sub(liquidity)
415                .ok_or(TempoPrecompileError::under_overflow())?,
416        )?;
417
418        // Update reserves with underflow checks
419        let user_amount: u128 = amount_user_token
420            .try_into()
421            .map_err(|_| TIPFeeAMMError::invalid_amount())?;
422        let validator_amount: u128 = amount_validator_token
423            .try_into()
424            .map_err(|_| TIPFeeAMMError::invalid_amount())?;
425
426        pool.reserve_user_token = pool
427            .reserve_user_token
428            .checked_sub(user_amount)
429            .ok_or(TIPFeeAMMError::insufficient_reserves())?;
430        pool.reserve_validator_token = pool
431            .reserve_validator_token
432            .checked_sub(validator_amount)
433            .ok_or(TIPFeeAMMError::insufficient_reserves())?;
434        self.pools[pool_id].write(pool)?;
435
436        // Transfer tokens to user
437        let _ = TIP20Token::from_address(user_token)?.transfer(
438            self.address,
439            ITIP20::transferCall {
440                to,
441                amount: amount_user_token,
442            },
443        )?;
444
445        let _ = TIP20Token::from_address(validator_token)?.transfer(
446            self.address,
447            ITIP20::transferCall {
448                to,
449                amount: amount_validator_token,
450            },
451        )?;
452
453        // Emit Burn event
454        self.emit_event(TIPFeeAMMEvent::burn(
455            msg_sender,
456            user_token,
457            validator_token,
458            amount_user_token,
459            amount_validator_token,
460            liquidity,
461            to,
462        ))?;
463
464        Ok((amount_user_token, amount_validator_token))
465    }
466
467    /// Calculate burn amounts for liquidity withdrawal
468    fn calculate_burn_amounts(
469        &self,
470        pool: &Pool,
471        pool_id: B256,
472        liquidity: U256,
473    ) -> Result<(U256, U256)> {
474        let total_supply = self.get_total_supply(pool_id)?;
475        let amount_user_token = liquidity
476            .checked_mul(U256::from(pool.reserve_user_token))
477            .and_then(|product| product.checked_div(total_supply))
478            .ok_or(TempoPrecompileError::under_overflow())?;
479        let amount_validator_token = liquidity
480            .checked_mul(U256::from(pool.reserve_validator_token))
481            .and_then(|product| product.checked_div(total_supply))
482            .ok_or(TempoPrecompileError::under_overflow())?;
483
484        Ok((amount_user_token, amount_validator_token))
485    }
486
487    /// Plans the AMM path needed to swap `max_amount` of `user_token` into `validator_token`
488    /// under the active hardfork. Read-only; does not reserve.
489    ///
490    /// On T5+ falls back to a two-hop path through `userToken.quoteToken()` as per [TIP-1033].
491    /// Returns `(route, queried_intermediate, pools)`:
492    /// - `route` is `None` when no path has sufficient liquidity.
493    /// - `queried_intermediate` is `Some(addr)` whenever `userToken.quoteToken()` was read,
494    ///   regardless of whether the value is usable. Callers can cache it to skip the cold
495    ///   storage read on subsequent admissions.
496    /// - `pools` lists every pool slot read during planning, paired with its observed reserve.
497    ///
498    /// # Errors
499    /// - `InvalidToken` — `user_token` does not have a valid TIP-20 prefix
500    /// - `UnderOverflow` — fee-amount arithmetic overflows
501    ///
502    /// [TIP-1033]: <https://docs.tempo.xyz/protocol/tips/tip-1033>
503    pub fn plan_fee_route(
504        &self,
505        user_token: Address,
506        validator_token: Address,
507        max_amount: U256,
508    ) -> Result<(Option<FeeRoute>, Option<Address>, Vec<PoolData>)> {
509        let mut data = Vec::new();
510
511        if user_token == validator_token {
512            return Ok((Some(FeeRoute::SameToken), None, data));
513        }
514
515        // Direct (single-hop) path — always checked.
516        let direct = self.pools[self.pool_id(user_token, validator_token)].read()?;
517        data.push((
518            (user_token, validator_token),
519            direct.reserve_validator_token,
520        ));
521        let amount_out = compute_amount_out(max_amount)?;
522        if amount_out <= U256::from(direct.reserve_validator_token) {
523            return Ok((Some(FeeRoute::Direct), None, data));
524        }
525
526        // T5+: two-hop fallback through `userToken.quoteToken()`.
527        if !self.storage.spec().is_t5() {
528            return Ok((None, None, data));
529        }
530
531        // TIP-20 token graph forbids self-quoting, so `intermediate == user_token` is unreachable.
532        let mid_token = TIP20Token::from_address(user_token)?.quote_token()?;
533        if mid_token.is_zero() || mid_token == validator_token {
534            return Ok((None, Some(mid_token), data));
535        }
536
537        // First leg: user_token -> intermediate.
538        let leg1 = self.pools[self.pool_id(user_token, mid_token)].read()?;
539        data.push(((user_token, mid_token), leg1.reserve_validator_token));
540        if amount_out > U256::from(leg1.reserve_validator_token) {
541            return Ok((None, Some(mid_token), data));
542        }
543
544        // Second leg: intermediate -> validator_token.
545        let amount_out2 = compute_amount_out(amount_out)?;
546        let leg2 = self.pools[self.pool_id(mid_token, validator_token)].read()?;
547        data.push(((mid_token, validator_token), leg2.reserve_validator_token));
548        if amount_out2 > U256::from(leg2.reserve_validator_token) {
549            return Ok((None, Some(mid_token), data));
550        }
551
552        Ok((Some(FeeRoute::TwoHop(mid_token)), Some(mid_token), data))
553    }
554
555    /// Executes a fee swap, converting `user_token` to `validator_token` at a fixed rate m = 0.997
556    /// Called internally by [`TipFeeManager::collect_fee_post_tx`] during post-tx fee collection.
557    ///
558    /// # Errors
559    /// - `InsufficientLiquidity` — pool validator-token reserve is below the required output
560    /// - `UnderOverflow` — reserve arithmetic overflows or amounts exceed `u128`
561    pub fn execute_fee_swap(
562        &mut self,
563        user_token: Address,
564        validator_token: Address,
565        amount_in: U256,
566    ) -> Result<U256> {
567        let pool_id = self.pool_id(user_token, validator_token);
568        let mut pool = self.pools[pool_id].read()?;
569
570        // Calculate output at fixed price m = 0.9970
571        let amount_out = compute_amount_out(amount_in)?;
572
573        // Check if there's enough validatorToken available
574        if amount_out > U256::from(pool.reserve_validator_token) {
575            return Err(TIPFeeAMMError::insufficient_liquidity().into());
576        }
577
578        // Update reserves
579        let amount_in_u128: u128 = amount_in
580            .try_into()
581            .map_err(|_| TempoPrecompileError::under_overflow())?;
582        let amount_out_u128: u128 = amount_out
583            .try_into()
584            .map_err(|_| TempoPrecompileError::under_overflow())?;
585
586        pool.reserve_user_token = pool
587            .reserve_user_token
588            .checked_add(amount_in_u128)
589            .ok_or(TempoPrecompileError::under_overflow())?;
590        pool.reserve_validator_token = pool
591            .reserve_validator_token
592            .checked_sub(amount_out_u128)
593            .ok_or(TempoPrecompileError::under_overflow())?;
594
595        self.pools[pool_id].write(pool)?;
596
597        Ok(amount_out)
598    }
599
600    /// Returns the total supply of LP tokens for the given pool.
601    pub fn get_total_supply(&self, pool_id: B256) -> Result<U256> {
602        self.total_supply[pool_id].read()
603    }
604
605    /// Set total supply of LP tokens for a pool
606    fn set_total_supply(&mut self, pool_id: B256, total_supply: U256) -> Result<()> {
607        self.total_supply[pool_id].write(total_supply)
608    }
609
610    /// Returns the LP token balance for `user` in the given pool.
611    pub fn get_liquidity_balances(&self, pool_id: B256, user: Address) -> Result<U256> {
612        self.liquidity_balances[pool_id][user].read()
613    }
614
615    /// Set user's LP token balance
616    fn set_liquidity_balances(
617        &mut self,
618        pool_id: B256,
619        user: Address,
620        balance: U256,
621    ) -> Result<()> {
622        self.liquidity_balances[pool_id][user].write(balance)
623    }
624}
625
626#[cfg(test)]
627mod tests {
628    use alloy::primitives::Address;
629    use tempo_chainspec::hardfork::TempoHardfork;
630    use tempo_contracts::precompiles::TIP20Error;
631
632    use super::*;
633    use crate::{
634        error::TempoPrecompileError,
635        storage::{ContractStorage, StorageCtx, hashmap::HashMapStorageProvider},
636        test_util::TIP20Setup,
637        tip_fee_manager::TIPFeeAMMError,
638    };
639
640    /// Integer square root using the Babylonian method
641    fn sqrt(x: U256) -> U256 {
642        if x == U256::ZERO {
643            return U256::ZERO;
644        }
645        let mut z = (x + U256::ONE) / uint!(2_U256);
646        let mut y = x;
647        while z < y {
648            y = z;
649            z = (x / z + z) / uint!(2_U256);
650        }
651        y
652    }
653
654    /// Sets up a pool with initial liquidity for testing
655    fn setup_pool_with_liquidity(
656        amm: &mut TipFeeManager,
657        user_token: Address,
658        validator_token: Address,
659        user_amount: U256,
660        validator_amount: U256,
661    ) -> Result<B256> {
662        let pool_id = amm.pool_id(user_token, validator_token);
663        let pool = Pool {
664            reserve_user_token: user_amount.try_into().unwrap(),
665            reserve_validator_token: validator_amount.try_into().unwrap(),
666        };
667        amm.pools[pool_id].write(pool)?;
668        let liquidity = sqrt(user_amount * validator_amount);
669        amm.total_supply[pool_id].write(liquidity)?;
670        Ok(pool_id)
671    }
672
673    #[test]
674    fn test_mint_identical_addresses() -> eyre::Result<()> {
675        let mut storage = HashMapStorageProvider::new(1);
676        let admin = Address::random();
677        StorageCtx::enter(&mut storage, || {
678            let token = TIP20Setup::create("Test", "TST", admin).apply()?;
679            let mut amm = TipFeeManager::new();
680            let result = amm.mint(
681                admin,
682                token.address(),
683                token.address(),
684                U256::from(1000),
685                admin,
686            );
687            assert!(matches!(
688                result,
689                Err(TempoPrecompileError::TIPFeeAMMError(
690                    TIPFeeAMMError::IdenticalAddresses(_)
691                ))
692            ));
693            Ok(())
694        })
695    }
696
697    #[test]
698    fn test_burn_identical_addresses() -> eyre::Result<()> {
699        let mut storage = HashMapStorageProvider::new(1);
700        let admin = Address::random();
701        StorageCtx::enter(&mut storage, || {
702            let token = TIP20Setup::create("Test", "TST", admin).apply()?;
703            let mut amm = TipFeeManager::new();
704            let result = amm.burn(
705                admin,
706                token.address(),
707                token.address(),
708                U256::from(1000),
709                admin,
710            );
711            assert!(matches!(
712                result,
713                Err(TempoPrecompileError::TIPFeeAMMError(
714                    TIPFeeAMMError::IdenticalAddresses(_)
715                ))
716            ));
717            Ok(())
718        })
719    }
720
721    #[test]
722    fn test_rebalance_swap_insufficient_funds() -> eyre::Result<()> {
723        let mut storage = HashMapStorageProvider::new(1);
724        let admin = Address::random();
725        let to = Address::random();
726        StorageCtx::enter(&mut storage, || {
727            let user_token = TIP20Setup::create("UserToken", "UTK", admin).apply()?;
728            let validator_token = TIP20Setup::create("ValidatorToken", "VTK", admin).apply()?;
729
730            let mut amm = TipFeeManager::new();
731            let amount = uint!(100000_U256) * uint!(10_U256).pow(U256::from(6));
732            setup_pool_with_liquidity(
733                &mut amm,
734                user_token.address(),
735                validator_token.address(),
736                amount,
737                amount,
738            )?;
739
740            let result = amm.rebalance_swap(
741                admin,
742                user_token.address(),
743                validator_token.address(),
744                amount + U256::ONE,
745                to,
746            );
747            assert!(matches!(
748                result,
749                Err(TempoPrecompileError::TIPFeeAMMError(
750                    TIPFeeAMMError::InvalidAmount(_)
751                ))
752            ));
753            Ok(())
754        })
755    }
756
757    #[test]
758    fn test_mint_rejects_non_usd_user_token() -> eyre::Result<()> {
759        let mut storage = HashMapStorageProvider::new(1);
760        let admin = Address::random();
761        StorageCtx::enter(&mut storage, || {
762            let eur_token = TIP20Setup::create("EuroToken", "EUR", admin)
763                .currency("EUR")
764                .apply()?;
765            let usd_token = TIP20Setup::create("USDToken", "USD", admin).apply()?;
766            let mut amm = TipFeeManager::new();
767
768            let result = amm.mint(
769                admin,
770                eur_token.address(),
771                usd_token.address(),
772                U256::from(1000),
773                admin,
774            );
775            assert!(matches!(
776                result,
777                Err(TempoPrecompileError::TIP20(TIP20Error::InvalidCurrency(_)))
778            ));
779
780            let result = amm.mint(
781                admin,
782                usd_token.address(),
783                eur_token.address(),
784                U256::from(1000),
785                admin,
786            );
787            assert!(matches!(
788                result,
789                Err(TempoPrecompileError::TIP20(TIP20Error::InvalidCurrency(_)))
790            ));
791            Ok(())
792        })
793    }
794
795    #[test]
796    fn test_burn_rejects_non_usd_tokens() -> eyre::Result<()> {
797        let mut storage = HashMapStorageProvider::new(1);
798        let admin = Address::random();
799        StorageCtx::enter(&mut storage, || {
800            let eur_token = TIP20Setup::create("EuroToken", "EUR", admin)
801                .currency("EUR")
802                .apply()?;
803            let usd_token = TIP20Setup::create("USDToken", "USD", admin).apply()?;
804            let mut amm = TipFeeManager::new();
805
806            let result = amm.burn(
807                admin,
808                eur_token.address(),
809                usd_token.address(),
810                U256::from(1000),
811                admin,
812            );
813            assert!(matches!(
814                result,
815                Err(TempoPrecompileError::TIP20(TIP20Error::InvalidCurrency(_)))
816            ));
817
818            let result = amm.burn(
819                admin,
820                usd_token.address(),
821                eur_token.address(),
822                U256::from(1000),
823                admin,
824            );
825            assert!(matches!(
826                result,
827                Err(TempoPrecompileError::TIP20(TIP20Error::InvalidCurrency(_)))
828            ));
829            Ok(())
830        })
831    }
832
833    #[test]
834    fn test_mint_insufficient_amount() -> eyre::Result<()> {
835        let mut storage = HashMapStorageProvider::new(1);
836        let admin = Address::random();
837        StorageCtx::enter(&mut storage, || {
838            let user_token = TIP20Setup::create("UserToken", "UTK", admin).apply()?;
839            let validator_token = TIP20Setup::create("ValidatorToken", "VTK", admin).apply()?;
840            let mut amm = TipFeeManager::new();
841
842            // MIN_LIQUIDITY = 1000, amount/2 must be > 1000, so 2000 should fail
843            let insufficient = uint!(2000_U256);
844            let result = amm.mint(
845                admin,
846                user_token.address(),
847                validator_token.address(),
848                insufficient,
849                admin,
850            );
851            assert!(matches!(
852                result,
853                Err(TempoPrecompileError::TIPFeeAMMError(
854                    TIPFeeAMMError::InsufficientLiquidity(_)
855                ))
856            ));
857            Ok(())
858        })
859    }
860
861    #[test]
862    fn test_add_liquidity() -> eyre::Result<()> {
863        let mut storage = HashMapStorageProvider::new(1);
864        let admin = Address::random();
865
866        StorageCtx::enter(&mut storage, || {
867            let mint_amount = uint!(10000000_U256);
868            let token1 = TIP20Setup::create("Token1", "TK1", admin)
869                .with_issuer(admin)
870                .with_mint(admin, mint_amount)
871                .apply()?
872                .address();
873            let token2 = TIP20Setup::create("Token2", "TK2", admin)
874                .with_issuer(admin)
875                .with_mint(admin, mint_amount)
876                .apply()?
877                .address();
878
879            let mut amm = TipFeeManager::new();
880            let amount = uint!(10000_U256);
881            let result = amm.mint(admin, token1, token2, amount, admin)?;
882            let expected_mean = amount / uint!(2_U256);
883            let expected_liquidity = expected_mean - MIN_LIQUIDITY;
884
885            assert_eq!(result, expected_liquidity,);
886
887            Ok(())
888        })
889    }
890
891    #[test]
892    fn test_calculate_burn_amounts() -> eyre::Result<()> {
893        let mut storage = HashMapStorageProvider::new(1);
894
895        StorageCtx::enter(&mut storage, || {
896            let mut amm = TipFeeManager::new();
897
898            let pool = Pool {
899                reserve_user_token: 1000,
900                reserve_validator_token: 1000,
901            };
902            let pool_id = B256::ZERO;
903            amm.set_total_supply(pool_id, uint!(1000000000000000_U256))?;
904
905            let liquidity = uint!(1_U256);
906            let result = amm.calculate_burn_amounts(&pool, pool_id, liquidity);
907
908            assert!(result.is_ok());
909            let (amount_user, amount_validator) = result?;
910            assert_eq!(amount_user, U256::ZERO);
911            assert_eq!(amount_validator, U256::ZERO);
912
913            Ok(())
914        })
915    }
916
917    /// Test execute_fee_swap executes swap immediately and updates reserves
918    #[test]
919    fn test_execute_fee_swap_immediate() -> eyre::Result<()> {
920        let mut storage = HashMapStorageProvider::new(1);
921        let admin = Address::random();
922
923        StorageCtx::enter(&mut storage, || {
924            let user_token = TIP20Setup::create("UserToken", "UTK", admin)
925                .apply()?
926                .address();
927            let validator_token = TIP20Setup::create("ValidatorToken", "VTK", admin)
928                .apply()?
929                .address();
930
931            let mut amm = TipFeeManager::new();
932
933            // Setup pool with 1000 tokens each
934            let liquidity_amount = uint!(1000_U256);
935            let pool_id = setup_pool_with_liquidity(
936                &mut amm,
937                user_token,
938                validator_token,
939                liquidity_amount,
940                liquidity_amount,
941            )?;
942
943            // Execute fee swap for 100 tokens
944            let amount_in = uint!(100_U256);
945            let expected_out = (amount_in * M) / SCALE; // 100 * 9970 / 10000 = 99
946
947            let amount_out = amm.execute_fee_swap(user_token, validator_token, amount_in)?;
948
949            assert_eq!(amount_out, expected_out);
950
951            // Verify reserves updated immediately
952            let pool = amm.pools[pool_id].read()?;
953            assert_eq!(
954                U256::from(pool.reserve_user_token),
955                liquidity_amount + amount_in
956            );
957            assert_eq!(
958                U256::from(pool.reserve_validator_token),
959                liquidity_amount - expected_out
960            );
961
962            Ok(())
963        })
964    }
965
966    /// Test execute_fee_swap fails with insufficient liquidity
967    #[test]
968    fn test_execute_fee_swap_insufficient_liquidity() -> eyre::Result<()> {
969        let mut storage = HashMapStorageProvider::new(1);
970        let admin = Address::random();
971
972        StorageCtx::enter(&mut storage, || {
973            let user_token = TIP20Setup::create("UserToken", "UTK", admin)
974                .apply()?
975                .address();
976            let validator_token = TIP20Setup::create("ValidatorToken", "VTK", admin)
977                .apply()?
978                .address();
979
980            let mut amm = TipFeeManager::new();
981
982            // Setup pool with only 100 tokens each
983            let small_liquidity = uint!(100_U256);
984            setup_pool_with_liquidity(
985                &mut amm,
986                user_token,
987                validator_token,
988                small_liquidity,
989                small_liquidity,
990            )?;
991
992            // Try to swap 200 tokens (would need ~199 output, but only 100 available)
993            let too_large_amount = uint!(200_U256);
994
995            let result = amm.execute_fee_swap(user_token, validator_token, too_large_amount);
996
997            assert!(matches!(
998                result,
999                Err(TempoPrecompileError::TIPFeeAMMError(
1000                    TIPFeeAMMError::InsufficientLiquidity(_)
1001                ))
1002            ));
1003
1004            Ok(())
1005        })
1006    }
1007
1008    /// Test fee swap rounding consistency across multiple swaps
1009    #[test]
1010    fn test_fee_swap_rounding_consistency() -> eyre::Result<()> {
1011        let mut storage = HashMapStorageProvider::new(1);
1012        let admin = Address::random();
1013
1014        StorageCtx::enter(&mut storage, || {
1015            let user_token = TIP20Setup::create("UserToken", "UTK", admin)
1016                .apply()?
1017                .address();
1018            let validator_token = TIP20Setup::create("ValidatorToken", "VTK", admin)
1019                .apply()?
1020                .address();
1021
1022            let mut amm = TipFeeManager::new();
1023            let liquidity = uint!(100000_U256) * uint!(10_U256).pow(U256::from(6));
1024            let pool_id = setup_pool_with_liquidity(
1025                &mut amm,
1026                user_token,
1027                validator_token,
1028                liquidity,
1029                liquidity,
1030            )?;
1031
1032            let amount_in = uint!(10000_U256) * uint!(10_U256).pow(U256::from(6));
1033            let expected_out = (amount_in * M) / SCALE;
1034
1035            let actual_out = amm.execute_fee_swap(user_token, validator_token, amount_in)?;
1036            assert_eq!(actual_out, expected_out, "Output should match expected");
1037
1038            let pool = amm.pools[pool_id].read()?;
1039            assert_eq!(
1040                U256::from(pool.reserve_user_token),
1041                liquidity + amount_in,
1042                "User reserve should increase"
1043            );
1044            assert_eq!(
1045                U256::from(pool.reserve_validator_token),
1046                liquidity - actual_out,
1047                "Validator reserve should decrease"
1048            );
1049
1050            Ok(())
1051        })
1052    }
1053
1054    /// Test multiple consecutive fee swaps update reserves correctly
1055    #[test]
1056    fn test_multiple_consecutive_fee_swaps() -> eyre::Result<()> {
1057        let mut storage = HashMapStorageProvider::new(1);
1058        let admin = Address::random();
1059
1060        StorageCtx::enter(&mut storage, || {
1061            let user_token = TIP20Setup::create("UserToken", "UTK", admin)
1062                .apply()?
1063                .address();
1064            let validator_token = TIP20Setup::create("ValidatorToken", "VTK", admin)
1065                .apply()?
1066                .address();
1067
1068            let mut amm = TipFeeManager::new();
1069            let initial = uint!(100000_U256) * uint!(10_U256).pow(U256::from(6));
1070            let pool_id =
1071                setup_pool_with_liquidity(&mut amm, user_token, validator_token, initial, initial)?;
1072
1073            let swap1 = uint!(1000_U256) * uint!(10_U256).pow(U256::from(6));
1074            let swap2 = uint!(2000_U256) * uint!(10_U256).pow(U256::from(6));
1075            let swap3 = uint!(3000_U256) * uint!(10_U256).pow(U256::from(6));
1076
1077            let out1 = amm.execute_fee_swap(user_token, validator_token, swap1)?;
1078            let out2 = amm.execute_fee_swap(user_token, validator_token, swap2)?;
1079            let out3 = amm.execute_fee_swap(user_token, validator_token, swap3)?;
1080
1081            let total_in = swap1 + swap2 + swap3;
1082            let total_out = out1 + out2 + out3;
1083
1084            // Each swap output should be amount_in * M / SCALE
1085            assert_eq!(out1, (swap1 * M) / SCALE);
1086            assert_eq!(out2, (swap2 * M) / SCALE);
1087            assert_eq!(out3, (swap3 * M) / SCALE);
1088
1089            let pool = amm.pools[pool_id].read()?;
1090            assert_eq!(U256::from(pool.reserve_user_token), initial + total_in);
1091            assert_eq!(
1092                U256::from(pool.reserve_validator_token),
1093                initial - total_out
1094            );
1095
1096            Ok(())
1097        })
1098    }
1099
1100    /// Test pool boundary condition
1101    #[test]
1102    fn test_pool_liquidity_boundary() -> eyre::Result<()> {
1103        let mut storage = HashMapStorageProvider::new(1);
1104        let admin = Address::random();
1105
1106        StorageCtx::enter(&mut storage, || {
1107            let user_token = TIP20Setup::create("UserToken", "UTK", admin)
1108                .apply()?
1109                .address();
1110            let validator_token = TIP20Setup::create("ValidatorToken", "VTK", admin)
1111                .apply()?
1112                .address();
1113
1114            let mut amm = TipFeeManager::new();
1115            let liquidity = uint!(100_U256) * uint!(10_U256).pow(U256::from(6));
1116            let pool_id = setup_pool_with_liquidity(
1117                &mut amm,
1118                user_token,
1119                validator_token,
1120                liquidity,
1121                liquidity,
1122            )?;
1123
1124            let reserve = U256::from(amm.pools[pool_id].read()?.reserve_validator_token);
1125
1126            // Exactly at boundary should succeed (100 * 0.997 = 99.7, which is < 100)
1127            let ok_amount = uint!(100_U256) * uint!(10_U256).pow(U256::from(6));
1128            assert!(reserve >= compute_amount_out(ok_amount)?);
1129
1130            // Just over boundary should fail (101 * 0.997 = 100.697, which is > 100)
1131            let too_much = uint!(101_U256) * uint!(10_U256).pow(U256::from(6));
1132            assert!(reserve < compute_amount_out(too_much)?);
1133
1134            Ok(())
1135        })
1136    }
1137
1138    /// Test zero liquidity burn
1139    #[test]
1140    fn test_burn_zero_liquidity() -> eyre::Result<()> {
1141        let mut storage = HashMapStorageProvider::new(1);
1142        let admin = Address::random();
1143
1144        StorageCtx::enter(&mut storage, || {
1145            let user_token = TIP20Setup::create("UserToken", "UTK", admin)
1146                .apply()?
1147                .address();
1148            let validator_token = TIP20Setup::create("ValidatorToken", "VTK", admin)
1149                .apply()?
1150                .address();
1151
1152            let mut amm = TipFeeManager::new();
1153
1154            let result = amm.burn(admin, user_token, validator_token, U256::ZERO, admin);
1155
1156            assert!(matches!(
1157                result,
1158                Err(TempoPrecompileError::TIPFeeAMMError(
1159                    TIPFeeAMMError::InvalidAmount(_)
1160                ))
1161            ));
1162
1163            Ok(())
1164        })
1165    }
1166
1167    /// Test zero amount validator token
1168    #[test]
1169    fn test_mint_zero_amount_validator_token() -> eyre::Result<()> {
1170        let mut storage = HashMapStorageProvider::new(1);
1171        let admin = Address::random();
1172
1173        StorageCtx::enter(&mut storage, || {
1174            let user_token = TIP20Setup::create("UserToken", "UTK", admin)
1175                .apply()?
1176                .address();
1177            let validator_token = TIP20Setup::create("ValidatorToken", "VTK", admin)
1178                .apply()?
1179                .address();
1180
1181            let mut amm = TipFeeManager::new();
1182
1183            let result = amm.mint(admin, user_token, validator_token, U256::ZERO, admin);
1184
1185            assert!(matches!(
1186                result,
1187                Err(TempoPrecompileError::TIPFeeAMMError(
1188                    TIPFeeAMMError::InvalidAmount(_)
1189                ))
1190            ));
1191
1192            Ok(())
1193        })
1194    }
1195
1196    #[test]
1197    fn test_rebalance_swap() -> eyre::Result<()> {
1198        let mut storage = HashMapStorageProvider::new(1);
1199        let admin = Address::random();
1200        let recipient = Address::random();
1201
1202        StorageCtx::enter(&mut storage, || {
1203            let mint_amount = uint!(10000000_U256);
1204            let mut amm = TipFeeManager::new();
1205            let amm_address = amm.address;
1206
1207            let user_token = TIP20Setup::create("UserToken", "UTK", admin)
1208                .with_issuer(admin)
1209                .with_mint(admin, mint_amount)
1210                .with_mint(amm_address, mint_amount)
1211                .apply()?
1212                .address();
1213            let validator_token = TIP20Setup::create("ValidatorToken", "VTK", admin)
1214                .with_issuer(admin)
1215                .with_mint(admin, mint_amount)
1216                .apply()?
1217                .address();
1218
1219            let liquidity = uint!(100000_U256);
1220            let pool_id = setup_pool_with_liquidity(
1221                &mut amm,
1222                user_token,
1223                validator_token,
1224                liquidity,
1225                liquidity,
1226            )?;
1227
1228            let amount_out = uint!(1000_U256);
1229            let expected_in = (amount_out * N) / SCALE + U256::ONE;
1230
1231            let amount_in =
1232                amm.rebalance_swap(admin, user_token, validator_token, amount_out, recipient)?;
1233
1234            assert_eq!(amount_in, expected_in);
1235
1236            let pool = amm.pools[pool_id].read()?;
1237            assert_eq!(U256::from(pool.reserve_user_token), liquidity - amount_out);
1238            assert_eq!(
1239                U256::from(pool.reserve_validator_token),
1240                liquidity + amount_in
1241            );
1242
1243            Ok(())
1244        })
1245    }
1246
1247    #[test]
1248    fn test_mint_subsequent_deposit() -> eyre::Result<()> {
1249        let mut storage = HashMapStorageProvider::new(1);
1250        let admin = Address::random();
1251        let second_user = Address::random();
1252
1253        StorageCtx::enter(&mut storage, || {
1254            let mint_amount = uint!(100000000_U256);
1255            let user_token = TIP20Setup::create("UserToken", "UTK", admin)
1256                .with_issuer(admin)
1257                .with_mint(admin, mint_amount)
1258                .with_mint(second_user, mint_amount)
1259                .apply()?
1260                .address();
1261            let validator_token = TIP20Setup::create("ValidatorToken", "VTK", admin)
1262                .with_issuer(admin)
1263                .with_mint(admin, mint_amount)
1264                .with_mint(second_user, mint_amount)
1265                .apply()?
1266                .address();
1267
1268            let mut amm = TipFeeManager::new();
1269
1270            let initial_amount = uint!(100000_U256);
1271            let first_liquidity =
1272                amm.mint(admin, user_token, validator_token, initial_amount, admin)?;
1273
1274            let expected_first_liquidity = initial_amount / uint!(2_U256) - MIN_LIQUIDITY;
1275            assert_eq!(first_liquidity, expected_first_liquidity);
1276
1277            let pool_id = amm.pool_id(user_token, validator_token);
1278            let total_supply_after_first = amm.get_total_supply(pool_id)?;
1279            assert_eq!(total_supply_after_first, first_liquidity + MIN_LIQUIDITY);
1280
1281            let pool_after_first = amm.pools[pool_id].read()?;
1282            let reserve_val = U256::from(pool_after_first.reserve_validator_token);
1283
1284            let second_amount = uint!(50000_U256);
1285            let second_liquidity = amm.mint(
1286                second_user,
1287                user_token,
1288                validator_token,
1289                second_amount,
1290                second_user,
1291            )?;
1292
1293            let expected_second_liquidity = second_amount * total_supply_after_first / reserve_val;
1294            assert_eq!(second_liquidity, expected_second_liquidity);
1295
1296            let total_supply_after_second = amm.get_total_supply(pool_id)?;
1297            assert_eq!(
1298                total_supply_after_second,
1299                total_supply_after_first + second_liquidity
1300            );
1301
1302            let admin_balance = amm.get_liquidity_balances(pool_id, admin)?;
1303            let second_user_balance = amm.get_liquidity_balances(pool_id, second_user)?;
1304            assert_eq!(admin_balance, first_liquidity);
1305            assert_eq!(second_user_balance, second_liquidity);
1306
1307            Ok(())
1308        })
1309    }
1310
1311    #[test]
1312    fn test_burn() -> eyre::Result<()> {
1313        let mut storage = HashMapStorageProvider::new(1);
1314        let admin = Address::random();
1315        let recipient = Address::random();
1316
1317        StorageCtx::enter(&mut storage, || {
1318            let mint_amount = uint!(100000000_U256);
1319            let user_token = TIP20Setup::create("UserToken", "UTK", admin)
1320                .with_issuer(admin)
1321                .with_mint(admin, mint_amount)
1322                .apply()?
1323                .address();
1324            let validator_token = TIP20Setup::create("ValidatorToken", "VTK", admin)
1325                .with_issuer(admin)
1326                .with_mint(admin, mint_amount)
1327                .apply()?
1328                .address();
1329
1330            let mut amm = TipFeeManager::new();
1331
1332            let deposit_amount = uint!(100000_U256);
1333            let liquidity = amm.mint(admin, user_token, validator_token, deposit_amount, admin)?;
1334
1335            let expected_liquidity = deposit_amount / uint!(2_U256) - MIN_LIQUIDITY;
1336            assert_eq!(liquidity, expected_liquidity);
1337
1338            let pool_id = amm.pool_id(user_token, validator_token);
1339            let pool_before = amm.pools[pool_id].read()?;
1340            let total_supply_before = amm.get_total_supply(pool_id)?;
1341
1342            let burn_amount = liquidity / uint!(2_U256);
1343            let (amount_user, amount_validator) =
1344                amm.burn(admin, user_token, validator_token, burn_amount, recipient)?;
1345
1346            let expected_user =
1347                burn_amount * U256::from(pool_before.reserve_user_token) / total_supply_before;
1348            let expected_validator =
1349                burn_amount * U256::from(pool_before.reserve_validator_token) / total_supply_before;
1350            assert_eq!(amount_user, expected_user);
1351            assert_eq!(amount_validator, expected_validator);
1352
1353            let pool_after = amm.pools[pool_id].read()?;
1354            let total_supply_after = amm.get_total_supply(pool_id)?;
1355
1356            assert_eq!(total_supply_after, total_supply_before - burn_amount);
1357
1358            let admin_balance = amm.get_liquidity_balances(pool_id, admin)?;
1359            assert_eq!(admin_balance, liquidity - burn_amount);
1360
1361            assert_eq!(
1362                U256::from(pool_after.reserve_user_token),
1363                U256::from(pool_before.reserve_user_token) - amount_user
1364            );
1365            assert_eq!(
1366                U256::from(pool_after.reserve_validator_token),
1367                U256::from(pool_before.reserve_validator_token) - amount_validator
1368            );
1369
1370            Ok(())
1371        })
1372    }
1373
1374    #[test]
1375    fn test_burn_insufficient_balance() -> eyre::Result<()> {
1376        let mut storage = HashMapStorageProvider::new(1);
1377        let admin = Address::random();
1378        let other_user = Address::random();
1379
1380        StorageCtx::enter(&mut storage, || {
1381            let mint_amount = uint!(100000000_U256);
1382            let user_token = TIP20Setup::create("UserToken", "UTK", admin)
1383                .with_issuer(admin)
1384                .with_mint(admin, mint_amount)
1385                .apply()?
1386                .address();
1387            let validator_token = TIP20Setup::create("ValidatorToken", "VTK", admin)
1388                .with_issuer(admin)
1389                .with_mint(admin, mint_amount)
1390                .apply()?
1391                .address();
1392
1393            let mut amm = TipFeeManager::new();
1394
1395            let deposit_amount = uint!(100000_U256);
1396            let liquidity = amm.mint(admin, user_token, validator_token, deposit_amount, admin)?;
1397
1398            let result = amm.burn(
1399                other_user,
1400                user_token,
1401                validator_token,
1402                liquidity,
1403                other_user,
1404            );
1405
1406            assert!(matches!(
1407                result,
1408                Err(TempoPrecompileError::TIPFeeAMMError(
1409                    TIPFeeAMMError::InsufficientLiquidity(_)
1410                ))
1411            ));
1412
1413            Ok(())
1414        })
1415    }
1416
1417    // Test zero amount rebalance swap
1418    #[test]
1419    fn test_rebalance_swap_zero_amount_out() -> eyre::Result<()> {
1420        let mut storage = HashMapStorageProvider::new(1);
1421        let admin = Address::random();
1422        let to = Address::random();
1423
1424        StorageCtx::enter(&mut storage, || {
1425            let user_token = TIP20Setup::create("UserToken", "UTK", admin)
1426                .apply()?
1427                .address();
1428            let validator_token = TIP20Setup::create("ValidatorToken", "VTK", admin)
1429                .apply()?
1430                .address();
1431
1432            let mut amm = TipFeeManager::new();
1433
1434            let result = amm.rebalance_swap(admin, user_token, validator_token, U256::ZERO, to);
1435
1436            assert!(matches!(
1437                result,
1438                Err(TempoPrecompileError::TIPFeeAMMError(
1439                    TIPFeeAMMError::InvalidAmount(_)
1440                ))
1441            ));
1442
1443            Ok(())
1444        })
1445    }
1446
1447    #[test]
1448    fn test_t1c_reserve_pool_liquidity() -> eyre::Result<()> {
1449        let mut storage = HashMapStorageProvider::new_with_spec(1, TempoHardfork::T1C);
1450        let admin = Address::random();
1451
1452        StorageCtx::enter(&mut storage, || {
1453            let mint_amount = uint!(10000000_U256);
1454            let user_token = TIP20Setup::create("UserToken", "UTK", admin)
1455                .with_issuer(admin)
1456                .with_mint(admin, mint_amount)
1457                .apply()?
1458                .address();
1459            let validator_token = TIP20Setup::create("ValidatorToken", "VTK", admin)
1460                .with_issuer(admin)
1461                .with_mint(admin, mint_amount)
1462                .apply()?
1463                .address();
1464
1465            let mut amm = TipFeeManager::new();
1466            let liquidity = uint!(100000_U256);
1467            let pool_id = setup_pool_with_liquidity(
1468                &mut amm,
1469                user_token,
1470                validator_token,
1471                liquidity,
1472                liquidity,
1473            )?;
1474
1475            let max_amount = uint!(10000_U256);
1476            let amount_out: u128 = compute_amount_out(max_amount)?.try_into().unwrap();
1477            amm.reserve_pool_liquidity(pool_id, amount_out)?;
1478
1479            let reserved = amm.pending_fee_swap_reservation[pool_id].t_read()?;
1480            let expected_reserved: u128 = compute_amount_out(max_amount)?.try_into().unwrap();
1481            assert_eq!(reserved, expected_reserved);
1482
1483            Ok(())
1484        })
1485    }
1486
1487    #[test]
1488    fn test_t1c_burn_respects_reservation() -> eyre::Result<()> {
1489        let mut storage = HashMapStorageProvider::new_with_spec(1, TempoHardfork::T1C);
1490        let admin = Address::random();
1491        let recipient = Address::random();
1492
1493        StorageCtx::enter(&mut storage, || {
1494            let mint_amount = uint!(100000000_U256);
1495            let user_token = TIP20Setup::create("UserToken", "UTK", admin)
1496                .with_issuer(admin)
1497                .with_mint(admin, mint_amount)
1498                .apply()?
1499                .address();
1500            let validator_token = TIP20Setup::create("ValidatorToken", "VTK", admin)
1501                .with_issuer(admin)
1502                .with_mint(admin, mint_amount)
1503                .apply()?
1504                .address();
1505
1506            let mut amm = TipFeeManager::new();
1507
1508            let deposit_amount = uint!(100000_U256);
1509            let liquidity = amm.mint(admin, user_token, validator_token, deposit_amount, admin)?;
1510
1511            let pool_id = amm.pool_id(user_token, validator_token);
1512            let pool = amm.pools[pool_id].read()?;
1513
1514            // Reserve most of the validator token liquidity
1515            let reserve_amount = U256::from(pool.reserve_validator_token) - uint!(100_U256);
1516            let amount_out: u128 = compute_amount_out(reserve_amount)?.try_into().unwrap();
1517            amm.reserve_pool_liquidity(pool_id, amount_out)?;
1518
1519            let result = amm.burn(admin, user_token, validator_token, liquidity, recipient);
1520            assert!(matches!(
1521                result,
1522                Err(TempoPrecompileError::TIPFeeAMMError(
1523                    TIPFeeAMMError::InsufficientLiquidity(_)
1524                ))
1525            ));
1526
1527            Ok(())
1528        })
1529    }
1530
1531    #[test]
1532    fn test_t1c_partial_burn_with_reservation() -> eyre::Result<()> {
1533        let mut storage = HashMapStorageProvider::new_with_spec(1, TempoHardfork::T1C);
1534        let admin = Address::random();
1535        let recipient = Address::random();
1536
1537        StorageCtx::enter(&mut storage, || {
1538            let mint_amount = uint!(100000000_U256);
1539            let user_token = TIP20Setup::create("UserToken", "UTK", admin)
1540                .with_issuer(admin)
1541                .with_mint(admin, mint_amount)
1542                .apply()?
1543                .address();
1544            let validator_token = TIP20Setup::create("ValidatorToken", "VTK", admin)
1545                .with_issuer(admin)
1546                .with_mint(admin, mint_amount)
1547                .apply()?
1548                .address();
1549
1550            let mut amm = TipFeeManager::new();
1551
1552            let deposit_amount = uint!(100000_U256);
1553            let liquidity = amm.mint(admin, user_token, validator_token, deposit_amount, admin)?;
1554
1555            let pool_id = amm.pool_id(user_token, validator_token);
1556            let small_reserve = uint!(1000_U256);
1557            let amount_out: u128 = compute_amount_out(small_reserve)?.try_into().unwrap();
1558            amm.reserve_pool_liquidity(pool_id, amount_out)?;
1559
1560            let small_burn = liquidity / uint!(10_U256);
1561            let result = amm.burn(admin, user_token, validator_token, small_burn, recipient);
1562
1563            assert!(result.is_ok());
1564
1565            Ok(())
1566        })
1567    }
1568
1569    #[test]
1570    fn test_t1c_rebalance_swap_respects_reservation() -> eyre::Result<()> {
1571        let mut storage = HashMapStorageProvider::new_with_spec(1, TempoHardfork::T1C);
1572        let admin = Address::random();
1573        let to = Address::random();
1574
1575        StorageCtx::enter(&mut storage, || {
1576            let mint_amount = uint!(100000000_U256);
1577            let mut amm = TipFeeManager::new();
1578            let amm_address = amm.address;
1579            let user_token = TIP20Setup::create("UserToken", "UTK", admin)
1580                .with_issuer(admin)
1581                .with_mint(admin, mint_amount)
1582                .with_mint(amm_address, mint_amount)
1583                .apply()?
1584                .address();
1585            let validator_token = TIP20Setup::create("ValidatorToken", "VTK", admin)
1586                .with_issuer(admin)
1587                .with_mint(admin, mint_amount)
1588                .apply()?
1589                .address();
1590
1591            let liq = uint!(100000_U256);
1592            let pool_id =
1593                setup_pool_with_liquidity(&mut amm, user_token, validator_token, liq, liq)?;
1594
1595            let amount_out: u128 = compute_amount_out(uint!(50000_U256))?.try_into().unwrap();
1596            amm.reserve_pool_liquidity(pool_id, amount_out)?;
1597
1598            amm.rebalance_swap(admin, user_token, validator_token, uint!(5000_U256), to)?;
1599            let pool = amm.pools[pool_id].read()?;
1600            let reserved = amm.pending_fee_swap_reservation[pool_id].t_read()?;
1601            assert!(pool.reserve_validator_token >= reserved);
1602
1603            Ok(())
1604        })
1605    }
1606
1607    #[test]
1608    fn test_pre_t1c_rebalance_swap_skips_reservation() -> eyre::Result<()> {
1609        let mut storage = HashMapStorageProvider::new_with_spec(1, TempoHardfork::T1B);
1610        let admin = Address::random();
1611        let to = Address::random();
1612
1613        StorageCtx::enter(&mut storage, || {
1614            let mint_amount = uint!(100000000_U256);
1615            let mut amm = TipFeeManager::new();
1616            let amm_address = amm.address;
1617            let user_token = TIP20Setup::create("UserToken", "UTK", admin)
1618                .with_issuer(admin)
1619                .with_mint(admin, mint_amount)
1620                .with_mint(amm_address, mint_amount)
1621                .apply()?
1622                .address();
1623            let validator_token = TIP20Setup::create("ValidatorToken", "VTK", admin)
1624                .with_issuer(admin)
1625                .with_mint(admin, mint_amount)
1626                .apply()?
1627                .address();
1628
1629            let liq = uint!(100000_U256);
1630            setup_pool_with_liquidity(&mut amm, user_token, validator_token, liq, liq)?;
1631            assert!(
1632                amm.rebalance_swap(admin, user_token, validator_token, uint!(5000_U256), to)
1633                    .is_ok()
1634            );
1635
1636            Ok(())
1637        })
1638    }
1639
1640    #[test]
1641    fn test_pre_t1c_no_reservation() -> eyre::Result<()> {
1642        let mut storage = HashMapStorageProvider::new_with_spec(1, TempoHardfork::T1B);
1643        let admin = Address::random();
1644        let recipient = Address::random();
1645
1646        StorageCtx::enter(&mut storage, || {
1647            let mint_amount = uint!(100000000_U256);
1648            let user_token = TIP20Setup::create("UserToken", "UTK", admin)
1649                .with_issuer(admin)
1650                .with_mint(admin, mint_amount)
1651                .apply()?
1652                .address();
1653            let validator_token = TIP20Setup::create("ValidatorToken", "VTK", admin)
1654                .with_issuer(admin)
1655                .with_mint(admin, mint_amount)
1656                .apply()?
1657                .address();
1658
1659            let mut amm = TipFeeManager::new();
1660
1661            let deposit_amount = uint!(100000_U256);
1662            let liquidity = amm.mint(admin, user_token, validator_token, deposit_amount, admin)?;
1663
1664            let result = amm.burn(admin, user_token, validator_token, liquidity, recipient);
1665            assert!(result.is_ok());
1666
1667            Ok(())
1668        })
1669    }
1670}