Skip to main content

tycho_simulation/evm/
tycho_models.rs

1use std::collections::HashMap;
2
3use alloy::primitives::{Address, U256};
4use serde::{Deserialize, Serialize};
5use tycho_common::simulation::errors::SimulationError;
6pub use tycho_common::{dto::ChangeType, models::Chain};
7
8use crate::{
9    evm::protocol::{u256_num, utils::bytes_to_address},
10    serde_helpers::{hex_bytes, hex_bytes_option},
11};
12
13#[derive(PartialEq, Serialize, Deserialize, Clone, Debug)]
14pub struct AccountUpdate {
15    pub address: Address,
16    pub chain: Chain,
17    pub slots: HashMap<U256, U256>,
18    pub balance: Option<U256>,
19    #[serde(with = "hex_bytes_option")]
20    pub code: Option<Vec<u8>>,
21    pub change: ChangeType,
22}
23
24impl AccountUpdate {
25    pub fn new(
26        address: Address,
27        chain: Chain,
28        slots: HashMap<U256, U256>,
29        balance: Option<U256>,
30        code: Option<Vec<u8>>,
31        change: ChangeType,
32    ) -> Self {
33        Self { address, chain, slots, balance, code, change }
34    }
35}
36
37impl TryFrom<tycho_common::dto::AccountUpdate> for AccountUpdate {
38    type Error = SimulationError;
39
40    fn try_from(value: tycho_common::dto::AccountUpdate) -> Result<Self, Self::Error> {
41        Ok(Self {
42            chain: value.chain.into(),
43            address: bytes_to_address(&value.address)?,
44            slots: u256_num::map_slots_to_u256(value.slots),
45            balance: value
46                .balance
47                .map(|balance| u256_num::bytes_to_u256(balance.into())),
48            code: value.code.map(|code| code.to_vec()),
49            change: value.change,
50        })
51    }
52}
53
54#[derive(PartialEq, Clone, Serialize, Deserialize, Default)]
55#[serde(rename = "Account")]
56pub struct ResponseAccount {
57    pub chain: Chain,
58    pub address: Address,
59    pub title: String,
60    pub slots: HashMap<U256, U256>,
61    pub native_balance: U256,
62    pub token_balances: HashMap<Address, U256>,
63    #[serde(with = "hex_bytes")]
64    pub code: Vec<u8>,
65}
66
67impl ResponseAccount {
68    pub fn new(
69        chain: Chain,
70        address: Address,
71        title: String,
72        slots: HashMap<U256, U256>,
73        native_balance: U256,
74        token_balances: HashMap<Address, U256>,
75        code: Vec<u8>,
76    ) -> Self {
77        Self { chain, address, title, slots, native_balance, token_balances, code }
78    }
79}
80
81/// Implement Debug for ResponseAccount manually to avoid printing the code field.
82impl std::fmt::Debug for ResponseAccount {
83    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
84        f.debug_struct("ResponseAccount")
85            .field("chain", &self.chain)
86            .field("address", &self.address)
87            .field("title", &self.title)
88            .field("slots", &self.slots)
89            .field("native_balance", &self.native_balance)
90            .field("token_balances", &self.token_balances)
91            .field("code", &format!("[{} bytes]", self.code.len()))
92            .finish()
93    }
94}
95
96impl TryFrom<tycho_common::dto::ResponseAccount> for ResponseAccount {
97    type Error = SimulationError;
98
99    #[allow(deprecated)]
100    fn try_from(value: tycho_common::dto::ResponseAccount) -> Result<Self, Self::Error> {
101        let token_balances = value
102            .token_balances
103            .into_iter()
104            .map(|(address, balance)| {
105                Ok((bytes_to_address(&address)?, u256_num::bytes_to_u256(balance.into())))
106            })
107            .collect::<Result<HashMap<_, _>, SimulationError>>()?;
108
109        Ok(Self {
110            chain: value.chain.into(),
111            address: bytes_to_address(&value.address)?,
112            title: value.title.clone(),
113            slots: u256_num::map_slots_to_u256(value.slots),
114            native_balance: u256_num::bytes_to_u256(value.native_balance.into()),
115            token_balances,
116            code: value.code.to_vec(),
117        })
118    }
119}
120
121impl From<tycho_common::models::contract::Account> for ResponseAccount {
122    fn from(value: tycho_common::models::contract::Account) -> Self {
123        Self {
124            chain: value.chain,
125            address: Address::from_slice(&value.address[..20]),
126            title: value.title,
127            slots: u256_num::map_slots_to_u256(value.slots),
128            native_balance: u256_num::bytes_to_u256(value.native_balance.into()),
129            token_balances: value
130                .token_balances
131                .into_iter()
132                .map(|(addr, ab)| {
133                    (Address::from_slice(&addr[..20]), u256_num::bytes_to_u256(ab.balance.into()))
134                })
135                .collect(),
136            code: value.code.to_vec(),
137        }
138    }
139}
140
141impl From<tycho_common::models::contract::AccountDelta> for AccountUpdate {
142    fn from(value: tycho_common::models::contract::AccountDelta) -> Self {
143        let code = value.code().clone().map(|c| c.to_vec());
144        let change = value.change_type().into();
145        Self {
146            chain: value.chain,
147            address: Address::from_slice(&value.address[..20]),
148            slots: value
149                .slots
150                .into_iter()
151                .map(|(k, v)| {
152                    (
153                        u256_num::bytes_to_u256(k.into()),
154                        u256_num::bytes_to_u256(v.unwrap_or_default().into()),
155                    )
156                })
157                .collect(),
158            balance: value
159                .balance
160                .map(|b| u256_num::bytes_to_u256(b.into())),
161            code,
162            change,
163        }
164    }
165}
166
167#[cfg(test)]
168mod tests {
169    use tycho_common::Bytes;
170
171    use super::*;
172
173    fn make_dto_response_account(address: Bytes) -> tycho_common::dto::ResponseAccount {
174        #[allow(deprecated)]
175        tycho_common::dto::ResponseAccount::new(
176            tycho_common::dto::Chain::Ethereum,
177            address,
178            "test".to_string(),
179            HashMap::new(),
180            Bytes::zero(32),
181            HashMap::new(),
182            Bytes::from(vec![0xDE, 0xAD]),
183            Bytes::from("0x00"),
184            Bytes::from("0x00"),
185            Bytes::from("0x00"),
186            None,
187        )
188    }
189
190    #[test]
191    fn test_response_account_conversion_succeeds() {
192        let dto = make_dto_response_account(Bytes::zero(20));
193
194        let result = ResponseAccount::try_from(dto).unwrap();
195
196        assert_eq!(result.address, Address::ZERO);
197        assert_eq!(result.code, vec![0xDE, 0xAD]);
198    }
199
200    #[test]
201    fn test_response_account_conversion_short_address_fails() {
202        let dto = make_dto_response_account(Bytes::from(vec![0x01]));
203
204        let result = ResponseAccount::try_from(dto);
205
206        assert!(result.is_err());
207    }
208}