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#[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 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 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 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 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 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 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, "e.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, "e.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 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 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 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(); 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 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 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(); 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}