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
81impl 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}