1use std::collections::HashMap;
22
23use alloy_primitives::{Address, B256, I256, Log, U256};
24use revm::Inspector;
25use revm::interpreter::InterpreterTypes;
26
27const TRANSFER_EVENT_SIGNATURE: B256 = B256::new([
31 0xdd, 0xf2, 0x52, 0xad, 0x1b, 0xe2, 0xc8, 0x9b, 0x69, 0xc2, 0xb0, 0x68, 0xfc, 0x37, 0x8d, 0xaa,
32 0x95, 0x2b, 0xa7, 0xf1, 0x63, 0xc4, 0xa1, 0x16, 0x28, 0xf5, 0x5a, 0x4d, 0xf5, 0x23, 0xb3, 0xef,
33]);
34
35#[derive(Clone, Debug, PartialEq, Eq)]
40pub struct TokenTransfer {
41 pub token: Address,
43 pub from: Address,
45 pub to: Address,
47 pub value: U256,
49}
50
51#[derive(Clone, Debug, Default)]
60pub struct TransferInspector {
61 pub transfers: Vec<TokenTransfer>,
63 pub logs: Vec<Log>,
65}
66
67impl TransferInspector {
68 pub fn new() -> Self {
78 Self {
79 transfers: Vec::new(),
80 logs: Vec::new(),
81 }
82 }
83
84 pub fn balance_deltas(&self, owner: Address) -> HashMap<Address, I256> {
89 let mut deltas: HashMap<Address, I256> = HashMap::new();
90
91 for transfer in &self.transfers {
92 if transfer.from == owner {
93 let entry = deltas.entry(transfer.token).or_insert(I256::ZERO);
95 *entry -= I256::from_raw(transfer.value);
96 }
97 if transfer.to == owner {
98 let entry = deltas.entry(transfer.token).or_insert(I256::ZERO);
100 *entry += I256::from_raw(transfer.value);
101 }
102 }
103
104 deltas
105 }
106
107 pub fn balance_deltas_for_tokens(
131 &self,
132 owner: Address,
133 tokens: impl IntoIterator<Item = Address>,
134 ) -> HashMap<Address, I256> {
135 let token_set: std::collections::HashSet<Address> = tokens.into_iter().collect();
136 let all_deltas = self.balance_deltas(owner);
137
138 all_deltas
139 .into_iter()
140 .filter(|(token, _)| token_set.contains(token))
141 .collect()
142 }
143
144 pub fn clear(&mut self) {
147 self.transfers.clear();
148 self.logs.clear();
149 }
150
151 pub fn parse_transfer(log: &Log) -> Option<TokenTransfer> {
153 let topics = log.topics();
160 if topics.len() < 3 {
161 return None;
162 }
163
164 if topics[0] != TRANSFER_EVENT_SIGNATURE {
166 return None;
167 }
168
169 let from = Address::from_word(topics[1]);
171 let to = Address::from_word(topics[2]);
172
173 let data = log.data.data.as_ref();
175 if data.len() < 32 {
176 return None;
177 }
178 let value = U256::from_be_slice(&data[..32]);
179
180 Some(TokenTransfer {
181 token: log.address,
182 from,
183 to,
184 value,
185 })
186 }
187}
188
189impl<CTX, INTR> Inspector<CTX, INTR> for TransferInspector
197where
198 INTR: InterpreterTypes,
199{
200 fn log(&mut self, _context: &mut CTX, log: Log) {
202 if let Some(transfer) = Self::parse_transfer(&log) {
204 self.transfers.push(transfer);
205 }
206
207 self.logs.push(log);
209 }
210}
211
212pub fn parse_receipt_deltas(receipt_logs: &[Log], owner: Address) -> HashMap<Address, I256> {
218 let mut deltas: HashMap<Address, I256> = HashMap::new();
219
220 for log in receipt_logs {
221 if let Some(transfer) = TransferInspector::parse_transfer(log) {
222 if transfer.from == owner {
223 let entry = deltas.entry(transfer.token).or_insert(I256::ZERO);
224 *entry -= I256::from_raw(transfer.value);
225 }
226 if transfer.to == owner {
227 let entry = deltas.entry(transfer.token).or_insert(I256::ZERO);
228 *entry += I256::from_raw(transfer.value);
229 }
230 }
231 }
232
233 deltas
234}
235
236#[cfg(test)]
237mod tests {
238 use super::*;
239 use alloy_primitives::{address, bytes};
240
241 fn make_transfer_log(token: Address, from: Address, to: Address, value: U256) -> Log {
242 Log::new(
243 token,
244 vec![TRANSFER_EVENT_SIGNATURE, from.into_word(), to.into_word()],
245 value.to_be_bytes_vec().into(),
246 )
247 .unwrap()
248 }
249
250 #[test]
251 fn test_parse_single_transfer() {
252 let token = address!("A0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48"); let from = address!("1111111111111111111111111111111111111111");
254 let to = address!("2222222222222222222222222222222222222222");
255 let value = U256::from(1000000u64); let log = make_transfer_log(token, from, to, value);
258 let transfer = TransferInspector::parse_transfer(&log).expect("should parse");
259
260 assert_eq!(transfer.token, token);
261 assert_eq!(transfer.from, from);
262 assert_eq!(transfer.to, to);
263 assert_eq!(transfer.value, value);
264 }
265
266 #[test]
267 fn test_balance_deltas_outgoing() {
268 let mut inspector = TransferInspector::new();
269 let token = address!("A0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48");
270 let owner = address!("1111111111111111111111111111111111111111");
271 let recipient = address!("2222222222222222222222222222222222222222");
272 let value = U256::from(1000000u64);
273
274 inspector.transfers.push(TokenTransfer {
275 token,
276 from: owner,
277 to: recipient,
278 value,
279 });
280
281 let deltas = inspector.balance_deltas(owner);
282 assert_eq!(deltas.get(&token), Some(&(-I256::from_raw(value))));
283 }
284
285 #[test]
286 fn test_balance_deltas_incoming() {
287 let mut inspector = TransferInspector::new();
288 let token = address!("A0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48");
289 let sender = address!("1111111111111111111111111111111111111111");
290 let owner = address!("2222222222222222222222222222222222222222");
291 let value = U256::from(1000000u64);
292
293 inspector.transfers.push(TokenTransfer {
294 token,
295 from: sender,
296 to: owner,
297 value,
298 });
299
300 let deltas = inspector.balance_deltas(owner);
301 assert_eq!(deltas.get(&token), Some(&I256::from_raw(value)));
302 }
303
304 #[test]
305 fn test_balance_deltas_multiple_transfers() {
306 let mut inspector = TransferInspector::new();
307 let token_a = address!("A0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48");
308 let token_b = address!("C02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2");
309 let owner = address!("1111111111111111111111111111111111111111");
310 let other = address!("2222222222222222222222222222222222222222");
311
312 inspector.transfers.push(TokenTransfer {
314 token: token_a,
315 from: owner,
316 to: other,
317 value: U256::from(100u64),
318 });
319
320 inspector.transfers.push(TokenTransfer {
322 token: token_b,
323 from: other,
324 to: owner,
325 value: U256::from(50u64),
326 });
327
328 inspector.transfers.push(TokenTransfer {
330 token: token_a,
331 from: other,
332 to: owner,
333 value: U256::from(25u64),
334 });
335
336 let deltas = inspector.balance_deltas(owner);
337
338 assert_eq!(deltas.get(&token_a), Some(&I256::try_from(-75i64).unwrap()));
340 assert_eq!(deltas.get(&token_b), Some(&I256::try_from(50i64).unwrap()));
342 }
343
344 #[test]
345 fn test_balance_deltas_for_tokens_filter() {
346 let mut inspector = TransferInspector::new();
347 let token_a = address!("A0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48");
348 let token_b = address!("C02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2");
349 let token_c = address!("6B175474E89094C44Da98b954EeecdE504D571D1");
350 let owner = address!("1111111111111111111111111111111111111111");
351 let other = address!("2222222222222222222222222222222222222222");
352
353 inspector.transfers.push(TokenTransfer {
354 token: token_a,
355 from: owner,
356 to: other,
357 value: U256::from(100u64),
358 });
359
360 inspector.transfers.push(TokenTransfer {
361 token: token_b,
362 from: other,
363 to: owner,
364 value: U256::from(50u64),
365 });
366
367 inspector.transfers.push(TokenTransfer {
368 token: token_c,
369 from: other,
370 to: owner,
371 value: U256::from(200u64),
372 });
373
374 let deltas = inspector.balance_deltas_for_tokens(owner, vec![token_a, token_b]);
376
377 assert_eq!(deltas.len(), 2);
378 assert!(deltas.contains_key(&token_a));
379 assert!(deltas.contains_key(&token_b));
380 assert!(!deltas.contains_key(&token_c));
381 }
382
383 #[test]
384 fn test_non_transfer_log_ignored() {
385 let log = Log::new(
387 address!("A0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48"),
388 vec![
389 B256::ZERO, address!("1111111111111111111111111111111111111111").into_word(),
391 address!("2222222222222222222222222222222222222222").into_word(),
392 ],
393 U256::from(100u64).to_be_bytes_vec().into(),
394 )
395 .unwrap();
396
397 assert!(TransferInspector::parse_transfer(&log).is_none());
398 }
399
400 #[test]
401 fn test_inspector_clear() {
402 let mut inspector = TransferInspector::new();
403 inspector.transfers.push(TokenTransfer {
404 token: Address::ZERO,
405 from: Address::ZERO,
406 to: Address::ZERO,
407 value: U256::ZERO,
408 });
409 inspector
410 .logs
411 .push(Log::new_unchecked(Address::ZERO, vec![], bytes!("")));
412
413 assert!(!inspector.transfers.is_empty());
414 assert!(!inspector.logs.is_empty());
415
416 inspector.clear();
417
418 assert!(inspector.transfers.is_empty());
419 assert!(inspector.logs.is_empty());
420 }
421
422 #[test]
423 fn test_parse_receipt_deltas() {
424 let token_a = address!("A0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48");
425 let token_b = address!("C02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2");
426 let owner = address!("1111111111111111111111111111111111111111");
427 let other = address!("2222222222222222222222222222222222222222");
428
429 let logs = vec![
430 make_transfer_log(token_a, owner, other, U256::from(100u64)),
432 make_transfer_log(token_b, other, owner, U256::from(50u64)),
434 Log::new(
436 token_a,
437 vec![B256::ZERO, owner.into_word(), other.into_word()],
438 U256::from(999u64).to_be_bytes_vec().into(),
439 )
440 .unwrap(),
441 make_transfer_log(token_a, other, owner, U256::from(25u64)),
443 ];
444
445 let deltas = super::parse_receipt_deltas(&logs, owner);
446
447 assert_eq!(deltas.get(&token_a), Some(&I256::try_from(-75i64).unwrap()));
449 assert_eq!(deltas.get(&token_b), Some(&I256::try_from(50i64).unwrap()));
451 }
452
453 #[test]
454 fn test_parse_receipt_deltas_empty() {
455 let owner = address!("1111111111111111111111111111111111111111");
456 let deltas = super::parse_receipt_deltas(&[], owner);
457 assert!(deltas.is_empty());
458 }
459}