tycho_simulation/evm/protocol/lido/
state.rs

1use std::{any::Any, collections::HashMap};
2
3use num_bigint::BigUint;
4use num_traits::Zero;
5use tycho_common::{
6    dto::ProtocolStateDelta,
7    models::token::Token,
8    simulation::{
9        errors::{SimulationError, TransitionError},
10        protocol_sim::{Balances, GetAmountOutResult, ProtocolSim},
11    },
12    Bytes,
13};
14
15use crate::evm::protocol::{
16    safe_math::{safe_add_u256, safe_div_u256, safe_mul_u256},
17    u256_num::{biguint_to_u256, u256_to_biguint, u256_to_f64},
18};
19
20#[derive(Debug, Clone, PartialEq, Eq)]
21pub enum LidoPoolType {
22    StEth,
23    WStEth,
24}
25
26#[derive(Clone, Copy, Debug, PartialEq, Eq)]
27pub enum StakingStatus {
28    Limited = 0,
29    Paused = 1,
30    Unlimited = 2,
31}
32
33impl StakingStatus {
34    pub fn as_str_name(&self) -> &'static str {
35        match self {
36            StakingStatus::Limited => "Limited",
37            StakingStatus::Paused => "Paused",
38            StakingStatus::Unlimited => "Unlimited",
39        }
40    }
41}
42
43// see here https://github.com/lidofinance/core/blob/cca04b42123735714d8c60a73c2f7af949e989db/contracts/0.4.24/lib/StakeLimitUtils.sol#L38
44#[derive(Debug, Clone, PartialEq, Eq)]
45pub struct StakeLimitState {
46    pub staking_status: StakingStatus,
47    pub staking_limit: BigUint,
48}
49
50impl StakeLimitState {
51    fn get_limit(&self) -> BigUint {
52        // https://github.com/lidofinance/core/blob/cca04b42123735714d8c60a73c2f7af949e989db/contracts/0.4.24/lib/StakeLimitUtils.sol#L98
53        match self.staking_status {
54            StakingStatus::Limited => self.staking_limit.clone(),
55            StakingStatus::Paused => BigUint::zero(),
56            StakingStatus::Unlimited => BigUint::from(u128::MAX),
57        }
58    }
59}
60
61pub const DEFAULT_GAS: u64 = 60000;
62
63#[derive(Clone, Debug, PartialEq, Eq)]
64pub struct LidoState {
65    pub pool_type: LidoPoolType,
66    pub total_shares: BigUint,
67    pub total_pooled_eth: BigUint,
68    pub total_wrapped_st_eth: Option<BigUint>,
69    pub id: Bytes,
70    pub native_address: Bytes,
71    pub stake_limits_state: StakeLimitState,
72    pub tokens: [Bytes; 2],
73    pub token_to_track_total_pooled_eth: Bytes,
74}
75
76impl LidoState {
77    fn steth_swap(&self, amount_in: BigUint) -> Result<GetAmountOutResult, SimulationError> {
78        let shares = safe_div_u256(
79            safe_mul_u256(biguint_to_u256(&amount_in), biguint_to_u256(&self.total_shares))?,
80            biguint_to_u256(&self.total_pooled_eth),
81        )?;
82
83        // During the lido submit step, the share count and total pooled ETH are updated,
84        // so all subsequent casts are calculated on the new values
85        let new_shares = safe_add_u256(biguint_to_u256(&self.total_shares.clone()), shares)?;
86        let new_total_pooled_eth = safe_add_u256(
87            biguint_to_u256(&self.total_pooled_eth.clone()),
88            biguint_to_u256(&amount_in),
89        )?;
90
91        // Add extra casts to mirror the additional token transfer performed by the tycho router.
92        let amount_from_shares = u256_to_biguint(safe_div_u256(
93            safe_mul_u256(shares, new_total_pooled_eth)?,
94            new_shares,
95        )?);
96
97        let shares_to_transfer = safe_div_u256(
98            safe_mul_u256(biguint_to_u256(&amount_from_shares), new_shares)?,
99            new_total_pooled_eth,
100        )?;
101
102        let amount_transferred = u256_to_biguint(safe_div_u256(
103            safe_mul_u256(shares_to_transfer, new_total_pooled_eth)?,
104            new_shares,
105        )?);
106
107        Ok(GetAmountOutResult {
108            amount: amount_transferred,
109            gas: BigUint::from(DEFAULT_GAS),
110            new_state: Box::new(Self {
111                pool_type: self.pool_type.clone(),
112                total_shares: u256_to_biguint(new_shares),
113                total_pooled_eth: u256_to_biguint(new_total_pooled_eth),
114                total_wrapped_st_eth: None,
115                id: self.id.clone(),
116                native_address: self.native_address.clone(),
117                stake_limits_state: self.stake_limits_state.clone(),
118                tokens: self.tokens.clone(),
119                token_to_track_total_pooled_eth: self
120                    .token_to_track_total_pooled_eth
121                    .clone(),
122            }),
123        })
124    }
125
126    fn wrap_steth(&self, amount_in: BigUint) -> Result<GetAmountOutResult, SimulationError> {
127        if amount_in.is_zero() {
128            return Err(SimulationError::InvalidInput("Cannot wrap 0 stETH ".to_string(), None))
129        }
130
131        let amount_out = u256_to_biguint(safe_div_u256(
132            safe_mul_u256(biguint_to_u256(&amount_in), biguint_to_u256(&self.total_shares))?,
133            biguint_to_u256(&self.total_pooled_eth),
134        )?);
135
136        let new_total_wrapped_st_eth = self
137            .total_wrapped_st_eth
138            .as_ref()
139            .expect("total_wrapped_st_eth must be present for wrapped staked ETH pool") +
140            &amount_out;
141
142        Ok(GetAmountOutResult {
143            amount: amount_out.clone(),
144            gas: BigUint::from(DEFAULT_GAS),
145            new_state: Box::new(Self {
146                pool_type: self.pool_type.clone(),
147                total_shares: self.total_shares.clone(),
148                total_pooled_eth: self.total_pooled_eth.clone(),
149                total_wrapped_st_eth: Some(new_total_wrapped_st_eth),
150                id: self.id.clone(),
151                native_address: self.native_address.clone(),
152                stake_limits_state: self.stake_limits_state.clone(),
153                tokens: self.tokens.clone(),
154                token_to_track_total_pooled_eth: self
155                    .token_to_track_total_pooled_eth
156                    .clone(),
157            }),
158        })
159    }
160
161    fn unwrap_steth(&self, amount_in: BigUint) -> Result<GetAmountOutResult, SimulationError> {
162        if amount_in.is_zero() {
163            return Err(SimulationError::InvalidInput("Cannot unwrap 0 wstETH ".to_string(), None))
164        }
165
166        let amount_out = u256_to_biguint(safe_div_u256(
167            safe_mul_u256(biguint_to_u256(&amount_in), biguint_to_u256(&self.total_pooled_eth))?,
168            biguint_to_u256(&self.total_shares),
169        )?);
170
171        let shares_from_amount_out = safe_div_u256(
172            safe_mul_u256(biguint_to_u256(&amount_out), biguint_to_u256(&self.total_shares))?,
173            biguint_to_u256(&self.total_pooled_eth),
174        )?;
175
176        let amount_transferred = u256_to_biguint(safe_div_u256(
177            safe_mul_u256(shares_from_amount_out, biguint_to_u256(&self.total_pooled_eth))?,
178            biguint_to_u256(&self.total_shares),
179        )?);
180
181        let new_total_wrapped_st_eth = self
182            .total_wrapped_st_eth
183            .as_ref()
184            .expect("total_wrapped_st_eth must be present for wrapped staked ETH pool") -
185            &amount_in;
186
187        Ok(GetAmountOutResult {
188            amount: amount_transferred.clone(),
189            gas: BigUint::from(DEFAULT_GAS),
190            new_state: Box::new(Self {
191                pool_type: self.pool_type.clone(),
192                total_shares: self.total_shares.clone(),
193                total_pooled_eth: self.total_pooled_eth.clone(),
194                total_wrapped_st_eth: Some(new_total_wrapped_st_eth),
195                id: self.id.clone(),
196                native_address: self.native_address.clone(),
197                stake_limits_state: self.stake_limits_state.clone(),
198                tokens: self.tokens.clone(),
199                token_to_track_total_pooled_eth: self
200                    .token_to_track_total_pooled_eth
201                    .clone(),
202            }),
203        })
204    }
205
206    fn zero2one(&self, sell_token: &Bytes, buy_token: &Bytes) -> Result<bool, SimulationError> {
207        let second_token = self
208            .tokens
209            .iter()
210            .find(|t| **t != self.token_to_track_total_pooled_eth)
211            .expect("No second token found");
212
213        if buy_token == second_token && *sell_token == self.token_to_track_total_pooled_eth {
214            Ok(true)
215        } else if *buy_token == self.token_to_track_total_pooled_eth && sell_token == second_token {
216            Ok(false)
217        } else {
218            Err(SimulationError::InvalidInput(
219                format!(
220                    "Invalid combination of tokens for type {:?}: {:?}, {:?}",
221                    self.pool_type, buy_token, sell_token
222                ),
223                None,
224            ))
225        }
226    }
227
228    fn st_eth_limits(
229        &self,
230        sell_token: Bytes,
231        buy_token: Bytes,
232    ) -> Result<(BigUint, BigUint), SimulationError> {
233        if self.zero2one(&sell_token, &buy_token)? {
234            let limit = self.stake_limits_state.get_limit();
235
236            let shares = safe_div_u256(
237                safe_mul_u256(biguint_to_u256(&limit), biguint_to_u256(&self.total_shares))?,
238                biguint_to_u256(&self.total_pooled_eth),
239            )?;
240
241            let amount_out = safe_div_u256(
242                safe_mul_u256(shares, biguint_to_u256(&self.total_pooled_eth))?,
243                biguint_to_u256(&self.total_shares),
244            )?;
245
246            Ok((limit.clone(), u256_to_biguint(amount_out)))
247        } else {
248            Ok((BigUint::zero(), BigUint::zero()))
249        }
250    }
251
252    fn wst_eth_limits(
253        &self,
254        sell_token: Bytes,
255        buy_token: Bytes,
256    ) -> Result<(BigUint, BigUint), SimulationError> {
257        if !self.zero2one(&sell_token, &buy_token)? {
258            //amount of wsteth, wstETH -> stETH
259
260            let total_wrapped_eth = self
261                .total_wrapped_st_eth
262                .clone()
263                .expect("total_wrapped_st_eth must be present for wrapped staked ETH pool");
264
265            let amount_out = u256_to_biguint(safe_div_u256(
266                safe_mul_u256(
267                    biguint_to_u256(&total_wrapped_eth),
268                    biguint_to_u256(&self.total_pooled_eth),
269                )?,
270                biguint_to_u256(&self.total_shares),
271            )?);
272            Ok((
273                self.total_wrapped_st_eth
274                    .clone()
275                    .expect("total_wrapped_st_eth must be present for wrapped staked ETH pool"),
276                amount_out,
277            ))
278        } else {
279            // total_shares - wstETH, stETH -> wstETH
280
281            let limit_for_wrapping = &self.total_shares -
282                self.total_wrapped_st_eth
283                    .as_ref()
284                    .expect("total_wrapped_st_eth must be present for wrapped staked ETH pool");
285
286            let amount_in = u256_to_biguint(safe_div_u256(
287                safe_mul_u256(
288                    biguint_to_u256(&limit_for_wrapping),
289                    biguint_to_u256(&self.total_shares),
290                )?,
291                biguint_to_u256(&self.total_pooled_eth),
292            )?);
293
294            Ok((amount_in, limit_for_wrapping))
295        }
296    }
297
298    fn st_eth_delta_transition(
299        &mut self,
300        delta: ProtocolStateDelta,
301    ) -> Result<(), TransitionError<String>> {
302        self.total_shares = BigUint::from_bytes_be(
303            delta
304                .updated_attributes
305                .get("total_shares")
306                .ok_or(TransitionError::MissingAttribute(
307                    "total_shares field is missing".to_owned(),
308                ))?,
309        );
310
311        let staking_status = delta
312            .updated_attributes
313            .get("staking_status")
314            .ok_or(TransitionError::MissingAttribute(
315                "Staking_status field is missing".to_owned(),
316            ))?;
317
318        let staking_status_parsed = if let Ok(status_as_str) = std::str::from_utf8(staking_status) {
319            match status_as_str {
320                "Limited" => StakingStatus::Limited,
321                "Paused" => StakingStatus::Paused,
322                "Unlimited" => StakingStatus::Unlimited,
323                _ => {
324                    return Err(TransitionError::DecodeError(
325                        "status_as_str parsed to invalid status".to_owned(),
326                    ))
327                }
328            }
329        } else {
330            return Err(TransitionError::DecodeError("status_as_str cannot be parsed".to_owned()))
331        };
332
333        let staking_limit = delta
334            .updated_attributes
335            .get("staking_limit")
336            .ok_or(TransitionError::MissingAttribute(
337                "Staking_limit field is missing".to_owned(),
338            ))?;
339
340        self.stake_limits_state = StakeLimitState {
341            staking_status: staking_status_parsed,
342            staking_limit: BigUint::from_bytes_be(staking_limit),
343        };
344        Ok(())
345    }
346
347    fn st_eth_balance_transition(&mut self, balances: &HashMap<Bytes, Bytes>) {
348        for (token, balance) in balances.iter() {
349            if token == &self.token_to_track_total_pooled_eth {
350                self.total_pooled_eth = BigUint::from_bytes_be(balance)
351            }
352        }
353    }
354
355    fn wst_eth_delta_transition(
356        &mut self,
357        delta: ProtocolStateDelta,
358    ) -> Result<(), TransitionError<String>> {
359        self.total_shares = BigUint::from_bytes_be(
360            delta
361                .updated_attributes
362                .get("total_shares")
363                .ok_or(TransitionError::MissingAttribute(
364                    "total_shares field is missing".to_owned(),
365                ))?,
366        );
367        self.total_wrapped_st_eth = Some(BigUint::from_bytes_be(
368            delta
369                .updated_attributes
370                .get("total_wstETH")
371                .ok_or(TransitionError::MissingAttribute(
372                    "total_wstETH field is missing".to_owned(),
373                ))?,
374        ));
375
376        Ok(())
377    }
378
379    fn wst_eth_balance_transition(&mut self, balances: &HashMap<Bytes, Bytes>) {
380        for (token, balance) in balances.iter() {
381            if token == &self.token_to_track_total_pooled_eth {
382                self.total_pooled_eth = BigUint::from_bytes_be(balance)
383            }
384        }
385    }
386}
387
388impl ProtocolSim for LidoState {
389    fn fee(&self) -> f64 {
390        // there is no fee when swapping
391        0.0
392    }
393
394    fn spot_price(&self, base: &Token, quote: &Token) -> Result<f64, SimulationError> {
395        match self.pool_type {
396            LidoPoolType::StEth => {
397                if self.zero2one(&base.address, &quote.address)? {
398                    let total_shares_f64 = u256_to_f64(biguint_to_u256(&self.total_shares))?;
399                    let total_pooled_eth_f64 =
400                        u256_to_f64(biguint_to_u256(&self.total_pooled_eth))?;
401
402                    Ok(total_pooled_eth_f64 / total_shares_f64 * total_shares_f64 /
403                        total_pooled_eth_f64)
404                } else {
405                    Err(SimulationError::InvalidInput(
406                        format!(
407                            "Spot_price: Invalid combination of tokens for type {:?}: {:?}, {:?}",
408                            self.pool_type, base, quote
409                        ),
410                        None,
411                    ))
412                }
413            }
414            LidoPoolType::WStEth => {
415                if self.zero2one(&base.address, &quote.address)? {
416                    let total_shares_f64 = u256_to_f64(biguint_to_u256(&self.total_shares))?;
417                    let total_pooled_eth_f64 =
418                        u256_to_f64(biguint_to_u256(&self.total_pooled_eth))?;
419
420                    Ok(total_shares_f64 / total_pooled_eth_f64)
421                } else {
422                    let total_shares_f64 = u256_to_f64(biguint_to_u256(&self.total_shares))?;
423                    let total_pooled_eth_f64 =
424                        u256_to_f64(biguint_to_u256(&self.total_pooled_eth))?;
425
426                    Ok(total_pooled_eth_f64 / total_shares_f64)
427                }
428            }
429        }
430    }
431
432    fn get_amount_out(
433        &self,
434        amount_in: BigUint,
435        token_in: &Token,
436        token_out: &Token,
437    ) -> Result<GetAmountOutResult, SimulationError> {
438        // check the pool type and the token in and token out
439        // if it's stETH type and the token out is ETH this is not allowed
440        // call the corresponding swap method
441        match self.pool_type {
442            LidoPoolType::StEth => {
443                if self.zero2one(&token_in.address, &token_out.address)? {
444                    Ok(self.steth_swap(amount_in)?)
445                } else {
446                    Err(SimulationError::InvalidInput(
447                        format!(
448                            "Invalid combination of tokens for type {:?}: {:?}, {:?}",
449                            self.pool_type, token_in, token_out
450                        ),
451                        None,
452                    ))
453                }
454            }
455            LidoPoolType::WStEth => {
456                if self.zero2one(&token_in.address, &token_out.address)? {
457                    self.wrap_steth(amount_in)
458                } else {
459                    self.unwrap_steth(amount_in)
460                }
461            }
462        }
463    }
464
465    fn get_limits(
466        &self,
467        sell_token: Bytes,
468        buy_token: Bytes,
469    ) -> Result<(BigUint, BigUint), SimulationError> {
470        // If it's the stETH type:
471        //   - and the buy token is ETH, the limits are 0
472        //   - sell token is ETH: use the StakeLimitState
473        // If it's wstETH: rely on the total supply
474        match self.pool_type {
475            LidoPoolType::StEth => self.st_eth_limits(sell_token, buy_token),
476            LidoPoolType::WStEth => self.wst_eth_limits(sell_token, buy_token),
477        }
478    }
479
480    fn delta_transition(
481        &mut self,
482        delta: ProtocolStateDelta,
483        _tokens: &HashMap<Bytes, Token>,
484        balances: &Balances,
485    ) -> Result<(), TransitionError<String>> {
486        for (component_id, balances) in balances.component_balances.iter() {
487            if Bytes::from(component_id.as_str()) == self.id {
488                match self.pool_type {
489                    LidoPoolType::StEth => self.st_eth_balance_transition(balances),
490                    LidoPoolType::WStEth => self.wst_eth_balance_transition(balances),
491                }
492            } else {
493                return Err(TransitionError::DecodeError(format!(
494                    "Invalid component id in balances: {:?}",
495                    component_id,
496                )))
497            }
498        }
499
500        if Bytes::from(delta.component_id.as_str()) == self.id {
501            match self.pool_type {
502                LidoPoolType::StEth => self.st_eth_delta_transition(delta),
503                LidoPoolType::WStEth => self.wst_eth_delta_transition(delta),
504            }
505        } else {
506            Err(TransitionError::DecodeError(format!(
507                "Invalid component id in delta: {:?}",
508                delta.component_id,
509            )))
510        }
511    }
512
513    fn clone_box(&self) -> Box<dyn ProtocolSim> {
514        Box::new(self.clone())
515    }
516
517    fn as_any(&self) -> &dyn Any {
518        self
519    }
520
521    fn as_any_mut(&mut self) -> &mut dyn Any {
522        self
523    }
524
525    fn eq(&self, other: &dyn ProtocolSim) -> bool {
526        if let Some(other_state) = other.as_any().downcast_ref::<Self>() {
527            self == other_state
528        } else {
529            false
530        }
531    }
532}
533
534#[cfg(test)]
535mod tests {
536    use std::{collections::HashSet, str::FromStr};
537
538    use num_bigint::BigUint;
539    use rstest::rstest;
540    use tycho_common::{
541        hex_bytes::Bytes,
542        models::{token::Token, Chain},
543    };
544
545    use super::*;
546
547    const ST_ETH_ADDRESS_PROXY: &str = "0xae7ab96520de3a18e5e111b5eaab095312d7fe84";
548    const WST_ETH_ADDRESS: &str = "0x7f39c581f595b53c5cb19bd0b3f8da6c935e2ca0";
549    const ETH_ADDRESS: &str = "0x0000000000000000000000000000000000000000";
550
551    fn from_hex_str_to_biguint(input: &str) -> BigUint {
552        let bytes = hex::decode(input).unwrap();
553        BigUint::from_bytes_be(&bytes)
554    }
555
556    fn lido_state_steth() -> LidoState {
557        let total_shares_start = from_hex_str_to_biguint(
558            "00000000000000000000000000000000000000000005dc41ec2e3ba19cf3ea6d",
559        );
560        let total_pooled_eth_start = from_hex_str_to_biguint("072409d75ebf50c5534125");
561        let staking_limit = from_hex_str_to_biguint("1fc3842bd1f071c00000");
562
563        LidoState {
564            pool_type: LidoPoolType::StEth,
565            total_shares: total_shares_start.clone(),
566            total_pooled_eth: total_pooled_eth_start.clone(),
567            total_wrapped_st_eth: None,
568            id: ST_ETH_ADDRESS_PROXY.into(),
569            native_address: ETH_ADDRESS.into(),
570            stake_limits_state: StakeLimitState {
571                staking_status: crate::evm::protocol::lido::state::StakingStatus::Limited,
572                staking_limit,
573            },
574            tokens: [
575                Bytes::from("0x0000000000000000000000000000000000000000"),
576                Bytes::from("0xae7ab96520de3a18e5e111b5eaab095312d7fe84"),
577            ],
578            token_to_track_total_pooled_eth: Bytes::from(ETH_ADDRESS),
579        }
580    }
581
582    fn lido_state_wsteth() -> LidoState {
583        let total_shares_start = from_hex_str_to_biguint(
584            "00000000000000000000000000000000000000000005dc41ec2e3ba19cf3ea6d",
585        );
586        let total_pooled_eth_start = from_hex_str_to_biguint("072409d75ebf50c5534125");
587        let total_wsteth_start = from_hex_str_to_biguint(
588            "00000000000000000000000000000000000000000002be110f2a220611513da6",
589        );
590
591        LidoState {
592            pool_type: LidoPoolType::WStEth,
593            total_shares: total_shares_start,
594            total_pooled_eth: total_pooled_eth_start.clone(),
595            total_wrapped_st_eth: Some(total_wsteth_start),
596            id: WST_ETH_ADDRESS.into(),
597            native_address: ETH_ADDRESS.into(),
598            stake_limits_state: StakeLimitState {
599                staking_status: crate::evm::protocol::lido::state::StakingStatus::Limited,
600                staking_limit: BigUint::zero(),
601            },
602            tokens: [
603                Bytes::from("0x7f39c581f595b53c5cb19bd0b3f8da6c935e2ca0"),
604                Bytes::from("0xae7ab96520de3a18e5e111b5eaab095312d7fe84"),
605            ],
606            token_to_track_total_pooled_eth: Bytes::from(ST_ETH_ADDRESS_PROXY),
607        }
608    }
609
610    fn bytes_st_eth() -> Bytes {
611        Bytes::from(ST_ETH_ADDRESS_PROXY)
612    }
613
614    fn bytes_wst_eth() -> Bytes {
615        Bytes::from(WST_ETH_ADDRESS)
616    }
617
618    fn bytes_eth() -> Bytes {
619        Bytes::from(ETH_ADDRESS)
620    }
621
622    fn token_st_eth() -> Token {
623        Token::new(
624            &Bytes::from_str(ST_ETH_ADDRESS_PROXY).unwrap(),
625            "stETH",
626            18,
627            0,
628            &[Some(44000)],
629            Chain::Ethereum,
630            10,
631        )
632    }
633
634    fn token_wst_eth() -> Token {
635        Token::new(
636            &Bytes::from_str(WST_ETH_ADDRESS).unwrap(),
637            "wstETH",
638            18,
639            0,
640            &[Some(44000)],
641            Chain::Ethereum,
642            10,
643        )
644    }
645
646    fn token_eth() -> Token {
647        Token::new(
648            &Bytes::from_str(ETH_ADDRESS).unwrap(),
649            "ETH",
650            18,
651            0,
652            &[Some(44000)],
653            Chain::Ethereum,
654            100,
655        )
656    }
657
658    #[test]
659    fn test_lido_get_amount_out() {
660        // total pooled eth: 0x072409d75ebf50c5534125, 8632667470434094430765349
661        // total shares: 0x00000000000000000000000000000000000000000005dc41ec2e3ba19cf3ea6d
662        // tx 0x1953b525c8640c2709e984ebc28bb1f2180dd72759bb2aac7413e94b602b0d53
663        // total_shares_after: 0x00000000000000000000000000000000000000000005dc41ec487a31b7865d5e
664        // total pooled eth after: 0x072409d77eb9c55db21616
665        let token_eth = token_eth();
666        let token_st_eth = token_st_eth();
667        let state = lido_state_steth();
668
669        let amount_in = BigUint::from_str("9001102957532401").unwrap();
670        let res = state
671            .get_amount_out(amount_in.clone(), &token_eth, &token_st_eth)
672            .unwrap();
673
674        let exp = BigUint::from_str("9001102957532398").unwrap(); // diff in total pooled eth; rounding error because of multiple casts
675        assert_eq!(res.amount, exp);
676
677        let total_shares_after = from_hex_str_to_biguint(
678            "00000000000000000000000000000000000000000005dc41ec487a31b7865d5e",
679        );
680
681        let total_pooled_eth_after = from_hex_str_to_biguint("072409d77eb9c55db21616");
682
683        let new_state = res
684            .new_state
685            .as_any()
686            .downcast_ref::<LidoState>()
687            .unwrap();
688        assert_eq!(new_state.total_shares, total_shares_after);
689        assert_eq!(new_state.total_pooled_eth, total_pooled_eth_after);
690    }
691
692    #[test]
693    fn test_lido_wrapping_get_amount_out() {
694        // total pooled eth: 072409d88cbb5e48a01616
695        // total shares: 00000000000000000000000000000000000000000005dc41ed2611c5fd46f034
696        // tx 0xce9418dd0e74cdf738362bee9428da73c047049c34f28ff8a18b19f047c27c53
697        // ws eth before 00000000000000000000000000000000000000000002be10e0f61dc56f6f85dc
698        // ws eth after 00000000000000000000000000000000000000000002be11ccda98eef241c759
699
700        let token_st_eth = token_st_eth();
701        let token_wst_eth = token_wst_eth();
702        let mut state = lido_state_wsteth();
703
704        let total_wsteth_start = from_hex_str_to_biguint(
705            "00000000000000000000000000000000000000000002be10e0f61dc56f6f85dc",
706        );
707        state.total_wrapped_st_eth = Some(total_wsteth_start);
708
709        let total_pooled_eth = from_hex_str_to_biguint("072409d88cbb5e48a01616");
710        state.total_pooled_eth = total_pooled_eth;
711
712        let total_shares = from_hex_str_to_biguint(
713            "00000000000000000000000000000000000000000005dc41ed2611c5fd46f034",
714        );
715        state.total_shares = total_shares;
716
717        let amount_in = BigUint::from_str("20711588703656141053").unwrap();
718        let res = state
719            .get_amount_out(amount_in.clone(), &token_st_eth, &token_wst_eth)
720            .unwrap();
721        let exp = BigUint::from_str("16997846311821787517").unwrap();
722        assert_eq!(res.amount, exp);
723
724        let total_wsteth_after = from_hex_str_to_biguint(
725            "00000000000000000000000000000000000000000002be11ccda98eef241c759",
726        );
727        let new_state = res
728            .new_state
729            .as_any()
730            .downcast_ref::<LidoState>()
731            .unwrap();
732        assert_eq!(new_state.total_wrapped_st_eth, Some(total_wsteth_after));
733
734        assert!(state
735            .get_amount_out(BigUint::zero(), &token_st_eth, &token_wst_eth)
736            .is_err());
737    }
738
739    #[test]
740
741    fn test_lido_unwrapping_get_amount_out() {
742        // total pooled eth: 0x072409d75ebf50c5534125, 8632667470434094430765349
743        // total shares: 0x00000000000000000000000000000000000000000005dc41ec2e3ba19cf3ea6d
744        // tx 0xa49316d76b7cf2ba9f81c7b84868faaa6306eef5a15f194f55b3675bce89367a
745        // ws eth after 0x00000000000000000000000000000000000000000002be10e0f61dc56f6f85dc
746        // ws eth before 0x00000000000000000000000000000000000000000002be110f2a220611513da6
747
748        let token_st_eth = token_st_eth();
749        let token_wst_eth = token_wst_eth();
750        let mut state = lido_state_wsteth();
751
752        let total_wsteth_start = from_hex_str_to_biguint(
753            "00000000000000000000000000000000000000000002be110f2a220611513da6",
754        );
755        state.total_wrapped_st_eth = Some(total_wsteth_start);
756
757        let amount_in = BigUint::from_str("3329290700173981642").unwrap();
758        let res = state
759            .get_amount_out(amount_in.clone(), &token_wst_eth, &token_st_eth)
760            .unwrap();
761        let exp = BigUint::from_str("4056684499432944067").unwrap(); // diff in total pooled eth; rounding error because of multiple casts
762        assert_eq!(res.amount, exp);
763
764        let total_wsteth_after = from_hex_str_to_biguint(
765            "00000000000000000000000000000000000000000002be10e0f61dc56f6f85dc",
766        );
767
768        let new_state = res
769            .new_state
770            .as_any()
771            .downcast_ref::<LidoState>()
772            .unwrap();
773
774        assert_eq!(new_state.total_wrapped_st_eth, Some(total_wsteth_after));
775    }
776
777    #[test]
778    fn test_lido_spot_price() {
779        let token_st_eth = token_st_eth();
780        let token_wst_eth = token_wst_eth();
781        let token_eth = token_eth();
782
783        let st_state = lido_state_steth();
784        let wst_state = lido_state_wsteth();
785
786        let res = st_state
787            .spot_price(&token_eth, &token_st_eth)
788            .unwrap();
789        let exp = 1.0000000000000002;
790        assert_eq!(res, exp);
791
792        let res = st_state.spot_price(&token_st_eth, &token_wst_eth);
793        assert!(res.is_err());
794
795        let res = wst_state
796            .spot_price(&token_st_eth, &token_wst_eth)
797            .unwrap();
798        assert_eq!(res, 0.8206925386086495);
799
800        let res = wst_state
801            .spot_price(&token_wst_eth, &token_st_eth)
802            .unwrap();
803        assert_eq!(res, 1.2184831139019945);
804
805        let res = wst_state.spot_price(&token_eth, &token_st_eth);
806        assert!(res.is_err());
807    }
808
809    #[test]
810    fn test_lido_get_limits() {
811        let token_st_eth = bytes_st_eth();
812        let token_wst_eth = bytes_wst_eth();
813        let token_eth = bytes_eth();
814
815        let st_state = lido_state_steth();
816        let wst_state = lido_state_wsteth();
817
818        let res = st_state
819            .get_limits(token_eth.clone(), token_st_eth.clone())
820            .unwrap();
821        let exp = (
822            st_state
823                .stake_limits_state
824                .staking_limit
825                .clone(),
826            BigUint::from_str("149999999999999999999999").unwrap(),
827        );
828        assert_eq!(res, exp);
829
830        let res = st_state
831            .get_limits(token_st_eth.clone(), token_eth.clone())
832            .unwrap();
833        let exp = (BigUint::zero(), BigUint::zero());
834        assert_eq!(res, exp);
835
836        let res = st_state.get_limits(token_wst_eth.clone(), token_eth.clone());
837        assert!(res.is_err());
838
839        let res = wst_state
840            .get_limits(token_st_eth.clone(), token_wst_eth.clone())
841            .unwrap();
842        let allowed_to_wrap = wst_state.total_shares.clone() -
843            wst_state
844                .total_wrapped_st_eth
845                .clone()
846                .unwrap();
847        let exp = (BigUint::from_str("3093477275082723426591391").unwrap(), allowed_to_wrap);
848
849        assert_eq!(res, exp);
850
851        let res = wst_state
852            .get_limits(token_wst_eth.clone(), token_st_eth.clone())
853            .unwrap();
854        let total_wrapped = wst_state
855            .total_wrapped_st_eth
856            .clone()
857            .unwrap();
858        let exp = (total_wrapped, BigUint::from_str("4039778360807033131920717").unwrap());
859
860        assert_eq!(res, exp);
861
862        let res = wst_state.get_limits(token_wst_eth.clone(), token_eth.clone());
863        assert!(res.is_err());
864    }
865
866    #[test]
867    fn test_lido_st_delta_transition() {
868        let mut st_state = lido_state_steth();
869
870        let total_shares_after =
871            "0x00000000000000000000000000000000000000000005d9d75ae42b4ba9c04d1a";
872        let staking_status_after = "0x4c696d69746564";
873        let staking_limit_after = "0x1fc3842bd1f071c00000";
874
875        let mut updated_attributes: HashMap<String, Bytes> = HashMap::new();
876        updated_attributes.insert("total_shares".to_owned(), Bytes::from(total_shares_after));
877        updated_attributes.insert("staking_status".to_owned(), Bytes::from(staking_status_after));
878        updated_attributes.insert("staking_limit".to_owned(), Bytes::from(staking_limit_after));
879
880        let staking_state_delta = ProtocolStateDelta {
881            component_id: ST_ETH_ADDRESS_PROXY.to_owned(),
882            updated_attributes,
883            deleted_attributes: HashSet::new(),
884        };
885
886        let mut component_balances = HashMap::new();
887        let mut component_balances_inner = HashMap::new();
888        component_balances_inner.insert(
889            Bytes::from_str(ETH_ADDRESS).unwrap(),
890            Bytes::from_str("0x072409d88cbb5e48a01616").unwrap(),
891        );
892        component_balances.insert(ST_ETH_ADDRESS_PROXY.to_owned(), component_balances_inner);
893
894        let balances = Balances { component_balances, account_balances: HashMap::new() };
895
896        st_state
897            .delta_transition(staking_state_delta.clone(), &HashMap::new(), &balances)
898            .unwrap();
899
900        let exp = LidoState {
901            pool_type: LidoPoolType::StEth,
902            total_shares: BigUint::from_bytes_be(
903                &hex::decode("00000000000000000000000000000000000000000005d9d75ae42b4ba9c04d1a")
904                    .unwrap(),
905            ),
906            total_pooled_eth: from_hex_str_to_biguint("072409d88cbb5e48a01616"),
907            total_wrapped_st_eth: None,
908            id: ST_ETH_ADDRESS_PROXY.into(),
909            native_address: ETH_ADDRESS.into(),
910            stake_limits_state: StakeLimitState {
911                staking_status: crate::evm::protocol::lido::state::StakingStatus::Limited,
912                staking_limit: BigUint::from_bytes_be(
913                    &hex::decode("1fc3842bd1f071c00000").unwrap(),
914                ),
915            },
916            tokens: [
917                Bytes::from("0x0000000000000000000000000000000000000000"),
918                Bytes::from("0xae7ab96520de3a18e5e111b5eaab095312d7fe84"),
919            ],
920            token_to_track_total_pooled_eth: Bytes::from(ETH_ADDRESS),
921        };
922        assert_eq!(st_state, exp);
923    }
924
925    #[rstest]
926    #[case::missing_total_shares("total_shares")]
927    #[case::missing_staking_status("staking_status")]
928    #[case::missing_staking_limit("staking_limit")]
929    fn test_lido_st_delta_transition_missing_arg(#[case] missing_attribute: &str) {
930        let mut st_state = lido_state_steth();
931
932        let total_shares_after =
933            "0x00000000000000000000000000000000000000000005d9d75ae42b4ba9c04d1a";
934        let staking_status_after = "0x4c696d69746564";
935        let staking_limit_after = "0x1fc3842bd1f071c00000";
936
937        let mut updated_attributes: HashMap<String, Bytes> = HashMap::new();
938        updated_attributes.insert("total_shares".to_owned(), Bytes::from(total_shares_after));
939        updated_attributes.insert("staking_status".to_owned(), Bytes::from(staking_status_after));
940        updated_attributes.insert("staking_limit".to_owned(), Bytes::from(staking_limit_after));
941
942        let mut staking_state_delta = ProtocolStateDelta {
943            component_id: ST_ETH_ADDRESS_PROXY.to_owned(),
944            updated_attributes,
945            deleted_attributes: HashSet::new(),
946        };
947
948        let balances =
949            Balances { component_balances: HashMap::new(), account_balances: HashMap::new() };
950
951        staking_state_delta
952            .updated_attributes
953            .remove(missing_attribute);
954
955        assert!(st_state
956            .delta_transition(staking_state_delta.clone(), &HashMap::new(), &balances)
957            .is_err());
958    }
959
960    #[test]
961    fn test_lido_wst_delta_transition() {
962        let mut wst_state = lido_state_wsteth();
963
964        let total_shares_after =
965            "0x00000000000000000000000000000000000000000005d9d75ae42b4ba9c04d1a";
966        let total_ws_eth_after =
967            "0x00000000000000000000000000000000000000000002ba6f7b9af3c7a7b749e2";
968
969        let mut updated_attributes: HashMap<String, Bytes> = HashMap::new();
970        updated_attributes.insert("total_shares".to_owned(), Bytes::from(total_shares_after));
971        updated_attributes.insert("total_wstETH".to_owned(), Bytes::from(total_ws_eth_after));
972
973        let staking_state_delta = ProtocolStateDelta {
974            component_id: WST_ETH_ADDRESS.to_owned(),
975            updated_attributes,
976            deleted_attributes: HashSet::new(),
977        };
978
979        let mut component_balances = HashMap::new();
980        let mut component_balances_inner = HashMap::new();
981        component_balances_inner.insert(
982            Bytes::from_str(ST_ETH_ADDRESS_PROXY).unwrap(),
983            Bytes::from_str("0x072409d88cbb5e48a01616").unwrap(),
984        );
985        component_balances.insert(WST_ETH_ADDRESS.to_owned(), component_balances_inner);
986
987        let balances = Balances { component_balances, account_balances: HashMap::new() };
988
989        wst_state
990            .delta_transition(staking_state_delta, &HashMap::new(), &balances)
991            .unwrap();
992
993        let exp = LidoState {
994            pool_type: LidoPoolType::WStEth,
995            total_shares: BigUint::from_bytes_be(
996                &hex::decode("00000000000000000000000000000000000000000005d9d75ae42b4ba9c04d1a")
997                    .unwrap(),
998            ),
999            total_pooled_eth: from_hex_str_to_biguint("072409d88cbb5e48a01616"),
1000            total_wrapped_st_eth: Some(BigUint::from_bytes_be(
1001                &hex::decode("00000000000000000000000000000000000000000002ba6f7b9af3c7a7b749e2")
1002                    .unwrap(),
1003            )),
1004            id: WST_ETH_ADDRESS.into(),
1005            native_address: ETH_ADDRESS.into(),
1006            stake_limits_state: wst_state.stake_limits_state.clone(),
1007            tokens: [
1008                Bytes::from("0x7f39c581f595b53c5cb19bd0b3f8da6c935e2ca0"),
1009                Bytes::from("0xae7ab96520de3a18e5e111b5eaab095312d7fe84"),
1010            ],
1011            token_to_track_total_pooled_eth: Bytes::from(ST_ETH_ADDRESS_PROXY),
1012        };
1013
1014        assert_eq!(wst_state, exp);
1015    }
1016
1017    #[rstest]
1018    #[case::missing_total_shares("total_shares")]
1019    #[case::missing_total_ws_eth("total_wstETH")]
1020    fn test_lido_wst_delta_transition_missing_arg(#[case] missing_attribute: &str) {
1021        let mut wst_state = lido_state_wsteth();
1022
1023        let total_shares_after =
1024            "0x00000000000000000000000000000000000000000005d9d75ae42b4ba9c04d1a";
1025        let total_ws_eth_after =
1026            "0x00000000000000000000000000000000000000000002ba6f7b9af3c7a7b749e2";
1027
1028        let mut updated_attributes: HashMap<String, Bytes> = HashMap::new();
1029        updated_attributes.insert("total_shares".to_owned(), Bytes::from(total_shares_after));
1030        updated_attributes.insert("total_wstETH".to_owned(), Bytes::from(total_ws_eth_after));
1031
1032        let mut staking_state_delta = ProtocolStateDelta {
1033            component_id: WST_ETH_ADDRESS.to_owned(),
1034            updated_attributes,
1035            deleted_attributes: HashSet::new(),
1036        };
1037
1038        let balances =
1039            Balances { component_balances: HashMap::new(), account_balances: HashMap::new() };
1040
1041        staking_state_delta
1042            .updated_attributes
1043            .remove(missing_attribute);
1044
1045        assert!(wst_state
1046            .delta_transition(staking_state_delta.clone(), &HashMap::new(), &balances)
1047            .is_err());
1048    }
1049}