Skip to main content

tycho_common/models/
blockchain.rs

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
43// Manual impl as `NaiveDateTime` structure referenced in `ts` does not implement DeepSizeOf
44impl 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    /// The index of the partial block. None if it's a full block.
104    #[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/// Changes grouped by their respective transaction.
209#[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    /// Merges this update with another one.
246    ///
247    /// The method combines two [`TxWithChanges`] instances if they are on the same block.
248    ///
249    /// NB: It is expected that `other` is a more recent update than `self` is and the two are
250    /// combined accordingly.
251    ///
252    /// # Errors
253    /// Returns a `MergeError` if any of the above conditions are violated.
254    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        // Merge new protocol components
273        // Log a warning if a new protocol component for the same id already exists, because this
274        // should never happen.
275        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        // Merge account deltas
291        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        // Merge protocol state updates
303        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        // Merge component balance changes
315        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        // Merge account balance changes
326        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        // Merge new entrypoints
337        for (component_id, entrypoints) in other.entrypoints {
338            self.entrypoints
339                .entry(component_id)
340                .or_default()
341                .extend(entrypoints);
342        }
343
344        // Merge new entrypoint params
345        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 block
384    Finalized,
385    /// Safe block
386    Safe,
387    /// Latest block
388    Latest,
389    /// Earliest block (genesis)
390    Earliest,
391    /// Pending block (not yet part of the blockchain)
392    Pending,
393    /// Block by number
394    Number(u64),
395}
396#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, DeepSizeOf)]
397pub struct EntryPoint {
398    /// Entry point id
399    pub external_id: String,
400    /// The address of the contract to trace.
401    pub target: Address,
402    /// The signature of the function to trace.
403    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/// A struct that combines an entry point with its associated tracing params.
419#[derive(Debug, Clone, PartialEq, Eq, Hash, DeepSizeOf)]
420pub struct EntryPointWithTracingParams {
421    /// The entry point to trace, containing the target contract address and function signature
422    pub entry_point: EntryPoint,
423    /// The tracing parameters for this entry point
424    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)]
470/// An entry point to trace. Different types of entry points tracing will be supported in the
471/// future. Like RPC debug tracing, symbolic execution, etc.
472pub enum TracingParams {
473    /// Uses RPC calls to retrieve the called addresses and retriggers
474    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    /// The caller address of the transaction, if not provided tracing will use the default value
530    /// for an address defined by the VM.
531    pub caller: Option<Address>,
532    /// The call data used for the tracing call, this needs to include the function selector
533    pub calldata: Bytes,
534    /// Optionally allow for state overrides so that the call works as expected
535    pub state_overrides: Option<BTreeMap<Address, AccountOverrides>>,
536    /// Addresses to prune from trace results. Useful for hooks that use mock
537    /// accounts/routers that shouldn't be tracked in the final DCI results.
538    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
601// Ensure serialization order, required by the storage layer
602impl Serialize for RPCTracerParams {
603    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
604    where
605        S: Serializer,
606    {
607        // Count fields: always serialize caller and calldata, plus optional fields
608        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        // Only serialize optional fields if they are present
621        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    /// A set of (address, storage slot) pairs representing state that contain a called address.
649    /// If any of these storage slots change, the execution path might change.
650    pub retriggers: HashSet<(Address, AddressStorageLocation)>,
651    /// A map of all addresses that were called during the trace with a list of storage slots that
652    /// were accessed.
653    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    /// Merges this tracing result with another one.
665    ///
666    /// The method combines two [`TracingResult`] instances.
667    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)]
679/// Represents a traced entry point and the results of the tracing operation.
680pub struct TracedEntryPoint {
681    /// The combined entry point and tracing params that was traced
682    pub entry_point_with_params: EntryPointWithTracingParams,
683    /// The block hash of the block that the entry point was traced on.
684    pub detection_block_hash: BlockHash,
685    /// The results of the tracing operation
686    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        // Verify that "caller" comes before "calldata" in the serialized output
973        assert!(serialized.find("\"caller\"").unwrap() < serialized.find("\"calldata\"").unwrap());
974
975        // Verify we can deserialize it back
976        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        // Verify retriggers were merged
1014        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        // Verify accessed slots were merged
1023        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        // Create tracing result with 2 retriggers and 3 accessed addresses
1082        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}