1use std::collections::{hash_map::Entry, BTreeMap, HashMap, HashSet};
2
3use chrono::NaiveDateTime;
4use deepsize::DeepSizeOf;
5use serde::{ser::SerializeStruct, Deserialize, Serialize, Serializer};
6use tracing::warn;
7
8use crate::{
9 dto,
10 models::{
11 contract::{AccountBalance, AccountChangesWithTx, AccountDelta},
12 protocol::{
13 ComponentBalance, ProtocolChangesWithTx, ProtocolComponent, ProtocolComponentStateDelta,
14 },
15 token::Token,
16 Address, Balance, BlockHash, Chain, Code, ComponentId, EntryPointId, MergeError, StoreKey,
17 StoreVal,
18 },
19 Bytes,
20};
21
22#[derive(Clone, Default, PartialEq, Serialize, Deserialize, Debug)]
23pub struct Block {
24 pub number: u64,
25 pub chain: Chain,
26 pub hash: Bytes,
27 pub parent_hash: Bytes,
28 pub ts: NaiveDateTime,
29}
30
31impl Block {
32 pub fn new(
33 number: u64,
34 chain: Chain,
35 hash: Bytes,
36 parent_hash: Bytes,
37 ts: NaiveDateTime,
38 ) -> Self {
39 Block { hash, parent_hash, number, chain, ts }
40 }
41}
42
43impl DeepSizeOf for Block {
45 fn deep_size_of_children(&self, context: &mut deepsize::Context) -> usize {
46 self.chain
47 .deep_size_of_children(context) +
48 self.hash.deep_size_of_children(context) +
49 self.parent_hash
50 .deep_size_of_children(context)
51 }
52}
53
54#[derive(Clone, Default, PartialEq, Debug, Eq, Hash, DeepSizeOf)]
55pub struct Transaction {
56 pub hash: Bytes,
57 pub block_hash: Bytes,
58 pub from: Bytes,
59 pub to: Option<Bytes>,
60 pub index: u64,
61}
62
63impl Transaction {
64 pub fn new(hash: Bytes, block_hash: Bytes, from: Bytes, to: Option<Bytes>, index: u64) -> Self {
65 Transaction { hash, block_hash, from, to, index }
66 }
67}
68
69pub struct BlockTransactionDeltas<T> {
70 pub extractor: String,
71 pub chain: Chain,
72 pub block: Block,
73 pub revert: bool,
74 pub deltas: Vec<TransactionDeltaGroup<T>>,
75}
76
77#[allow(dead_code)]
78pub struct TransactionDeltaGroup<T> {
79 changes: T,
80 protocol_component: HashMap<String, ProtocolComponent>,
81 component_balances: HashMap<String, ComponentBalance>,
82 component_tvl: HashMap<String, f64>,
83 tx: Transaction,
84}
85
86#[derive(Debug, Default, Clone, Serialize, Deserialize, PartialEq, DeepSizeOf)]
87pub struct BlockAggregatedChanges {
88 pub extractor: String,
89 pub chain: Chain,
90 pub block: Block,
91 pub finalized_block_height: u64,
92 pub db_committed_block_height: Option<u64>,
93 pub revert: bool,
94 pub state_deltas: HashMap<String, ProtocolComponentStateDelta>,
95 pub account_deltas: HashMap<Bytes, AccountDelta>,
96 pub new_tokens: HashMap<Address, Token>,
97 pub new_protocol_components: HashMap<String, ProtocolComponent>,
98 pub deleted_protocol_components: HashMap<String, ProtocolComponent>,
99 pub component_balances: HashMap<ComponentId, HashMap<Bytes, ComponentBalance>>,
100 pub account_balances: HashMap<Address, HashMap<Address, AccountBalance>>,
101 pub component_tvl: HashMap<String, f64>,
102 pub dci_update: DCIUpdate,
103 #[serde(default, skip_serializing_if = "Option::is_none")]
105 pub partial_block_index: Option<u32>,
106}
107
108impl BlockAggregatedChanges {
109 #[allow(clippy::too_many_arguments)]
110 pub fn new(
111 extractor: &str,
112 chain: Chain,
113 block: Block,
114 db_committed_block_height: Option<u64>,
115 finalized_block_height: u64,
116 revert: bool,
117 state_deltas: HashMap<String, ProtocolComponentStateDelta>,
118 account_deltas: HashMap<Bytes, AccountDelta>,
119 new_tokens: HashMap<Address, Token>,
120 new_components: HashMap<String, ProtocolComponent>,
121 deleted_components: HashMap<String, ProtocolComponent>,
122 component_balances: HashMap<ComponentId, HashMap<Bytes, ComponentBalance>>,
123 account_balances: HashMap<Address, HashMap<Address, AccountBalance>>,
124 component_tvl: HashMap<String, f64>,
125 dci_update: DCIUpdate,
126 ) -> Self {
127 Self {
128 extractor: extractor.to_string(),
129 chain,
130 block,
131 db_committed_block_height,
132 finalized_block_height,
133 revert,
134 state_deltas,
135 account_deltas,
136 new_tokens,
137 new_protocol_components: new_components,
138 deleted_protocol_components: deleted_components,
139 component_balances,
140 account_balances,
141 component_tvl,
142 dci_update,
143 partial_block_index: None,
144 }
145 }
146
147 pub fn drop_state(&self) -> Self {
148 Self {
149 extractor: self.extractor.clone(),
150 chain: self.chain,
151 block: self.block.clone(),
152 db_committed_block_height: self.db_committed_block_height,
153 finalized_block_height: self.finalized_block_height,
154 revert: self.revert,
155 account_deltas: HashMap::new(),
156 state_deltas: HashMap::new(),
157 new_tokens: self.new_tokens.clone(),
158 new_protocol_components: self.new_protocol_components.clone(),
159 deleted_protocol_components: self.deleted_protocol_components.clone(),
160 component_balances: self.component_balances.clone(),
161 account_balances: self.account_balances.clone(),
162 component_tvl: self.component_tvl.clone(),
163 dci_update: self.dci_update.clone(),
164 partial_block_index: self.partial_block_index,
165 }
166 }
167
168 pub fn is_partial(&self) -> bool {
169 self.partial_block_index.is_some()
170 }
171}
172
173impl std::fmt::Display for BlockAggregatedChanges {
174 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
175 write!(f, "block_number: {}, extractor: {}", self.block.number, self.extractor)
176 }
177}
178
179pub trait BlockScoped {
180 fn block(&self) -> Block;
181}
182
183impl BlockScoped for BlockAggregatedChanges {
184 fn block(&self) -> Block {
185 self.block.clone()
186 }
187}
188
189impl From<dto::Block> for Block {
190 fn from(value: dto::Block) -> Self {
191 Self {
192 number: value.number,
193 chain: value.chain.into(),
194 hash: value.hash,
195 parent_hash: value.parent_hash,
196 ts: value.ts,
197 }
198 }
199}
200
201#[derive(Debug, Default, Clone, Serialize, Deserialize, PartialEq, DeepSizeOf)]
202pub struct DCIUpdate {
203 pub new_entrypoints: HashMap<ComponentId, HashSet<EntryPoint>>,
204 pub new_entrypoint_params: HashMap<EntryPointId, HashSet<(TracingParams, ComponentId)>>,
205 pub trace_results: HashMap<EntryPointId, TracingResult>,
206}
207
208#[derive(Debug, Clone, PartialEq, Default, DeepSizeOf)]
210pub struct TxWithChanges {
211 pub tx: Transaction,
212 pub protocol_components: HashMap<ComponentId, ProtocolComponent>,
213 pub account_deltas: HashMap<Address, AccountDelta>,
214 pub state_updates: HashMap<ComponentId, ProtocolComponentStateDelta>,
215 pub balance_changes: HashMap<ComponentId, HashMap<Address, ComponentBalance>>,
216 pub account_balance_changes: HashMap<Address, HashMap<Address, AccountBalance>>,
217 pub entrypoints: HashMap<ComponentId, HashSet<EntryPoint>>,
218 pub entrypoint_params: HashMap<EntryPointId, HashSet<(TracingParams, ComponentId)>>,
219}
220
221impl TxWithChanges {
222 #[allow(clippy::too_many_arguments)]
223 pub fn new(
224 tx: Transaction,
225 protocol_components: HashMap<ComponentId, ProtocolComponent>,
226 account_deltas: HashMap<Address, AccountDelta>,
227 protocol_states: HashMap<ComponentId, ProtocolComponentStateDelta>,
228 balance_changes: HashMap<ComponentId, HashMap<Address, ComponentBalance>>,
229 account_balance_changes: HashMap<Address, HashMap<Address, AccountBalance>>,
230 entrypoints: HashMap<ComponentId, HashSet<EntryPoint>>,
231 entrypoint_params: HashMap<EntryPointId, HashSet<(TracingParams, ComponentId)>>,
232 ) -> Self {
233 Self {
234 tx,
235 account_deltas,
236 protocol_components,
237 state_updates: protocol_states,
238 balance_changes,
239 account_balance_changes,
240 entrypoints,
241 entrypoint_params,
242 }
243 }
244
245 pub fn merge(&mut self, other: TxWithChanges) -> Result<(), MergeError> {
255 if self.tx.block_hash != other.tx.block_hash {
256 return Err(MergeError::BlockMismatch(
257 "TxWithChanges".to_string(),
258 self.tx.block_hash.clone(),
259 other.tx.block_hash,
260 ));
261 }
262 if self.tx.index > other.tx.index {
263 return Err(MergeError::TransactionOrderError(
264 "TxWithChanges".to_string(),
265 self.tx.index,
266 other.tx.index,
267 ));
268 }
269
270 self.tx = other.tx;
271
272 for (key, value) in other.protocol_components {
276 match self.protocol_components.entry(key) {
277 Entry::Occupied(mut entry) => {
278 warn!(
279 "Overwriting new protocol component for id {} with a new one. This should never happen! Please check logic",
280 entry.get().id
281 );
282 entry.insert(value);
283 }
284 Entry::Vacant(entry) => {
285 entry.insert(value);
286 }
287 }
288 }
289
290 for (address, update) in other.account_deltas.clone().into_iter() {
292 match self.account_deltas.entry(address) {
293 Entry::Occupied(mut e) => {
294 e.get_mut().merge(update)?;
295 }
296 Entry::Vacant(e) => {
297 e.insert(update);
298 }
299 }
300 }
301
302 for (key, value) in other.state_updates {
304 match self.state_updates.entry(key) {
305 Entry::Occupied(mut entry) => {
306 entry.get_mut().merge(value)?;
307 }
308 Entry::Vacant(entry) => {
309 entry.insert(value);
310 }
311 }
312 }
313
314 for (component_id, balance_changes) in other.balance_changes {
316 let token_balances = self
317 .balance_changes
318 .entry(component_id)
319 .or_default();
320 for (token, balance) in balance_changes {
321 token_balances.insert(token, balance);
322 }
323 }
324
325 for (account_addr, balance_changes) in other.account_balance_changes {
327 let token_balances = self
328 .account_balance_changes
329 .entry(account_addr)
330 .or_default();
331 for (token, balance) in balance_changes {
332 token_balances.insert(token, balance);
333 }
334 }
335
336 for (component_id, entrypoints) in other.entrypoints {
338 self.entrypoints
339 .entry(component_id)
340 .or_default()
341 .extend(entrypoints);
342 }
343
344 for (entrypoint_id, params) in other.entrypoint_params {
346 self.entrypoint_params
347 .entry(entrypoint_id)
348 .or_default()
349 .extend(params);
350 }
351
352 Ok(())
353 }
354}
355
356impl From<AccountChangesWithTx> for TxWithChanges {
357 fn from(value: AccountChangesWithTx) -> Self {
358 Self {
359 tx: value.tx,
360 protocol_components: value.protocol_components,
361 account_deltas: value.account_deltas,
362 balance_changes: value.component_balances,
363 account_balance_changes: value.account_balances,
364 ..Default::default()
365 }
366 }
367}
368
369impl From<ProtocolChangesWithTx> for TxWithChanges {
370 fn from(value: ProtocolChangesWithTx) -> Self {
371 Self {
372 tx: value.tx,
373 protocol_components: value.new_protocol_components,
374 state_updates: value.protocol_states,
375 balance_changes: value.balance_changes,
376 ..Default::default()
377 }
378 }
379}
380
381#[derive(Copy, Clone, Debug, PartialEq)]
382pub enum BlockTag {
383 Finalized,
385 Safe,
387 Latest,
389 Earliest,
391 Pending,
393 Number(u64),
395}
396#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, DeepSizeOf)]
397pub struct EntryPoint {
398 pub external_id: String,
400 pub target: Address,
402 pub signature: String,
404}
405
406impl EntryPoint {
407 pub fn new(external_id: String, target: Address, signature: String) -> Self {
408 Self { external_id, target, signature }
409 }
410}
411
412impl From<dto::EntryPoint> for EntryPoint {
413 fn from(value: dto::EntryPoint) -> Self {
414 Self { external_id: value.external_id, target: value.target, signature: value.signature }
415 }
416}
417
418#[derive(Debug, Clone, PartialEq, Eq, Hash, DeepSizeOf)]
420pub struct EntryPointWithTracingParams {
421 pub entry_point: EntryPoint,
423 pub params: TracingParams,
425}
426
427impl From<dto::EntryPointWithTracingParams> for EntryPointWithTracingParams {
428 fn from(value: dto::EntryPointWithTracingParams) -> Self {
429 match value.params {
430 dto::TracingParams::RPCTracer(ref tracer_params) => Self {
431 entry_point: EntryPoint {
432 external_id: value.entry_point.external_id,
433 target: value.entry_point.target,
434 signature: value.entry_point.signature,
435 },
436 params: TracingParams::RPCTracer(RPCTracerParams {
437 caller: tracer_params.caller.clone(),
438 calldata: tracer_params.calldata.clone(),
439 state_overrides: tracer_params
440 .state_overrides
441 .clone()
442 .map(|s| {
443 s.into_iter()
444 .map(|(k, v)| (k, v.into()))
445 .collect()
446 }),
447 prune_addresses: tracer_params.prune_addresses.clone(),
448 }),
449 },
450 }
451 }
452}
453
454impl EntryPointWithTracingParams {
455 pub fn new(entry_point: EntryPoint, params: TracingParams) -> Self {
456 Self { entry_point, params }
457 }
458}
459
460impl std::fmt::Display for EntryPointWithTracingParams {
461 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
462 let tracer_type = match &self.params {
463 TracingParams::RPCTracer(_) => "RPC",
464 };
465 write!(f, "{} [{}]", self.entry_point.external_id, tracer_type)
466 }
467}
468
469#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Eq, Hash, DeepSizeOf)]
470pub enum TracingParams {
473 RPCTracer(RPCTracerParams),
475}
476
477impl std::fmt::Display for TracingParams {
478 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
479 match self {
480 TracingParams::RPCTracer(params) => write!(f, "RPC: {params}"),
481 }
482 }
483}
484
485impl From<dto::TracingParams> for TracingParams {
486 fn from(value: dto::TracingParams) -> Self {
487 match value {
488 dto::TracingParams::RPCTracer(tracer_params) => {
489 TracingParams::RPCTracer(tracer_params.into())
490 }
491 }
492 }
493}
494
495#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Eq, Hash, DeepSizeOf)]
496pub enum StorageOverride {
497 Diff(BTreeMap<StoreKey, StoreVal>),
498 Replace(BTreeMap<StoreKey, StoreVal>),
499}
500
501impl From<dto::StorageOverride> for StorageOverride {
502 fn from(value: dto::StorageOverride) -> Self {
503 match value {
504 dto::StorageOverride::Diff(diff) => StorageOverride::Diff(diff),
505 dto::StorageOverride::Replace(replace) => StorageOverride::Replace(replace),
506 }
507 }
508}
509
510#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Eq, Hash, DeepSizeOf)]
511pub struct AccountOverrides {
512 pub slots: Option<StorageOverride>,
513 pub native_balance: Option<Balance>,
514 pub code: Option<Code>,
515}
516
517impl From<dto::AccountOverrides> for AccountOverrides {
518 fn from(value: dto::AccountOverrides) -> Self {
519 Self {
520 slots: value.slots.map(|s| s.into()),
521 native_balance: value.native_balance,
522 code: value.code,
523 }
524 }
525}
526
527#[derive(Debug, Clone, PartialEq, Deserialize, Eq, Hash, DeepSizeOf)]
528pub struct RPCTracerParams {
529 pub caller: Option<Address>,
532 pub calldata: Bytes,
534 pub state_overrides: Option<BTreeMap<Address, AccountOverrides>>,
536 pub prune_addresses: Option<Vec<Address>>,
539}
540
541impl From<dto::RPCTracerParams> for RPCTracerParams {
542 fn from(value: dto::RPCTracerParams) -> Self {
543 Self {
544 caller: value.caller,
545 calldata: value.calldata,
546 state_overrides: value.state_overrides.map(|overrides| {
547 overrides
548 .into_iter()
549 .map(|(address, account_overrides)| (address, account_overrides.into()))
550 .collect()
551 }),
552 prune_addresses: value.prune_addresses,
553 }
554 }
555}
556
557impl RPCTracerParams {
558 pub fn new(caller: Option<Address>, calldata: Bytes) -> Self {
559 Self { caller, calldata, state_overrides: None, prune_addresses: None }
560 }
561
562 pub fn with_state_overrides(mut self, state: BTreeMap<Address, AccountOverrides>) -> Self {
563 self.state_overrides = Some(state);
564 self
565 }
566
567 pub fn with_prune_addresses(mut self, addresses: Vec<Address>) -> Self {
568 self.prune_addresses = Some(addresses);
569 self
570 }
571}
572
573impl std::fmt::Display for RPCTracerParams {
574 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
575 let caller_str = match &self.caller {
576 Some(addr) => format!("caller={addr}"),
577 None => String::new(),
578 };
579
580 let calldata_str = if self.calldata.len() >= 8 {
581 format!(
582 "calldata=0x{}..({} bytes)",
583 hex::encode(&self.calldata[..8]),
584 self.calldata.len()
585 )
586 } else {
587 format!("calldata={}", self.calldata)
588 };
589
590 let overrides_str = match &self.state_overrides {
591 Some(overrides) if !overrides.is_empty() => {
592 format!(", {} state override(s)", overrides.len())
593 }
594 _ => String::new(),
595 };
596
597 write!(f, "{caller_str}, {calldata_str}{overrides_str}")
598 }
599}
600
601impl Serialize for RPCTracerParams {
603 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
604 where
605 S: Serializer,
606 {
607 let mut field_count = 2;
609 if self.state_overrides.is_some() {
610 field_count += 1;
611 }
612 if self.prune_addresses.is_some() {
613 field_count += 1;
614 }
615
616 let mut state = serializer.serialize_struct("RPCTracerEntryPoint", field_count)?;
617 state.serialize_field("caller", &self.caller)?;
618 state.serialize_field("calldata", &self.calldata)?;
619
620 if let Some(ref overrides) = self.state_overrides {
622 state.serialize_field("state_overrides", overrides)?;
623 }
624 if let Some(ref prune_addrs) = self.prune_addresses {
625 state.serialize_field("prune_addresses", prune_addrs)?;
626 }
627
628 state.end()
629 }
630}
631
632#[derive(
633 Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize, DeepSizeOf,
634)]
635pub struct AddressStorageLocation {
636 pub key: StoreKey,
637 pub offset: u8,
638}
639
640impl AddressStorageLocation {
641 pub fn new(key: StoreKey, offset: u8) -> Self {
642 Self { key, offset }
643 }
644}
645
646#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default, DeepSizeOf)]
647pub struct TracingResult {
648 pub retriggers: HashSet<(Address, AddressStorageLocation)>,
651 pub accessed_slots: HashMap<Address, HashSet<StoreKey>>,
654}
655
656impl TracingResult {
657 pub fn new(
658 retriggers: HashSet<(Address, AddressStorageLocation)>,
659 accessed_slots: HashMap<Address, HashSet<StoreKey>>,
660 ) -> Self {
661 Self { retriggers, accessed_slots }
662 }
663
664 pub fn merge(&mut self, other: TracingResult) {
668 self.retriggers.extend(other.retriggers);
669 for (address, slots) in other.accessed_slots {
670 self.accessed_slots
671 .entry(address)
672 .or_default()
673 .extend(slots);
674 }
675 }
676}
677
678#[derive(Debug, Clone, PartialEq, DeepSizeOf)]
679pub struct TracedEntryPoint {
681 pub entry_point_with_params: EntryPointWithTracingParams,
683 pub detection_block_hash: BlockHash,
685 pub tracing_result: TracingResult,
687}
688
689impl TracedEntryPoint {
690 pub fn new(
691 entry_point_with_params: EntryPointWithTracingParams,
692 detection_block_hash: BlockHash,
693 result: TracingResult,
694 ) -> Self {
695 Self { entry_point_with_params, detection_block_hash, tracing_result: result }
696 }
697
698 pub fn entry_point_id(&self) -> String {
699 self.entry_point_with_params
700 .entry_point
701 .external_id
702 .clone()
703 }
704}
705
706impl std::fmt::Display for TracedEntryPoint {
707 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
708 write!(
709 f,
710 "[{}: {} retriggers, {} accessed addresses]",
711 self.entry_point_id(),
712 self.tracing_result.retriggers.len(),
713 self.tracing_result.accessed_slots.len()
714 )
715 }
716}
717
718#[cfg(test)]
719pub mod fixtures {
720 use std::str::FromStr;
721
722 use rstest::rstest;
723
724 use super::*;
725 use crate::models::ChangeType;
726
727 pub fn transaction01() -> Transaction {
728 Transaction::new(
729 Bytes::zero(32),
730 Bytes::zero(32),
731 Bytes::zero(20),
732 Some(Bytes::zero(20)),
733 10,
734 )
735 }
736
737 pub fn create_transaction(hash: &str, block: &str, index: u64) -> Transaction {
738 Transaction::new(
739 hash.parse().unwrap(),
740 block.parse().unwrap(),
741 Bytes::zero(20),
742 Some(Bytes::zero(20)),
743 index,
744 )
745 }
746
747 #[test]
748 fn test_merge_tx_with_changes() {
749 let base_token = Bytes::from_str("C02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2").unwrap();
750 let quote_token = Bytes::from_str("A0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48").unwrap();
751 let contract_addr = Bytes::from_str("aaaaaaaaa24eeeb8d57d431224f73832bc34f688").unwrap();
752 let tx_hash0 = "0x2f6350a292c0fc918afe67cb893744a080dacb507b0cea4cc07437b8aff23cdb";
753 let tx_hash1 = "0x0d9e0da36cf9f305a189965b248fc79c923619801e8ab5ef158d4fd528a291ad";
754 let block = "0x0000000000000000000000000000000000000000000000000000000000000000";
755 let component = ProtocolComponent::new(
756 "ambient_USDC_ETH",
757 "test",
758 "vm:pool",
759 Chain::Ethereum,
760 vec![base_token.clone(), quote_token.clone()],
761 vec![contract_addr.clone()],
762 Default::default(),
763 Default::default(),
764 Bytes::from_str(tx_hash0).unwrap(),
765 Default::default(),
766 );
767 let account_delta = AccountDelta::new(
768 Chain::Ethereum,
769 contract_addr.clone(),
770 HashMap::new(),
771 None,
772 Some(vec![0, 0, 0, 0].into()),
773 ChangeType::Creation,
774 );
775
776 let mut changes1 = TxWithChanges::new(
777 create_transaction(tx_hash0, block, 1),
778 HashMap::from([(component.id.clone(), component.clone())]),
779 HashMap::from([(contract_addr.clone(), account_delta.clone())]),
780 HashMap::new(),
781 HashMap::from([(
782 component.id.clone(),
783 HashMap::from([(
784 base_token.clone(),
785 ComponentBalance {
786 token: base_token.clone(),
787 balance: Bytes::from(800_u64).lpad(32, 0),
788 balance_float: 800.0,
789 component_id: component.id.clone(),
790 modify_tx: Bytes::from_str(tx_hash0).unwrap(),
791 },
792 )]),
793 )]),
794 HashMap::from([(
795 contract_addr.clone(),
796 HashMap::from([(
797 base_token.clone(),
798 AccountBalance {
799 token: base_token.clone(),
800 balance: Bytes::from(800_u64).lpad(32, 0),
801 modify_tx: Bytes::from_str(tx_hash0).unwrap(),
802 account: contract_addr.clone(),
803 },
804 )]),
805 )]),
806 HashMap::from([(
807 component.id.clone(),
808 HashSet::from([EntryPoint::new(
809 "test".to_string(),
810 contract_addr.clone(),
811 "function()".to_string(),
812 )]),
813 )]),
814 HashMap::from([(
815 "test".to_string(),
816 HashSet::from([(
817 TracingParams::RPCTracer(RPCTracerParams::new(
818 None,
819 Bytes::from_str("0x000001ef").unwrap(),
820 )),
821 component.id.clone(),
822 )]),
823 )]),
824 );
825 let changes2 = TxWithChanges::new(
826 create_transaction(tx_hash1, block, 2),
827 HashMap::from([(
828 component.id.clone(),
829 ProtocolComponent {
830 creation_tx: Bytes::from_str(tx_hash1).unwrap(),
831 ..component.clone()
832 },
833 )]),
834 HashMap::from([(
835 contract_addr.clone(),
836 AccountDelta::new(
837 Chain::Ethereum,
838 contract_addr.clone(),
839 HashMap::from([(vec![0, 0, 0, 0].into(), Some(vec![0, 0, 0, 0].into()))]),
840 None,
841 None,
842 ChangeType::Update,
843 ),
844 )]),
845 HashMap::new(),
846 HashMap::from([(
847 component.id.clone(),
848 HashMap::from([(
849 base_token.clone(),
850 ComponentBalance {
851 token: base_token.clone(),
852 balance: Bytes::from(1000_u64).lpad(32, 0),
853 balance_float: 1000.0,
854 component_id: component.id.clone(),
855 modify_tx: Bytes::from_str(tx_hash1).unwrap(),
856 },
857 )]),
858 )]),
859 HashMap::from([(
860 contract_addr.clone(),
861 HashMap::from([(
862 base_token.clone(),
863 AccountBalance {
864 token: base_token.clone(),
865 balance: Bytes::from(1000_u64).lpad(32, 0),
866 modify_tx: Bytes::from_str(tx_hash1).unwrap(),
867 account: contract_addr.clone(),
868 },
869 )]),
870 )]),
871 HashMap::from([(
872 component.id.clone(),
873 HashSet::from([
874 EntryPoint::new(
875 "test".to_string(),
876 contract_addr.clone(),
877 "function()".to_string(),
878 ),
879 EntryPoint::new(
880 "test2".to_string(),
881 contract_addr.clone(),
882 "function_2()".to_string(),
883 ),
884 ]),
885 )]),
886 HashMap::from([(
887 "test2".to_string(),
888 HashSet::from([(
889 TracingParams::RPCTracer(RPCTracerParams::new(
890 None,
891 Bytes::from_str("0x000001").unwrap(),
892 )),
893 component.id.clone(),
894 )]),
895 )]),
896 );
897
898 assert!(changes1.merge(changes2).is_ok());
899 assert_eq!(
900 changes1
901 .account_balance_changes
902 .get(&contract_addr)
903 .unwrap()
904 .get(&base_token)
905 .unwrap()
906 .balance,
907 Bytes::from(1000_u64).lpad(32, 0),
908 );
909 assert_eq!(
910 changes1
911 .balance_changes
912 .get(&component.id)
913 .unwrap()
914 .get(&base_token)
915 .unwrap()
916 .balance,
917 Bytes::from(1000_u64).lpad(32, 0),
918 );
919 assert_eq!(changes1.tx.hash, Bytes::from_str(tx_hash1).unwrap(),);
920 assert_eq!(changes1.entrypoints.len(), 1);
921 assert_eq!(
922 changes1
923 .entrypoints
924 .get(&component.id)
925 .unwrap()
926 .len(),
927 2
928 );
929 let mut expected_entry_points = changes1
930 .entrypoints
931 .values()
932 .flat_map(|ep| ep.iter())
933 .map(|ep| ep.signature.clone())
934 .collect::<Vec<_>>();
935 expected_entry_points.sort();
936 assert_eq!(
937 expected_entry_points,
938 vec!["function()".to_string(), "function_2()".to_string()],
939 );
940 }
941
942 #[rstest]
943 #[case::mismatched_blocks(
944 fixtures::create_transaction("0x01", "0x0abc", 1),
945 fixtures::create_transaction("0x02", "0x0def", 2)
946 )]
947 #[case::older_transaction(
948 fixtures::create_transaction("0x02", "0x0abc", 2),
949 fixtures::create_transaction("0x01", "0x0abc", 1)
950 )]
951 fn test_merge_errors(#[case] tx1: Transaction, #[case] tx2: Transaction) {
952 let mut changes1 = TxWithChanges { tx: tx1, ..Default::default() };
953
954 let changes2 = TxWithChanges { tx: tx2, ..Default::default() };
955
956 assert!(changes1.merge(changes2).is_err());
957 }
958
959 #[test]
960 fn test_rpc_tracer_entry_point_serialization_order() {
961 use std::str::FromStr;
962
963 use serde_json;
964
965 let entry_point = RPCTracerParams::new(
966 Some(Address::from_str("0x1234567890123456789012345678901234567890").unwrap()),
967 Bytes::from_str("0xabcdef").unwrap(),
968 );
969
970 let serialized = serde_json::to_string(&entry_point).unwrap();
971
972 assert!(serialized.find("\"caller\"").unwrap() < serialized.find("\"calldata\"").unwrap());
974
975 let deserialized: RPCTracerParams = serde_json::from_str(&serialized).unwrap();
977 assert_eq!(entry_point, deserialized);
978 }
979
980 #[test]
981 fn test_tracing_result_merge() {
982 let address1 = Address::from_str("0x1234567890123456789012345678901234567890").unwrap();
983 let address2 = Address::from_str("0x2345678901234567890123456789012345678901").unwrap();
984 let address3 = Address::from_str("0x3456789012345678901234567890123456789012").unwrap();
985
986 let store_key1 = StoreKey::from(vec![1, 2, 3, 4]);
987 let store_key2 = StoreKey::from(vec![5, 6, 7, 8]);
988
989 let mut result1 = TracingResult::new(
990 HashSet::from([(
991 address1.clone(),
992 AddressStorageLocation::new(store_key1.clone(), 12),
993 )]),
994 HashMap::from([
995 (address2.clone(), HashSet::from([store_key1.clone()])),
996 (address3.clone(), HashSet::from([store_key2.clone()])),
997 ]),
998 );
999
1000 let result2 = TracingResult::new(
1001 HashSet::from([(
1002 address3.clone(),
1003 AddressStorageLocation::new(store_key2.clone(), 12),
1004 )]),
1005 HashMap::from([
1006 (address1.clone(), HashSet::from([store_key1.clone()])),
1007 (address2.clone(), HashSet::from([store_key2.clone()])),
1008 ]),
1009 );
1010
1011 result1.merge(result2);
1012
1013 assert_eq!(result1.retriggers.len(), 2);
1015 assert!(result1
1016 .retriggers
1017 .contains(&(address1.clone(), AddressStorageLocation::new(store_key1.clone(), 12))));
1018 assert!(result1
1019 .retriggers
1020 .contains(&(address3.clone(), AddressStorageLocation::new(store_key2.clone(), 12))));
1021
1022 assert_eq!(result1.accessed_slots.len(), 3);
1024 assert!(result1
1025 .accessed_slots
1026 .contains_key(&address1));
1027 assert!(result1
1028 .accessed_slots
1029 .contains_key(&address2));
1030 assert!(result1
1031 .accessed_slots
1032 .contains_key(&address3));
1033
1034 assert_eq!(
1035 result1
1036 .accessed_slots
1037 .get(&address2)
1038 .unwrap(),
1039 &HashSet::from([store_key1.clone(), store_key2.clone()])
1040 );
1041 }
1042
1043 #[test]
1044 fn test_entry_point_with_tracing_params_display() {
1045 use std::str::FromStr;
1046
1047 let entry_point = EntryPoint::new(
1048 "uniswap_v3_pool_swap".to_string(),
1049 Address::from_str("0x1234567890123456789012345678901234567890").unwrap(),
1050 "swapExactETHForTokens(uint256,address[],address,uint256)".to_string(),
1051 );
1052
1053 let tracing_params = TracingParams::RPCTracer(RPCTracerParams::new(
1054 Some(Address::from_str("0x9876543210987654321098765432109876543210").unwrap()),
1055 Bytes::from_str("0xabcdef").unwrap(),
1056 ));
1057
1058 let entry_point_with_params = EntryPointWithTracingParams::new(entry_point, tracing_params);
1059
1060 let display_output = entry_point_with_params.to_string();
1061 assert_eq!(display_output, "uniswap_v3_pool_swap [RPC]");
1062 }
1063
1064 #[test]
1065 fn test_traced_entry_point_display() {
1066 use std::str::FromStr;
1067
1068 let entry_point = EntryPoint::new(
1069 "uniswap_v3_pool_swap".to_string(),
1070 Address::from_str("0x1234567890123456789012345678901234567890").unwrap(),
1071 "swapExactETHForTokens(uint256,address[],address,uint256)".to_string(),
1072 );
1073
1074 let tracing_params = TracingParams::RPCTracer(RPCTracerParams::new(
1075 Some(Address::from_str("0x9876543210987654321098765432109876543210").unwrap()),
1076 Bytes::from_str("0xabcdef").unwrap(),
1077 ));
1078
1079 let entry_point_with_params = EntryPointWithTracingParams::new(entry_point, tracing_params);
1080
1081 let address1 = Address::from_str("0x1111111111111111111111111111111111111111").unwrap();
1083 let address2 = Address::from_str("0x2222222222222222222222222222222222222222").unwrap();
1084 let address3 = Address::from_str("0x3333333333333333333333333333333333333333").unwrap();
1085
1086 let store_key1 = StoreKey::from(vec![1, 2, 3, 4]);
1087 let store_key2 = StoreKey::from(vec![5, 6, 7, 8]);
1088
1089 let tracing_result = TracingResult::new(
1090 HashSet::from([
1091 (address1.clone(), AddressStorageLocation::new(store_key1.clone(), 0)),
1092 (address2.clone(), AddressStorageLocation::new(store_key2.clone(), 12)),
1093 ]),
1094 HashMap::from([
1095 (address1.clone(), HashSet::from([store_key1.clone()])),
1096 (address2.clone(), HashSet::from([store_key2.clone()])),
1097 (address3.clone(), HashSet::from([store_key1.clone()])),
1098 ]),
1099 );
1100
1101 let traced_entry_point = TracedEntryPoint::new(
1102 entry_point_with_params,
1103 Bytes::from_str("0xabcdef1234567890").unwrap(),
1104 tracing_result,
1105 );
1106
1107 let display_output = traced_entry_point.to_string();
1108 assert_eq!(display_output, "[uniswap_v3_pool_swap: 2 retriggers, 3 accessed addresses]");
1109 }
1110}