tycho_common/
dto.rs

1//! Data Transfer Objects (or structs)
2//!
3//! These structs serve to serialise and deserialize messages between server and client, they should
4//! be very simple and ideally not contain any business logic.
5//!
6//! Structs in here implement utoipa traits so they can be used to derive an OpenAPI schema.
7#![allow(deprecated)]
8use std::{
9    collections::{BTreeMap, HashMap, HashSet},
10    fmt,
11    hash::{Hash, Hasher},
12};
13
14use chrono::{NaiveDateTime, Utc};
15use deepsize::{Context, DeepSizeOf};
16use serde::{de, Deserialize, Deserializer, Serialize};
17use strum_macros::{Display, EnumString};
18use thiserror::Error;
19use utoipa::{IntoParams, ToSchema};
20use uuid::Uuid;
21
22use crate::{
23    models::{
24        self, blockchain::BlockAggregatedChanges, Address, Balance, Code, ComponentId, StoreKey,
25        StoreVal,
26    },
27    serde_primitives::{
28        hex_bytes, hex_bytes_option, hex_hashmap_key, hex_hashmap_key_value, hex_hashmap_value,
29    },
30    Bytes,
31};
32
33/// Currently supported Blockchains
34#[derive(
35    Debug,
36    Clone,
37    Copy,
38    PartialEq,
39    Eq,
40    Hash,
41    Serialize,
42    Deserialize,
43    EnumString,
44    Display,
45    Default,
46    ToSchema,
47    DeepSizeOf,
48)]
49#[serde(rename_all = "lowercase")]
50#[strum(serialize_all = "lowercase")]
51pub enum Chain {
52    #[default]
53    Ethereum,
54    Starknet,
55    ZkSync,
56    Arbitrum,
57    Base,
58    Bsc,
59    Unichain,
60}
61
62impl From<models::contract::Account> for ResponseAccount {
63    fn from(value: models::contract::Account) -> Self {
64        ResponseAccount::new(
65            value.chain.into(),
66            value.address,
67            value.title,
68            value.slots,
69            value.native_balance,
70            value
71                .token_balances
72                .into_iter()
73                .map(|(k, v)| (k, v.balance))
74                .collect(),
75            value.code,
76            value.code_hash,
77            value.balance_modify_tx,
78            value.code_modify_tx,
79            value.creation_tx,
80        )
81    }
82}
83
84impl From<models::Chain> for Chain {
85    fn from(value: models::Chain) -> Self {
86        match value {
87            models::Chain::Ethereum => Chain::Ethereum,
88            models::Chain::Starknet => Chain::Starknet,
89            models::Chain::ZkSync => Chain::ZkSync,
90            models::Chain::Arbitrum => Chain::Arbitrum,
91            models::Chain::Base => Chain::Base,
92            models::Chain::Bsc => Chain::Bsc,
93            models::Chain::Unichain => Chain::Unichain,
94        }
95    }
96}
97
98#[derive(
99    Debug,
100    PartialEq,
101    Default,
102    Copy,
103    Clone,
104    Deserialize,
105    Serialize,
106    ToSchema,
107    EnumString,
108    Display,
109    DeepSizeOf,
110)]
111pub enum ChangeType {
112    #[default]
113    Update,
114    Deletion,
115    Creation,
116    Unspecified,
117}
118
119impl From<models::ChangeType> for ChangeType {
120    fn from(value: models::ChangeType) -> Self {
121        match value {
122            models::ChangeType::Update => ChangeType::Update,
123            models::ChangeType::Creation => ChangeType::Creation,
124            models::ChangeType::Deletion => ChangeType::Deletion,
125        }
126    }
127}
128
129impl ChangeType {
130    pub fn merge(&self, other: &Self) -> Self {
131        if matches!(self, Self::Creation) {
132            Self::Creation
133        } else {
134            *other
135        }
136    }
137}
138
139#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, Hash, Default)]
140pub struct ExtractorIdentity {
141    pub chain: Chain,
142    pub name: String,
143}
144
145impl ExtractorIdentity {
146    pub fn new(chain: Chain, name: &str) -> Self {
147        Self { chain, name: name.to_owned() }
148    }
149}
150
151impl fmt::Display for ExtractorIdentity {
152    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
153        write!(f, "{}:{}", self.chain, self.name)
154    }
155}
156
157/// A command sent from the client to the server
158#[derive(Deserialize, Serialize, Debug, PartialEq, Eq)]
159#[serde(tag = "method", rename_all = "lowercase")]
160pub enum Command {
161    Subscribe {
162        extractor_id: ExtractorIdentity,
163        include_state: bool,
164        /// Enable zstd compression for messages in this subscription.
165        /// Defaults to false for backward compatibility.
166        #[serde(default)]
167        compression: bool,
168    },
169    Unsubscribe {
170        subscription_id: Uuid,
171    },
172}
173
174/// A easy serializable version of `models::error::WebsocketError`
175///
176/// This serves purely to transfer errors via websocket. It is meant to render
177/// similarly to the original struct but does not have server side debug information
178/// attached.
179///
180/// It should contain information needed to handle errors correctly on the client side.
181#[derive(Error, Debug, Serialize, Deserialize, Eq, PartialEq, Clone)]
182pub enum WebsocketError {
183    #[error("Extractor not found: {0}")]
184    ExtractorNotFound(ExtractorIdentity),
185
186    #[error("Subscription not found: {0}")]
187    SubscriptionNotFound(Uuid),
188
189    #[error("Failed to parse JSON: {1}, msg: {0}")]
190    ParseError(String, String),
191
192    #[error("Failed to subscribe to extractor: {0}")]
193    SubscribeError(ExtractorIdentity),
194
195    #[error("Failed to compress message for subscription: {0}, error: {1}")]
196    CompressionError(Uuid, String),
197}
198
199impl From<crate::models::error::WebsocketError> for WebsocketError {
200    fn from(value: crate::models::error::WebsocketError) -> Self {
201        match value {
202            crate::models::error::WebsocketError::ExtractorNotFound(eid) => {
203                Self::ExtractorNotFound(eid.into())
204            }
205            crate::models::error::WebsocketError::SubscriptionNotFound(sid) => {
206                Self::SubscriptionNotFound(sid)
207            }
208            crate::models::error::WebsocketError::ParseError(raw, error) => {
209                Self::ParseError(error.to_string(), raw)
210            }
211            crate::models::error::WebsocketError::SubscribeError(eid) => {
212                Self::SubscribeError(eid.into())
213            }
214            crate::models::error::WebsocketError::CompressionError(sid, error) => {
215                Self::CompressionError(sid, error.to_string())
216            }
217        }
218    }
219}
220
221/// A response sent from the server to the client
222#[derive(Deserialize, Serialize, Debug, PartialEq, Eq, Clone)]
223#[serde(tag = "method", rename_all = "lowercase")]
224pub enum Response {
225    NewSubscription { extractor_id: ExtractorIdentity, subscription_id: Uuid },
226    SubscriptionEnded { subscription_id: Uuid },
227    Error(WebsocketError),
228}
229
230/// A message sent from the server to the client
231#[allow(clippy::large_enum_variant)]
232#[derive(Serialize, Deserialize, Debug, Display, Clone)]
233#[serde(untagged)]
234pub enum WebSocketMessage {
235    BlockChanges { subscription_id: Uuid, deltas: BlockChanges },
236    Response(Response),
237}
238
239#[derive(Debug, PartialEq, Clone, Deserialize, Serialize, Default, ToSchema)]
240pub struct Block {
241    pub number: u64,
242    #[serde(with = "hex_bytes")]
243    pub hash: Bytes,
244    #[serde(with = "hex_bytes")]
245    pub parent_hash: Bytes,
246    pub chain: Chain,
247    pub ts: NaiveDateTime,
248}
249
250#[derive(Debug, Serialize, Deserialize, PartialEq, Clone, ToSchema, Eq, Hash, DeepSizeOf)]
251#[serde(deny_unknown_fields)]
252pub struct BlockParam {
253    #[schema(value_type=Option<String>)]
254    #[serde(with = "hex_bytes_option", default)]
255    pub hash: Option<Bytes>,
256    #[deprecated(
257        note = "The `chain` field is deprecated and will be removed in a future version."
258    )]
259    #[serde(default)]
260    pub chain: Option<Chain>,
261    #[serde(default)]
262    pub number: Option<i64>,
263}
264
265impl From<&Block> for BlockParam {
266    fn from(value: &Block) -> Self {
267        // The hash should uniquely identify a block across chains
268        BlockParam { hash: Some(value.hash.clone()), chain: None, number: None }
269    }
270}
271
272#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, Default)]
273pub struct TokenBalances(#[serde(with = "hex_hashmap_key")] pub HashMap<Bytes, ComponentBalance>);
274
275impl From<HashMap<Bytes, ComponentBalance>> for TokenBalances {
276    fn from(value: HashMap<Bytes, ComponentBalance>) -> Self {
277        TokenBalances(value)
278    }
279}
280
281#[derive(Debug, PartialEq, Clone, Default, Deserialize, Serialize)]
282pub struct Transaction {
283    #[serde(with = "hex_bytes")]
284    pub hash: Bytes,
285    #[serde(with = "hex_bytes")]
286    pub block_hash: Bytes,
287    #[serde(with = "hex_bytes")]
288    pub from: Bytes,
289    #[serde(with = "hex_bytes_option")]
290    pub to: Option<Bytes>,
291    pub index: u64,
292}
293
294impl Transaction {
295    pub fn new(hash: Bytes, block_hash: Bytes, from: Bytes, to: Option<Bytes>, index: u64) -> Self {
296        Self { hash, block_hash, from, to, index }
297    }
298}
299
300/// A container for updates grouped by account/component.
301#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, Default)]
302pub struct BlockChanges {
303    pub extractor: String,
304    pub chain: Chain,
305    pub block: Block,
306    pub finalized_block_height: u64,
307    pub revert: bool,
308    #[serde(with = "hex_hashmap_key", default)]
309    pub new_tokens: HashMap<Bytes, ResponseToken>,
310    #[serde(alias = "account_deltas", with = "hex_hashmap_key")]
311    pub account_updates: HashMap<Bytes, AccountUpdate>,
312    #[serde(alias = "state_deltas")]
313    pub state_updates: HashMap<String, ProtocolStateDelta>,
314    pub new_protocol_components: HashMap<String, ProtocolComponent>,
315    pub deleted_protocol_components: HashMap<String, ProtocolComponent>,
316    pub component_balances: HashMap<String, TokenBalances>,
317    pub account_balances: HashMap<Bytes, HashMap<Bytes, AccountBalance>>,
318    pub component_tvl: HashMap<String, f64>,
319    pub dci_update: DCIUpdate,
320}
321
322impl BlockChanges {
323    #[allow(clippy::too_many_arguments)]
324    pub fn new(
325        extractor: &str,
326        chain: Chain,
327        block: Block,
328        finalized_block_height: u64,
329        revert: bool,
330        account_updates: HashMap<Bytes, AccountUpdate>,
331        state_updates: HashMap<String, ProtocolStateDelta>,
332        new_protocol_components: HashMap<String, ProtocolComponent>,
333        deleted_protocol_components: HashMap<String, ProtocolComponent>,
334        component_balances: HashMap<String, HashMap<Bytes, ComponentBalance>>,
335        account_balances: HashMap<Bytes, HashMap<Bytes, AccountBalance>>,
336        dci_update: DCIUpdate,
337    ) -> Self {
338        BlockChanges {
339            extractor: extractor.to_owned(),
340            chain,
341            block,
342            finalized_block_height,
343            revert,
344            new_tokens: HashMap::new(),
345            account_updates,
346            state_updates,
347            new_protocol_components,
348            deleted_protocol_components,
349            component_balances: component_balances
350                .into_iter()
351                .map(|(k, v)| (k, v.into()))
352                .collect(),
353            account_balances,
354            component_tvl: HashMap::new(),
355            dci_update,
356        }
357    }
358
359    pub fn merge(mut self, other: Self) -> Self {
360        other
361            .account_updates
362            .into_iter()
363            .for_each(|(k, v)| {
364                self.account_updates
365                    .entry(k)
366                    .and_modify(|e| {
367                        e.merge(&v);
368                    })
369                    .or_insert(v);
370            });
371
372        other
373            .state_updates
374            .into_iter()
375            .for_each(|(k, v)| {
376                self.state_updates
377                    .entry(k)
378                    .and_modify(|e| {
379                        e.merge(&v);
380                    })
381                    .or_insert(v);
382            });
383
384        other
385            .component_balances
386            .into_iter()
387            .for_each(|(k, v)| {
388                self.component_balances
389                    .entry(k)
390                    .and_modify(|e| e.0.extend(v.0.clone()))
391                    .or_insert_with(|| v);
392            });
393
394        other
395            .account_balances
396            .into_iter()
397            .for_each(|(k, v)| {
398                self.account_balances
399                    .entry(k)
400                    .and_modify(|e| e.extend(v.clone()))
401                    .or_insert(v);
402            });
403
404        self.component_tvl
405            .extend(other.component_tvl);
406        self.new_protocol_components
407            .extend(other.new_protocol_components);
408        self.deleted_protocol_components
409            .extend(other.deleted_protocol_components);
410        self.revert = other.revert;
411        self.block = other.block;
412
413        self
414    }
415
416    pub fn get_block(&self) -> &Block {
417        &self.block
418    }
419
420    pub fn is_revert(&self) -> bool {
421        self.revert
422    }
423
424    pub fn filter_by_component<F: Fn(&str) -> bool>(&mut self, keep: F) {
425        self.state_updates
426            .retain(|k, _| keep(k));
427        self.component_balances
428            .retain(|k, _| keep(k));
429        self.component_tvl
430            .retain(|k, _| keep(k));
431    }
432
433    pub fn filter_by_contract<F: Fn(&Bytes) -> bool>(&mut self, keep: F) {
434        self.account_updates
435            .retain(|k, _| keep(k));
436        self.account_balances
437            .retain(|k, _| keep(k));
438    }
439
440    pub fn n_changes(&self) -> usize {
441        self.account_updates.len() + self.state_updates.len()
442    }
443
444    pub fn drop_state(&self) -> Self {
445        Self {
446            extractor: self.extractor.clone(),
447            chain: self.chain,
448            block: self.block.clone(),
449            finalized_block_height: self.finalized_block_height,
450            revert: self.revert,
451            new_tokens: self.new_tokens.clone(),
452            account_updates: HashMap::new(),
453            state_updates: HashMap::new(),
454            new_protocol_components: self.new_protocol_components.clone(),
455            deleted_protocol_components: self.deleted_protocol_components.clone(),
456            component_balances: self.component_balances.clone(),
457            account_balances: self.account_balances.clone(),
458            component_tvl: self.component_tvl.clone(),
459            dci_update: self.dci_update.clone(),
460        }
461    }
462}
463
464impl From<models::blockchain::Block> for Block {
465    fn from(value: models::blockchain::Block) -> Self {
466        Self {
467            number: value.number,
468            hash: value.hash,
469            parent_hash: value.parent_hash,
470            chain: value.chain.into(),
471            ts: value.ts,
472        }
473    }
474}
475
476impl From<models::protocol::ComponentBalance> for ComponentBalance {
477    fn from(value: models::protocol::ComponentBalance) -> Self {
478        Self {
479            token: value.token,
480            balance: value.balance,
481            balance_float: value.balance_float,
482            modify_tx: value.modify_tx,
483            component_id: value.component_id,
484        }
485    }
486}
487
488impl From<models::contract::AccountBalance> for AccountBalance {
489    fn from(value: models::contract::AccountBalance) -> Self {
490        Self {
491            account: value.account,
492            token: value.token,
493            balance: value.balance,
494            modify_tx: value.modify_tx,
495        }
496    }
497}
498
499impl From<BlockAggregatedChanges> for BlockChanges {
500    fn from(value: BlockAggregatedChanges) -> Self {
501        Self {
502            extractor: value.extractor,
503            chain: value.chain.into(),
504            block: value.block.into(),
505            finalized_block_height: value.finalized_block_height,
506            revert: value.revert,
507            account_updates: value
508                .account_deltas
509                .into_iter()
510                .map(|(k, v)| (k, v.into()))
511                .collect(),
512            state_updates: value
513                .state_deltas
514                .into_iter()
515                .map(|(k, v)| (k, v.into()))
516                .collect(),
517            new_protocol_components: value
518                .new_protocol_components
519                .into_iter()
520                .map(|(k, v)| (k, v.into()))
521                .collect(),
522            deleted_protocol_components: value
523                .deleted_protocol_components
524                .into_iter()
525                .map(|(k, v)| (k, v.into()))
526                .collect(),
527            component_balances: value
528                .component_balances
529                .into_iter()
530                .map(|(component_id, v)| {
531                    let balances: HashMap<Bytes, ComponentBalance> = v
532                        .into_iter()
533                        .map(|(k, v)| (k, ComponentBalance::from(v)))
534                        .collect();
535                    (component_id, balances.into())
536                })
537                .collect(),
538            account_balances: value
539                .account_balances
540                .into_iter()
541                .map(|(k, v)| {
542                    (
543                        k,
544                        v.into_iter()
545                            .map(|(k, v)| (k, v.into()))
546                            .collect(),
547                    )
548                })
549                .collect(),
550            dci_update: value.dci_update.into(),
551            new_tokens: value
552                .new_tokens
553                .into_iter()
554                .map(|(k, v)| (k, v.into()))
555                .collect(),
556            component_tvl: value.component_tvl,
557        }
558    }
559}
560
561#[derive(PartialEq, Serialize, Deserialize, Clone, Debug, ToSchema)]
562pub struct AccountUpdate {
563    #[serde(with = "hex_bytes")]
564    #[schema(value_type=Vec<String>)]
565    pub address: Bytes,
566    pub chain: Chain,
567    #[serde(with = "hex_hashmap_key_value")]
568    #[schema(value_type=HashMap<String, String>)]
569    pub slots: HashMap<Bytes, Bytes>,
570    #[serde(with = "hex_bytes_option")]
571    #[schema(value_type=Option<String>)]
572    pub balance: Option<Bytes>,
573    #[serde(with = "hex_bytes_option")]
574    #[schema(value_type=Option<String>)]
575    pub code: Option<Bytes>,
576    pub change: ChangeType,
577}
578
579impl AccountUpdate {
580    pub fn new(
581        address: Bytes,
582        chain: Chain,
583        slots: HashMap<Bytes, Bytes>,
584        balance: Option<Bytes>,
585        code: Option<Bytes>,
586        change: ChangeType,
587    ) -> Self {
588        Self { address, chain, slots, balance, code, change }
589    }
590
591    pub fn merge(&mut self, other: &Self) {
592        self.slots.extend(
593            other
594                .slots
595                .iter()
596                .map(|(k, v)| (k.clone(), v.clone())),
597        );
598        self.balance.clone_from(&other.balance);
599        self.code.clone_from(&other.code);
600        self.change = self.change.merge(&other.change);
601    }
602}
603
604impl From<models::contract::AccountDelta> for AccountUpdate {
605    fn from(value: models::contract::AccountDelta) -> Self {
606        let code = value.code().clone();
607        let change_type = value.change_type().into();
608        AccountUpdate::new(
609            value.address,
610            value.chain.into(),
611            value
612                .slots
613                .into_iter()
614                .map(|(k, v)| (k, v.unwrap_or_default()))
615                .collect(),
616            value.balance,
617            code,
618            change_type,
619        )
620    }
621}
622
623/// Represents the static parts of a protocol component.
624#[derive(Debug, Clone, PartialEq, Default, Deserialize, Serialize, ToSchema)]
625pub struct ProtocolComponent {
626    /// Unique identifier for this component
627    pub id: String,
628    /// Protocol system this component is part of
629    pub protocol_system: String,
630    /// Type of the protocol system
631    pub protocol_type_name: String,
632    pub chain: Chain,
633    /// Token addresses the component operates on
634    #[schema(value_type=Vec<String>)]
635    pub tokens: Vec<Bytes>,
636    /// Contract addresses involved in the components operations (may be empty for
637    /// native implementations)
638    #[serde(alias = "contract_addresses")]
639    #[schema(value_type=Vec<String>)]
640    pub contract_ids: Vec<Bytes>,
641    /// Constant attributes of the component
642    #[serde(with = "hex_hashmap_value")]
643    #[schema(value_type=HashMap<String, String>)]
644    pub static_attributes: HashMap<String, Bytes>,
645    /// Indicates if last change was update, create or delete (for internal use only).
646    #[serde(default)]
647    pub change: ChangeType,
648    /// Transaction hash which created this component
649    #[serde(with = "hex_bytes")]
650    #[schema(value_type=String)]
651    pub creation_tx: Bytes,
652    /// Date time of creation in UTC time
653    pub created_at: NaiveDateTime,
654}
655
656// Manual impl as `NaiveDateTime` structure referenced in `created_at` does not implement DeepSizeOf
657impl DeepSizeOf for ProtocolComponent {
658    fn deep_size_of_children(&self, ctx: &mut Context) -> usize {
659        self.id.deep_size_of_children(ctx) +
660            self.protocol_system
661                .deep_size_of_children(ctx) +
662            self.protocol_type_name
663                .deep_size_of_children(ctx) +
664            self.chain.deep_size_of_children(ctx) +
665            self.tokens.deep_size_of_children(ctx) +
666            self.contract_ids
667                .deep_size_of_children(ctx) +
668            self.static_attributes
669                .deep_size_of_children(ctx) +
670            self.change.deep_size_of_children(ctx) +
671            self.creation_tx
672                .deep_size_of_children(ctx)
673    }
674}
675
676impl From<models::protocol::ProtocolComponent> for ProtocolComponent {
677    fn from(value: models::protocol::ProtocolComponent) -> Self {
678        Self {
679            id: value.id,
680            protocol_system: value.protocol_system,
681            protocol_type_name: value.protocol_type_name,
682            chain: value.chain.into(),
683            tokens: value.tokens,
684            contract_ids: value.contract_addresses,
685            static_attributes: value.static_attributes,
686            change: value.change.into(),
687            creation_tx: value.creation_tx,
688            created_at: value.created_at,
689        }
690    }
691}
692
693#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, Default)]
694pub struct ComponentBalance {
695    #[serde(with = "hex_bytes")]
696    pub token: Bytes,
697    pub balance: Bytes,
698    pub balance_float: f64,
699    #[serde(with = "hex_bytes")]
700    pub modify_tx: Bytes,
701    pub component_id: String,
702}
703
704#[derive(Debug, PartialEq, Clone, Default, Serialize, Deserialize, ToSchema)]
705/// Represents a change in protocol state.
706pub struct ProtocolStateDelta {
707    pub component_id: String,
708    #[schema(value_type=HashMap<String, String>)]
709    pub updated_attributes: HashMap<String, Bytes>,
710    pub deleted_attributes: HashSet<String>,
711}
712
713impl From<models::protocol::ProtocolComponentStateDelta> for ProtocolStateDelta {
714    fn from(value: models::protocol::ProtocolComponentStateDelta) -> Self {
715        Self {
716            component_id: value.component_id,
717            updated_attributes: value.updated_attributes,
718            deleted_attributes: value.deleted_attributes,
719        }
720    }
721}
722
723impl ProtocolStateDelta {
724    /// Merges 'other' into 'self'.
725    ///
726    ///
727    /// During merge of these deltas a special situation can arise when an attribute is present in
728    /// `self.deleted_attributes` and `other.update_attributes``. If we would just merge the sets
729    /// of deleted attributes or vice versa, it would be ambiguous and potential lead to a
730    /// deletion of an attribute that should actually be present, or retention of an actually
731    /// deleted attribute.
732    ///
733    /// This situation is handled the following way:
734    ///
735    ///     - If an attribute is deleted and in the next message recreated, it is removed from the
736    ///       set of deleted attributes and kept in updated_attributes. This way it's temporary
737    ///       deletion is never communicated to the final receiver.
738    ///     - If an attribute was updated and is deleted in the next message, it is removed from
739    ///       updated attributes and kept in deleted. This way the attributes temporary update (or
740    ///       potentially short-lived existence) before its deletion is never communicated to the
741    ///       final receiver.
742    pub fn merge(&mut self, other: &Self) {
743        // either updated and then deleted -> keep in deleted, remove from updated
744        self.updated_attributes
745            .retain(|k, _| !other.deleted_attributes.contains(k));
746
747        // or deleted and then updated/recreated -> remove from deleted and keep in updated
748        self.deleted_attributes.retain(|attr| {
749            !other
750                .updated_attributes
751                .contains_key(attr)
752        });
753
754        // simply merge updates
755        self.updated_attributes.extend(
756            other
757                .updated_attributes
758                .iter()
759                .map(|(k, v)| (k.clone(), v.clone())),
760        );
761
762        // simply merge deletions
763        self.deleted_attributes
764            .extend(other.deleted_attributes.iter().cloned());
765    }
766}
767
768/// Maximum page size for this endpoint is 100
769#[derive(
770    Clone, Serialize, Debug, Default, Deserialize, PartialEq, ToSchema, Eq, Hash, DeepSizeOf,
771)]
772#[serde(deny_unknown_fields)]
773pub struct StateRequestBody {
774    /// Filters response by contract addresses
775    #[serde(alias = "contractIds")]
776    #[schema(value_type=Option<Vec<String>>)]
777    pub contract_ids: Option<Vec<Bytes>>,
778    /// Does not filter response, only required to correctly apply unconfirmed state
779    /// from ReorgBuffers
780    #[serde(alias = "protocolSystem", default)]
781    pub protocol_system: String,
782    #[serde(default = "VersionParam::default")]
783    pub version: VersionParam,
784    #[serde(default)]
785    pub chain: Chain,
786    #[serde(default)]
787    pub pagination: PaginationParams,
788}
789
790impl StateRequestBody {
791    pub fn new(
792        contract_ids: Option<Vec<Bytes>>,
793        protocol_system: String,
794        version: VersionParam,
795        chain: Chain,
796        pagination: PaginationParams,
797    ) -> Self {
798        Self { contract_ids, protocol_system, version, chain, pagination }
799    }
800
801    pub fn from_block(protocol_system: &str, block: BlockParam) -> Self {
802        Self {
803            contract_ids: None,
804            protocol_system: protocol_system.to_string(),
805            version: VersionParam { timestamp: None, block: Some(block.clone()) },
806            chain: block.chain.unwrap_or_default(),
807            pagination: PaginationParams::default(),
808        }
809    }
810
811    pub fn from_timestamp(protocol_system: &str, timestamp: NaiveDateTime, chain: Chain) -> Self {
812        Self {
813            contract_ids: None,
814            protocol_system: protocol_system.to_string(),
815            version: VersionParam { timestamp: Some(timestamp), block: None },
816            chain,
817            pagination: PaginationParams::default(),
818        }
819    }
820}
821
822/// Response from Tycho server for a contract state request.
823#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, ToSchema, DeepSizeOf)]
824pub struct StateRequestResponse {
825    pub accounts: Vec<ResponseAccount>,
826    pub pagination: PaginationResponse,
827}
828
829impl StateRequestResponse {
830    pub fn new(accounts: Vec<ResponseAccount>, pagination: PaginationResponse) -> Self {
831        Self { accounts, pagination }
832    }
833}
834
835#[derive(PartialEq, Clone, Serialize, Deserialize, Default, ToSchema, DeepSizeOf)]
836#[serde(rename = "Account")]
837/// Account struct for the response from Tycho server for a contract state request.
838///
839/// Code is serialized as a hex string instead of a list of bytes.
840pub struct ResponseAccount {
841    pub chain: Chain,
842    /// The address of the account as hex encoded string
843    #[schema(value_type=String, example="0xc9f2e6ea1637E499406986ac50ddC92401ce1f58")]
844    #[serde(with = "hex_bytes")]
845    pub address: Bytes,
846    /// The title of the account usualy specifying its function within the protocol
847    #[schema(value_type=String, example="Protocol Vault")]
848    pub title: String,
849    /// Contract storage map of hex encoded string values
850    #[schema(value_type=HashMap<String, String>, example=json!({"0x....": "0x...."}))]
851    #[serde(with = "hex_hashmap_key_value")]
852    pub slots: HashMap<Bytes, Bytes>,
853    /// The balance of the account in the native token
854    #[schema(value_type=String, example="0x00")]
855    #[serde(with = "hex_bytes")]
856    pub native_balance: Bytes,
857    /// Balances of this account in other tokens (only tokens balance that are
858    /// relevant to the protocol are returned here)
859    #[schema(value_type=HashMap<String, String>, example=json!({"0x....": "0x...."}))]
860    #[serde(with = "hex_hashmap_key_value")]
861    pub token_balances: HashMap<Bytes, Bytes>,
862    /// The accounts code as hex encoded string
863    #[schema(value_type=String, example="0xBADBABE")]
864    #[serde(with = "hex_bytes")]
865    pub code: Bytes,
866    /// The hash of above code
867    #[schema(value_type=String, example="0x123456789")]
868    #[serde(with = "hex_bytes")]
869    pub code_hash: Bytes,
870    /// Transaction hash which last modified native balance
871    #[schema(value_type=String, example="0x8f1133bfb054a23aedfe5d25b1d81b96195396d8b88bd5d4bcf865fc1ae2c3f4")]
872    #[serde(with = "hex_bytes")]
873    pub balance_modify_tx: Bytes,
874    /// Transaction hash which last modified code
875    #[schema(value_type=String, example="0x8f1133bfb054a23aedfe5d25b1d81b96195396d8b88bd5d4bcf865fc1ae2c3f4")]
876    #[serde(with = "hex_bytes")]
877    pub code_modify_tx: Bytes,
878    /// Transaction hash which created the account
879    #[deprecated(note = "The `creation_tx` field is deprecated.")]
880    #[schema(value_type=Option<String>, example="0x8f1133bfb054a23aedfe5d25b1d81b96195396d8b88bd5d4bcf865fc1ae2c3f4")]
881    #[serde(with = "hex_bytes_option")]
882    pub creation_tx: Option<Bytes>,
883}
884
885impl ResponseAccount {
886    #[allow(clippy::too_many_arguments)]
887    pub fn new(
888        chain: Chain,
889        address: Bytes,
890        title: String,
891        slots: HashMap<Bytes, Bytes>,
892        native_balance: Bytes,
893        token_balances: HashMap<Bytes, Bytes>,
894        code: Bytes,
895        code_hash: Bytes,
896        balance_modify_tx: Bytes,
897        code_modify_tx: Bytes,
898        creation_tx: Option<Bytes>,
899    ) -> Self {
900        Self {
901            chain,
902            address,
903            title,
904            slots,
905            native_balance,
906            token_balances,
907            code,
908            code_hash,
909            balance_modify_tx,
910            code_modify_tx,
911            creation_tx,
912        }
913    }
914}
915
916/// Implement Debug for ResponseAccount manually to avoid printing the code field.
917impl fmt::Debug for ResponseAccount {
918    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
919        f.debug_struct("ResponseAccount")
920            .field("chain", &self.chain)
921            .field("address", &self.address)
922            .field("title", &self.title)
923            .field("slots", &self.slots)
924            .field("native_balance", &self.native_balance)
925            .field("token_balances", &self.token_balances)
926            .field("code", &format!("[{} bytes]", self.code.len()))
927            .field("code_hash", &self.code_hash)
928            .field("balance_modify_tx", &self.balance_modify_tx)
929            .field("code_modify_tx", &self.code_modify_tx)
930            .field("creation_tx", &self.creation_tx)
931            .finish()
932    }
933}
934
935#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, Default)]
936pub struct AccountBalance {
937    #[serde(with = "hex_bytes")]
938    pub account: Bytes,
939    #[serde(with = "hex_bytes")]
940    pub token: Bytes,
941    #[serde(with = "hex_bytes")]
942    pub balance: Bytes,
943    #[serde(with = "hex_bytes")]
944    pub modify_tx: Bytes,
945}
946
947#[derive(Debug, Serialize, Deserialize, PartialEq, Clone, ToSchema)]
948#[serde(deny_unknown_fields)]
949pub struct ContractId {
950    #[serde(with = "hex_bytes")]
951    #[schema(value_type=String)]
952    pub address: Bytes,
953    pub chain: Chain,
954}
955
956/// Uniquely identifies a contract on a specific chain.
957impl ContractId {
958    pub fn new(chain: Chain, address: Bytes) -> Self {
959        Self { address, chain }
960    }
961
962    pub fn address(&self) -> &Bytes {
963        &self.address
964    }
965}
966
967impl fmt::Display for ContractId {
968    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
969        write!(f, "{:?}: 0x{}", self.chain, hex::encode(&self.address))
970    }
971}
972
973/// The version of the requested state, given as either a timestamp or a block.
974///
975/// If block is provided, the state at that exact block is returned. Will error if the block
976/// has not been processed yet. If timestamp is provided, the state at the latest block before
977/// that timestamp is returned.
978/// Defaults to the current time.
979#[derive(Debug, Serialize, Deserialize, PartialEq, Clone, ToSchema, Eq, Hash)]
980#[serde(deny_unknown_fields)]
981pub struct VersionParam {
982    pub timestamp: Option<NaiveDateTime>,
983    pub block: Option<BlockParam>,
984}
985
986impl DeepSizeOf for VersionParam {
987    fn deep_size_of_children(&self, ctx: &mut Context) -> usize {
988        if let Some(block) = &self.block {
989            return block.deep_size_of_children(ctx);
990        }
991
992        0
993    }
994}
995
996impl VersionParam {
997    pub fn new(timestamp: Option<NaiveDateTime>, block: Option<BlockParam>) -> Self {
998        Self { timestamp, block }
999    }
1000}
1001
1002impl Default for VersionParam {
1003    fn default() -> Self {
1004        VersionParam { timestamp: Some(Utc::now().naive_utc()), block: None }
1005    }
1006}
1007
1008#[deprecated(note = "Use StateRequestBody instead")]
1009#[derive(Serialize, Deserialize, Default, Debug, IntoParams)]
1010pub struct StateRequestParameters {
1011    /// The minimum TVL of the protocol components to return, denoted in the chain's native token.
1012    #[param(default = 0)]
1013    pub tvl_gt: Option<u64>,
1014    /// The minimum inertia of the protocol components to return.
1015    #[param(default = 0)]
1016    pub inertia_min_gt: Option<u64>,
1017    /// Whether to include ERC20 balances in the response.
1018    #[serde(default = "default_include_balances_flag")]
1019    pub include_balances: bool,
1020    #[serde(default)]
1021    pub pagination: PaginationParams,
1022}
1023
1024impl StateRequestParameters {
1025    pub fn new(include_balances: bool) -> Self {
1026        Self {
1027            tvl_gt: None,
1028            inertia_min_gt: None,
1029            include_balances,
1030            pagination: PaginationParams::default(),
1031        }
1032    }
1033
1034    pub fn to_query_string(&self) -> String {
1035        let mut parts = vec![format!("include_balances={}", self.include_balances)];
1036
1037        if let Some(tvl_gt) = self.tvl_gt {
1038            parts.push(format!("tvl_gt={tvl_gt}"));
1039        }
1040
1041        if let Some(inertia) = self.inertia_min_gt {
1042            parts.push(format!("inertia_min_gt={inertia}"));
1043        }
1044
1045        let mut res = parts.join("&");
1046        if !res.is_empty() {
1047            res = format!("?{res}");
1048        }
1049        res
1050    }
1051}
1052
1053#[derive(
1054    Serialize, Deserialize, Debug, Default, PartialEq, ToSchema, Eq, Hash, Clone, DeepSizeOf,
1055)]
1056#[serde(deny_unknown_fields)]
1057pub struct TokensRequestBody {
1058    /// Filters tokens by addresses
1059    #[serde(alias = "tokenAddresses")]
1060    #[schema(value_type=Option<Vec<String>>)]
1061    pub token_addresses: Option<Vec<Bytes>>,
1062    /// Quality is between 0-100, where:
1063    ///  - 100: Normal ERC-20 Token behavior
1064    ///  - 75: Rebasing token
1065    ///  - 50: Fee-on-transfer token
1066    ///  - 10: Token analysis failed at first detection
1067    ///  - 5: Token analysis failed multiple times (after creation)
1068    ///  - 0: Failed to extract attributes, like Decimal or Symbol
1069    #[serde(default)]
1070    pub min_quality: Option<i32>,
1071    /// Filters tokens by recent trade activity
1072    #[serde(default)]
1073    pub traded_n_days_ago: Option<u64>,
1074    /// Max page size supported is 3000
1075    #[serde(default)]
1076    pub pagination: PaginationParams,
1077    /// Filter tokens by blockchain, default 'ethereum'
1078    #[serde(default)]
1079    pub chain: Chain,
1080}
1081
1082/// Response from Tycho server for a tokens request.
1083#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, ToSchema, Eq, Hash, DeepSizeOf)]
1084pub struct TokensRequestResponse {
1085    pub tokens: Vec<ResponseToken>,
1086    pub pagination: PaginationResponse,
1087}
1088
1089impl TokensRequestResponse {
1090    pub fn new(tokens: Vec<ResponseToken>, pagination_request: &PaginationResponse) -> Self {
1091        Self { tokens, pagination: pagination_request.clone() }
1092    }
1093}
1094
1095/// Pagination parameter
1096#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, ToSchema, Eq, Hash, DeepSizeOf)]
1097#[serde(deny_unknown_fields)]
1098pub struct PaginationParams {
1099    /// What page to retrieve
1100    #[serde(default)]
1101    pub page: i64,
1102    /// How many results to return per page
1103    #[serde(default)]
1104    #[schema(default = 10)]
1105    pub page_size: i64,
1106}
1107
1108impl PaginationParams {
1109    pub fn new(page: i64, page_size: i64) -> Self {
1110        Self { page, page_size }
1111    }
1112}
1113
1114impl Default for PaginationParams {
1115    fn default() -> Self {
1116        PaginationParams { page: 0, page_size: 20 }
1117    }
1118}
1119
1120#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, ToSchema, Eq, Hash, DeepSizeOf)]
1121#[serde(deny_unknown_fields)]
1122pub struct PaginationResponse {
1123    pub page: i64,
1124    pub page_size: i64,
1125    /// The total number of items available across all pages of results
1126    pub total: i64,
1127}
1128
1129/// Current pagination information
1130impl PaginationResponse {
1131    pub fn new(page: i64, page_size: i64, total: i64) -> Self {
1132        Self { page, page_size, total }
1133    }
1134
1135    pub fn total_pages(&self) -> i64 {
1136        // ceil(total / page_size)
1137        (self.total + self.page_size - 1) / self.page_size
1138    }
1139}
1140
1141#[derive(
1142    PartialEq, Debug, Clone, Serialize, Deserialize, Default, ToSchema, Eq, Hash, DeepSizeOf,
1143)]
1144#[serde(rename = "Token")]
1145/// Token struct for the response from Tycho server for a tokens request.
1146pub struct ResponseToken {
1147    pub chain: Chain,
1148    /// The address of this token as hex encoded string
1149    #[schema(value_type=String, example="0xc9f2e6ea1637E499406986ac50ddC92401ce1f58")]
1150    #[serde(with = "hex_bytes")]
1151    pub address: Bytes,
1152    /// A shorthand symbol for this token (not unique)
1153    #[schema(value_type=String, example="WETH")]
1154    pub symbol: String,
1155    /// The number of decimals used to represent token values
1156    pub decimals: u32,
1157    /// The tax this token charges on transfers in basis points
1158    pub tax: u64,
1159    /// Gas usage of the token, currently is always a single averaged value
1160    pub gas: Vec<Option<u64>>,
1161    /// Quality is between 0-100, where:
1162    ///  - 100: Normal ERC-20 Token behavior
1163    ///  - 75: Rebasing token
1164    ///  - 50: Fee-on-transfer token
1165    ///  - 10: Token analysis failed at first detection
1166    ///  - 5: Token analysis failed multiple times (after creation)
1167    ///  - 0: Failed to extract attributes, like Decimal or Symbol
1168    pub quality: u32,
1169}
1170
1171impl From<models::token::Token> for ResponseToken {
1172    fn from(value: models::token::Token) -> Self {
1173        Self {
1174            chain: value.chain.into(),
1175            address: value.address,
1176            symbol: value.symbol,
1177            decimals: value.decimals,
1178            tax: value.tax,
1179            gas: value.gas,
1180            quality: value.quality,
1181        }
1182    }
1183}
1184
1185#[derive(Serialize, Deserialize, Debug, Default, ToSchema, Clone, DeepSizeOf)]
1186#[serde(deny_unknown_fields)]
1187pub struct ProtocolComponentsRequestBody {
1188    /// Filters by protocol, required to correctly apply unconfirmed state from
1189    /// ReorgBuffers
1190    pub protocol_system: String,
1191    /// Filter by component ids
1192    #[serde(alias = "componentAddresses")]
1193    pub component_ids: Option<Vec<ComponentId>>,
1194    /// The minimum TVL of the protocol components to return, denoted in the chain's
1195    /// native token.
1196    #[serde(default)]
1197    pub tvl_gt: Option<f64>,
1198    #[serde(default)]
1199    pub chain: Chain,
1200    /// Max page size supported is 500
1201    #[serde(default)]
1202    pub pagination: PaginationParams,
1203}
1204
1205// Implement PartialEq where tvl is considered equal if the difference is less than 1e-6
1206impl PartialEq for ProtocolComponentsRequestBody {
1207    fn eq(&self, other: &Self) -> bool {
1208        let tvl_close_enough = match (self.tvl_gt, other.tvl_gt) {
1209            (Some(a), Some(b)) => (a - b).abs() < 1e-6,
1210            (None, None) => true,
1211            _ => false,
1212        };
1213
1214        self.protocol_system == other.protocol_system &&
1215            self.component_ids == other.component_ids &&
1216            tvl_close_enough &&
1217            self.chain == other.chain &&
1218            self.pagination == other.pagination
1219    }
1220}
1221
1222// Implement Eq without any new logic
1223impl Eq for ProtocolComponentsRequestBody {}
1224
1225impl Hash for ProtocolComponentsRequestBody {
1226    fn hash<H: Hasher>(&self, state: &mut H) {
1227        self.protocol_system.hash(state);
1228        self.component_ids.hash(state);
1229
1230        // Handle the f64 `tvl_gt` field by converting it into a hashable integer
1231        if let Some(tvl) = self.tvl_gt {
1232            // Convert f64 to bits and hash those bits
1233            tvl.to_bits().hash(state);
1234        } else {
1235            // Use a constant value to represent None
1236            state.write_u8(0);
1237        }
1238
1239        self.chain.hash(state);
1240        self.pagination.hash(state);
1241    }
1242}
1243
1244impl ProtocolComponentsRequestBody {
1245    pub fn system_filtered(system: &str, tvl_gt: Option<f64>, chain: Chain) -> Self {
1246        Self {
1247            protocol_system: system.to_string(),
1248            component_ids: None,
1249            tvl_gt,
1250            chain,
1251            pagination: Default::default(),
1252        }
1253    }
1254
1255    pub fn id_filtered(system: &str, ids: Vec<String>, chain: Chain) -> Self {
1256        Self {
1257            protocol_system: system.to_string(),
1258            component_ids: Some(ids),
1259            tvl_gt: None,
1260            chain,
1261            pagination: Default::default(),
1262        }
1263    }
1264}
1265
1266impl ProtocolComponentsRequestBody {
1267    pub fn new(
1268        protocol_system: String,
1269        component_ids: Option<Vec<String>>,
1270        tvl_gt: Option<f64>,
1271        chain: Chain,
1272        pagination: PaginationParams,
1273    ) -> Self {
1274        Self { protocol_system, component_ids, tvl_gt, chain, pagination }
1275    }
1276}
1277
1278#[deprecated(note = "Use ProtocolComponentsRequestBody instead")]
1279#[derive(Serialize, Deserialize, Default, Debug, IntoParams)]
1280pub struct ProtocolComponentRequestParameters {
1281    /// The minimum TVL of the protocol components to return, denoted in the chain's native token.
1282    #[param(default = 0)]
1283    pub tvl_gt: Option<f64>,
1284}
1285
1286impl ProtocolComponentRequestParameters {
1287    pub fn tvl_filtered(min_tvl: f64) -> Self {
1288        Self { tvl_gt: Some(min_tvl) }
1289    }
1290}
1291
1292impl ProtocolComponentRequestParameters {
1293    pub fn to_query_string(&self) -> String {
1294        if let Some(tvl_gt) = self.tvl_gt {
1295            return format!("?tvl_gt={tvl_gt}");
1296        }
1297        String::new()
1298    }
1299}
1300
1301/// Response from Tycho server for a protocol components request.
1302#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, ToSchema, DeepSizeOf)]
1303pub struct ProtocolComponentRequestResponse {
1304    pub protocol_components: Vec<ProtocolComponent>,
1305    pub pagination: PaginationResponse,
1306}
1307
1308impl ProtocolComponentRequestResponse {
1309    pub fn new(
1310        protocol_components: Vec<ProtocolComponent>,
1311        pagination: PaginationResponse,
1312    ) -> Self {
1313        Self { protocol_components, pagination }
1314    }
1315}
1316
1317#[derive(Debug, Serialize, Deserialize, PartialEq, Clone, ToSchema, Eq, Hash)]
1318#[serde(deny_unknown_fields)]
1319#[deprecated]
1320pub struct ProtocolId {
1321    pub id: String,
1322    pub chain: Chain,
1323}
1324
1325impl From<ProtocolId> for String {
1326    fn from(protocol_id: ProtocolId) -> Self {
1327        protocol_id.id
1328    }
1329}
1330
1331impl AsRef<str> for ProtocolId {
1332    fn as_ref(&self) -> &str {
1333        &self.id
1334    }
1335}
1336
1337/// Protocol State struct for the response from Tycho server for a protocol state request.
1338#[derive(Debug, Clone, PartialEq, Default, Deserialize, Serialize, ToSchema, DeepSizeOf)]
1339pub struct ResponseProtocolState {
1340    /// Component id this state belongs to
1341    pub component_id: String,
1342    /// Attributes of the component. If an attribute's value is a `bigint`,
1343    /// it will be encoded as a big endian signed hex string.
1344    #[schema(value_type=HashMap<String, String>)]
1345    #[serde(with = "hex_hashmap_value")]
1346    pub attributes: HashMap<String, Bytes>,
1347    /// Sum aggregated balances of the component
1348    #[schema(value_type=HashMap<String, String>)]
1349    #[serde(with = "hex_hashmap_key_value")]
1350    pub balances: HashMap<Bytes, Bytes>,
1351}
1352
1353impl From<models::protocol::ProtocolComponentState> for ResponseProtocolState {
1354    fn from(value: models::protocol::ProtocolComponentState) -> Self {
1355        Self {
1356            component_id: value.component_id,
1357            attributes: value.attributes,
1358            balances: value.balances,
1359        }
1360    }
1361}
1362
1363fn default_include_balances_flag() -> bool {
1364    true
1365}
1366
1367/// Max page size supported is 100
1368#[derive(Clone, Debug, Serialize, PartialEq, ToSchema, Default, Eq, Hash, DeepSizeOf)]
1369#[serde(deny_unknown_fields)]
1370pub struct ProtocolStateRequestBody {
1371    /// Filters response by protocol components ids
1372    #[serde(alias = "protocolIds")]
1373    pub protocol_ids: Option<Vec<String>>,
1374    /// Filters by protocol, required to correctly apply unconfirmed state from
1375    /// ReorgBuffers
1376    #[serde(alias = "protocolSystem")]
1377    pub protocol_system: String,
1378    #[serde(default)]
1379    pub chain: Chain,
1380    /// Whether to include account balances in the response. Defaults to true.
1381    #[serde(default = "default_include_balances_flag")]
1382    pub include_balances: bool,
1383    #[serde(default = "VersionParam::default")]
1384    pub version: VersionParam,
1385    #[serde(default)]
1386    pub pagination: PaginationParams,
1387}
1388
1389impl ProtocolStateRequestBody {
1390    pub fn id_filtered<I, T>(ids: I) -> Self
1391    where
1392        I: IntoIterator<Item = T>,
1393        T: Into<String>,
1394    {
1395        Self {
1396            protocol_ids: Some(
1397                ids.into_iter()
1398                    .map(Into::into)
1399                    .collect(),
1400            ),
1401            ..Default::default()
1402        }
1403    }
1404}
1405
1406/// Custom deserializer for ProtocolStateRequestBody to support backwards compatibility with the old
1407/// ProtocolIds format.
1408/// To be removed when the old format is no longer supported.
1409impl<'de> Deserialize<'de> for ProtocolStateRequestBody {
1410    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
1411    where
1412        D: Deserializer<'de>,
1413    {
1414        #[derive(Deserialize)]
1415        #[serde(untagged)]
1416        enum ProtocolIdOrString {
1417            Old(Vec<ProtocolId>),
1418            New(Vec<String>),
1419        }
1420
1421        struct ProtocolStateRequestBodyVisitor;
1422
1423        impl<'de> de::Visitor<'de> for ProtocolStateRequestBodyVisitor {
1424            type Value = ProtocolStateRequestBody;
1425
1426            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
1427                formatter.write_str("struct ProtocolStateRequestBody")
1428            }
1429
1430            fn visit_map<V>(self, mut map: V) -> Result<ProtocolStateRequestBody, V::Error>
1431            where
1432                V: de::MapAccess<'de>,
1433            {
1434                let mut protocol_ids = None;
1435                let mut protocol_system = None;
1436                let mut version = None;
1437                let mut chain = None;
1438                let mut include_balances = None;
1439                let mut pagination = None;
1440
1441                while let Some(key) = map.next_key::<String>()? {
1442                    match key.as_str() {
1443                        "protocol_ids" | "protocolIds" => {
1444                            let value: ProtocolIdOrString = map.next_value()?;
1445                            protocol_ids = match value {
1446                                ProtocolIdOrString::Old(ids) => {
1447                                    Some(ids.into_iter().map(|p| p.id).collect())
1448                                }
1449                                ProtocolIdOrString::New(ids_str) => Some(ids_str),
1450                            };
1451                        }
1452                        "protocol_system" | "protocolSystem" => {
1453                            protocol_system = Some(map.next_value()?);
1454                        }
1455                        "version" => {
1456                            version = Some(map.next_value()?);
1457                        }
1458                        "chain" => {
1459                            chain = Some(map.next_value()?);
1460                        }
1461                        "include_balances" => {
1462                            include_balances = Some(map.next_value()?);
1463                        }
1464                        "pagination" => {
1465                            pagination = Some(map.next_value()?);
1466                        }
1467                        _ => {
1468                            return Err(de::Error::unknown_field(
1469                                &key,
1470                                &[
1471                                    "contract_ids",
1472                                    "protocol_system",
1473                                    "version",
1474                                    "chain",
1475                                    "include_balances",
1476                                    "pagination",
1477                                ],
1478                            ))
1479                        }
1480                    }
1481                }
1482
1483                Ok(ProtocolStateRequestBody {
1484                    protocol_ids,
1485                    protocol_system: protocol_system.unwrap_or_default(),
1486                    version: version.unwrap_or_else(VersionParam::default),
1487                    chain: chain.unwrap_or_else(Chain::default),
1488                    include_balances: include_balances.unwrap_or(true),
1489                    pagination: pagination.unwrap_or_else(PaginationParams::default),
1490                })
1491            }
1492        }
1493
1494        deserializer.deserialize_struct(
1495            "ProtocolStateRequestBody",
1496            &[
1497                "contract_ids",
1498                "protocol_system",
1499                "version",
1500                "chain",
1501                "include_balances",
1502                "pagination",
1503            ],
1504            ProtocolStateRequestBodyVisitor,
1505        )
1506    }
1507}
1508
1509#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, ToSchema, DeepSizeOf)]
1510pub struct ProtocolStateRequestResponse {
1511    pub states: Vec<ResponseProtocolState>,
1512    pub pagination: PaginationResponse,
1513}
1514
1515impl ProtocolStateRequestResponse {
1516    pub fn new(states: Vec<ResponseProtocolState>, pagination: PaginationResponse) -> Self {
1517        Self { states, pagination }
1518    }
1519}
1520
1521#[derive(Serialize, Clone, PartialEq, Hash, Eq)]
1522pub struct ProtocolComponentId {
1523    pub chain: Chain,
1524    pub system: String,
1525    pub id: String,
1526}
1527
1528#[derive(Debug, Serialize, ToSchema)]
1529#[serde(tag = "status", content = "message")]
1530#[schema(example = json!({"status": "NotReady", "message": "No db connection"}))]
1531pub enum Health {
1532    Ready,
1533    Starting(String),
1534    NotReady(String),
1535}
1536
1537#[derive(Serialize, Deserialize, Debug, Default, PartialEq, ToSchema, Eq, Hash, Clone)]
1538#[serde(deny_unknown_fields)]
1539pub struct ProtocolSystemsRequestBody {
1540    #[serde(default)]
1541    pub chain: Chain,
1542    #[serde(default)]
1543    pub pagination: PaginationParams,
1544}
1545
1546#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, ToSchema, Eq, Hash)]
1547pub struct ProtocolSystemsRequestResponse {
1548    /// List of currently supported protocol systems
1549    pub protocol_systems: Vec<String>,
1550    pub pagination: PaginationResponse,
1551}
1552
1553impl ProtocolSystemsRequestResponse {
1554    pub fn new(protocol_systems: Vec<String>, pagination: PaginationResponse) -> Self {
1555        Self { protocol_systems, pagination }
1556    }
1557}
1558
1559#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, Default)]
1560pub struct DCIUpdate {
1561    /// Map of component id to the new entrypoints associated with the component
1562    pub new_entrypoints: HashMap<ComponentId, HashSet<EntryPoint>>,
1563    /// Map of entrypoint id to the new entrypoint params associtated with it (and optionally the
1564    /// component linked to those params)
1565    pub new_entrypoint_params: HashMap<String, HashSet<(TracingParams, Option<String>)>>,
1566    /// Map of entrypoint id to its trace result
1567    pub trace_results: HashMap<String, TracingResult>,
1568}
1569
1570impl From<models::blockchain::DCIUpdate> for DCIUpdate {
1571    fn from(value: models::blockchain::DCIUpdate) -> Self {
1572        Self {
1573            new_entrypoints: value
1574                .new_entrypoints
1575                .into_iter()
1576                .map(|(k, v)| {
1577                    (
1578                        k,
1579                        v.into_iter()
1580                            .map(|v| v.into())
1581                            .collect(),
1582                    )
1583                })
1584                .collect(),
1585            new_entrypoint_params: value
1586                .new_entrypoint_params
1587                .into_iter()
1588                .map(|(k, v)| {
1589                    (
1590                        k,
1591                        v.into_iter()
1592                            .map(|(params, i)| (params.into(), i))
1593                            .collect(),
1594                    )
1595                })
1596                .collect(),
1597            trace_results: value
1598                .trace_results
1599                .into_iter()
1600                .map(|(k, v)| (k, v.into()))
1601                .collect(),
1602        }
1603    }
1604}
1605
1606#[derive(Serialize, Deserialize, Debug, Default, PartialEq, ToSchema, Eq, Hash, Clone)]
1607#[serde(deny_unknown_fields)]
1608pub struct ComponentTvlRequestBody {
1609    #[serde(default)]
1610    pub chain: Chain,
1611    /// Filters protocol components by protocol system
1612    /// Useful when `component_ids` is omitted to fetch all components under a specific system.
1613    #[serde(alias = "protocolSystem")]
1614    pub protocol_system: Option<String>,
1615    #[serde(default)]
1616    pub component_ids: Option<Vec<String>>,
1617    #[serde(default)]
1618    pub pagination: PaginationParams,
1619}
1620
1621impl ComponentTvlRequestBody {
1622    pub fn system_filtered(system: &str, chain: Chain) -> Self {
1623        Self {
1624            chain,
1625            protocol_system: Some(system.to_string()),
1626            component_ids: None,
1627            pagination: Default::default(),
1628        }
1629    }
1630
1631    pub fn id_filtered(ids: Vec<String>, chain: Chain) -> Self {
1632        Self {
1633            chain,
1634            protocol_system: None,
1635            component_ids: Some(ids),
1636            pagination: Default::default(),
1637        }
1638    }
1639}
1640// #[derive(Clone, Debug, Serialize, Deserialize, PartialEq, ToSchema, Eq, Hash)]
1641#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, ToSchema)]
1642pub struct ComponentTvlRequestResponse {
1643    pub tvl: HashMap<String, f64>,
1644    pub pagination: PaginationResponse,
1645}
1646
1647impl ComponentTvlRequestResponse {
1648    pub fn new(tvl: HashMap<String, f64>, pagination: PaginationResponse) -> Self {
1649        Self { tvl, pagination }
1650    }
1651}
1652
1653#[derive(
1654    Serialize, Deserialize, Debug, Default, PartialEq, ToSchema, Eq, Hash, Clone, DeepSizeOf,
1655)]
1656pub struct TracedEntryPointRequestBody {
1657    #[serde(default)]
1658    pub chain: Chain,
1659    /// Filters by protocol, required to correctly apply unconfirmed state from
1660    /// ReorgBuffers
1661    pub protocol_system: String,
1662    /// Filter by component ids
1663    pub component_ids: Option<Vec<ComponentId>>,
1664    /// Max page size supported is 100
1665    #[serde(default)]
1666    pub pagination: PaginationParams,
1667}
1668
1669#[derive(Debug, Serialize, Deserialize, PartialEq, Clone, ToSchema, Eq, Hash, DeepSizeOf)]
1670pub struct EntryPoint {
1671    #[schema(example = "0xEdf63cce4bA70cbE74064b7687882E71ebB0e988:getRate()")]
1672    /// Entry point id.
1673    pub external_id: String,
1674    #[schema(value_type=String, example="0x8f4E8439b970363648421C692dd897Fb9c0Bd1D9")]
1675    #[serde(with = "hex_bytes")]
1676    /// The address of the contract to trace.
1677    pub target: Bytes,
1678    #[schema(example = "getRate()")]
1679    /// The signature of the function to trace.
1680    pub signature: String,
1681}
1682
1683#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, ToSchema, Eq, Hash, DeepSizeOf)]
1684pub enum StorageOverride {
1685    /// Applies changes incrementally to the existing account storage.
1686    /// Only modifies the specific storage slots provided in the map while
1687    /// preserving all other storage slots.
1688    Diff(BTreeMap<StoreKey, StoreVal>),
1689
1690    /// Completely replaces the account's storage state.
1691    /// Only the storage slots provided in the map will exist after the operation,
1692    /// and any existing storage slots not included will be cleared/zeroed.
1693    Replace(BTreeMap<StoreKey, StoreVal>),
1694}
1695
1696impl From<models::blockchain::StorageOverride> for StorageOverride {
1697    fn from(value: models::blockchain::StorageOverride) -> Self {
1698        match value {
1699            models::blockchain::StorageOverride::Diff(diff) => StorageOverride::Diff(diff),
1700            models::blockchain::StorageOverride::Replace(replace) => {
1701                StorageOverride::Replace(replace)
1702            }
1703        }
1704    }
1705}
1706
1707/// State overrides for an account.
1708///
1709/// Used to modify account state. Commonly used for testing contract interactions with specific
1710/// state conditions or simulating transactions with modified balances/code.
1711#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, ToSchema, Eq, Hash, DeepSizeOf)]
1712pub struct AccountOverrides {
1713    /// Storage slots to override
1714    pub slots: Option<StorageOverride>,
1715    /// Native token balance override
1716    pub native_balance: Option<Balance>,
1717    /// Contract code override
1718    pub code: Option<Code>,
1719}
1720
1721impl From<models::blockchain::AccountOverrides> for AccountOverrides {
1722    fn from(value: models::blockchain::AccountOverrides) -> Self {
1723        AccountOverrides {
1724            slots: value.slots.map(|s| s.into()),
1725            native_balance: value.native_balance,
1726            code: value.code,
1727        }
1728    }
1729}
1730
1731#[derive(Debug, Serialize, Deserialize, PartialEq, Clone, ToSchema, Eq, Hash, DeepSizeOf)]
1732pub struct RPCTracerParams {
1733    /// The caller address of the transaction, if not provided tracing uses the default value
1734    /// for an address defined by the VM.
1735    #[schema(value_type=Option<String>)]
1736    #[serde(with = "hex_bytes_option", default)]
1737    pub caller: Option<Bytes>,
1738    /// The call data used for the tracing call, this needs to include the function selector
1739    #[schema(value_type=String, example="0x679aefce")]
1740    #[serde(with = "hex_bytes")]
1741    pub calldata: Bytes,
1742    /// Optionally allow for state overrides so that the call works as expected
1743    pub state_overrides: Option<BTreeMap<Address, AccountOverrides>>,
1744    /// Addresses to prune from trace results. Useful for hooks that use mock
1745    /// accounts/routers that shouldn't be tracked in the final DCI results.
1746    #[schema(value_type=Option<Vec<String>>)]
1747    #[serde(default)]
1748    pub prune_addresses: Option<Vec<Address>>,
1749}
1750
1751impl From<models::blockchain::RPCTracerParams> for RPCTracerParams {
1752    fn from(value: models::blockchain::RPCTracerParams) -> Self {
1753        RPCTracerParams {
1754            caller: value.caller,
1755            calldata: value.calldata,
1756            state_overrides: value.state_overrides.map(|overrides| {
1757                overrides
1758                    .into_iter()
1759                    .map(|(address, account_overrides)| (address, account_overrides.into()))
1760                    .collect()
1761            }),
1762            prune_addresses: value.prune_addresses,
1763        }
1764    }
1765}
1766
1767#[derive(Deserialize, Serialize, Debug, PartialEq, Eq, Clone, Hash, DeepSizeOf)]
1768#[serde(tag = "method", rename_all = "lowercase")]
1769pub enum TracingParams {
1770    /// Uses RPC calls to retrieve the called addresses and retriggers
1771    RPCTracer(RPCTracerParams),
1772}
1773
1774impl From<models::blockchain::TracingParams> for TracingParams {
1775    fn from(value: models::blockchain::TracingParams) -> Self {
1776        match value {
1777            models::blockchain::TracingParams::RPCTracer(params) => {
1778                TracingParams::RPCTracer(params.into())
1779            }
1780        }
1781    }
1782}
1783
1784impl From<models::blockchain::EntryPoint> for EntryPoint {
1785    fn from(value: models::blockchain::EntryPoint) -> Self {
1786        Self { external_id: value.external_id, target: value.target, signature: value.signature }
1787    }
1788}
1789
1790#[derive(Serialize, Deserialize, Debug, PartialEq, ToSchema, Eq, Clone, DeepSizeOf)]
1791pub struct EntryPointWithTracingParams {
1792    /// The entry point object
1793    pub entry_point: EntryPoint,
1794    /// The parameters used
1795    pub params: TracingParams,
1796}
1797
1798impl From<models::blockchain::EntryPointWithTracingParams> for EntryPointWithTracingParams {
1799    fn from(value: models::blockchain::EntryPointWithTracingParams) -> Self {
1800        Self { entry_point: value.entry_point.into(), params: value.params.into() }
1801    }
1802}
1803
1804#[derive(
1805    Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash, Serialize, Deserialize, DeepSizeOf,
1806)]
1807pub struct AddressStorageLocation {
1808    pub key: StoreKey,
1809    pub offset: u8,
1810}
1811
1812impl AddressStorageLocation {
1813    pub fn new(key: StoreKey, offset: u8) -> Self {
1814        Self { key, offset }
1815    }
1816}
1817
1818impl From<models::blockchain::AddressStorageLocation> for AddressStorageLocation {
1819    fn from(value: models::blockchain::AddressStorageLocation) -> Self {
1820        Self { key: value.key, offset: value.offset }
1821    }
1822}
1823
1824fn deserialize_retriggers_from_value(
1825    value: &serde_json::Value,
1826) -> Result<HashSet<(StoreKey, AddressStorageLocation)>, String> {
1827    use serde::Deserialize;
1828    use serde_json::Value;
1829
1830    let mut result = HashSet::new();
1831
1832    if let Value::Array(items) = value {
1833        for item in items {
1834            if let Value::Array(pair) = item {
1835                if pair.len() == 2 {
1836                    let key = StoreKey::deserialize(&pair[0])
1837                        .map_err(|e| format!("Failed to deserialize key: {}", e))?;
1838
1839                    // Handle both old format (string) and new format (AddressStorageLocation)
1840                    let addr_storage = match &pair[1] {
1841                        Value::String(_) => {
1842                            // Old format: just a string key with offset defaulted to 0
1843                            let storage_key = StoreKey::deserialize(&pair[1]).map_err(|e| {
1844                                format!("Failed to deserialize old format storage key: {}", e)
1845                            })?;
1846                            AddressStorageLocation::new(storage_key, 12)
1847                        }
1848                        Value::Object(_) => {
1849                            // New format: AddressStorageLocation struct
1850                            AddressStorageLocation::deserialize(&pair[1]).map_err(|e| {
1851                                format!("Failed to deserialize AddressStorageLocation: {}", e)
1852                            })?
1853                        }
1854                        _ => return Err("Invalid retrigger format".to_string()),
1855                    };
1856
1857                    result.insert((key, addr_storage));
1858                }
1859            }
1860        }
1861    }
1862
1863    Ok(result)
1864}
1865
1866#[derive(Serialize, Debug, Default, PartialEq, ToSchema, Eq, Clone, DeepSizeOf)]
1867pub struct TracingResult {
1868    #[schema(value_type=HashSet<(String, String)>)]
1869    pub retriggers: HashSet<(StoreKey, AddressStorageLocation)>,
1870    #[schema(value_type=HashMap<String,HashSet<String>>)]
1871    pub accessed_slots: HashMap<Address, HashSet<StoreKey>>,
1872}
1873
1874/// Deserialize TracingResult with backward compatibility for retriggers
1875/// TODO: remove this after offset detection is deployed in production
1876impl<'de> Deserialize<'de> for TracingResult {
1877    fn deserialize<D>(deserializer: D) -> Result<TracingResult, D::Error>
1878    where
1879        D: Deserializer<'de>,
1880    {
1881        use serde::de::Error;
1882        use serde_json::Value;
1883
1884        let value = Value::deserialize(deserializer)?;
1885        let mut result = TracingResult::default();
1886
1887        if let Value::Object(map) = value {
1888            // Deserialize retriggers using our custom deserializer
1889            if let Some(retriggers_value) = map.get("retriggers") {
1890                result.retriggers =
1891                    deserialize_retriggers_from_value(retriggers_value).map_err(|e| {
1892                        D::Error::custom(format!("Failed to deserialize retriggers: {}", e))
1893                    })?;
1894            }
1895
1896            // Deserialize accessed_slots normally
1897            if let Some(accessed_slots_value) = map.get("accessed_slots") {
1898                result.accessed_slots = serde_json::from_value(accessed_slots_value.clone())
1899                    .map_err(|e| {
1900                        D::Error::custom(format!("Failed to deserialize accessed_slots: {}", e))
1901                    })?;
1902            }
1903        }
1904
1905        Ok(result)
1906    }
1907}
1908
1909impl From<models::blockchain::TracingResult> for TracingResult {
1910    fn from(value: models::blockchain::TracingResult) -> Self {
1911        TracingResult {
1912            retriggers: value
1913                .retriggers
1914                .into_iter()
1915                .map(|(k, v)| (k, v.into()))
1916                .collect(),
1917            accessed_slots: value.accessed_slots,
1918        }
1919    }
1920}
1921
1922#[derive(Serialize, PartialEq, ToSchema, Eq, Clone, Debug, Deserialize, DeepSizeOf)]
1923pub struct TracedEntryPointRequestResponse {
1924    /// Map of protocol component id to a list of a tuple containing each entry point with its
1925    /// tracing parameters and its corresponding tracing results.
1926    pub traced_entry_points:
1927        HashMap<ComponentId, Vec<(EntryPointWithTracingParams, TracingResult)>>,
1928    pub pagination: PaginationResponse,
1929}
1930impl From<TracedEntryPointRequestResponse> for DCIUpdate {
1931    fn from(response: TracedEntryPointRequestResponse) -> Self {
1932        let mut new_entrypoints = HashMap::new();
1933        let mut new_entrypoint_params = HashMap::new();
1934        let mut trace_results = HashMap::new();
1935
1936        for (component, traces) in response.traced_entry_points {
1937            let mut entrypoints = HashSet::new();
1938
1939            for (entrypoint, trace) in traces {
1940                let entrypoint_id = entrypoint
1941                    .entry_point
1942                    .external_id
1943                    .clone();
1944
1945                // Collect entrypoints
1946                entrypoints.insert(entrypoint.entry_point.clone());
1947
1948                // Collect entrypoint params
1949                new_entrypoint_params
1950                    .entry(entrypoint_id.clone())
1951                    .or_insert_with(HashSet::new)
1952                    .insert((entrypoint.params, Some(component.clone())));
1953
1954                // Collect trace results
1955                trace_results
1956                    .entry(entrypoint_id)
1957                    .and_modify(|existing_trace: &mut TracingResult| {
1958                        // Merge traces for the same entrypoint
1959                        existing_trace
1960                            .retriggers
1961                            .extend(trace.retriggers.clone());
1962                        for (address, slots) in trace.accessed_slots.clone() {
1963                            existing_trace
1964                                .accessed_slots
1965                                .entry(address)
1966                                .or_default()
1967                                .extend(slots);
1968                        }
1969                    })
1970                    .or_insert(trace);
1971            }
1972
1973            if !entrypoints.is_empty() {
1974                new_entrypoints.insert(component, entrypoints);
1975            }
1976        }
1977
1978        DCIUpdate { new_entrypoints, new_entrypoint_params, trace_results }
1979    }
1980}
1981
1982#[derive(Serialize, Deserialize, Debug, Default, PartialEq, ToSchema, Eq, Clone)]
1983pub struct AddEntryPointRequestBody {
1984    #[serde(default)]
1985    pub chain: Chain,
1986    #[schema(value_type=String)]
1987    #[serde(default)]
1988    pub block_hash: Bytes,
1989    /// The map of component ids to their tracing params to insert
1990    pub entry_points_with_tracing_data: Vec<(ComponentId, Vec<EntryPointWithTracingParams>)>,
1991}
1992
1993#[derive(Serialize, PartialEq, ToSchema, Eq, Clone, Debug, Deserialize)]
1994pub struct AddEntryPointRequestResponse {
1995    /// Map of protocol component id to a list of a tuple containing each entry point with its
1996    /// tracing parameters and its corresponding tracing results.
1997    pub traced_entry_points:
1998        HashMap<ComponentId, Vec<(EntryPointWithTracingParams, TracingResult)>>,
1999}
2000
2001#[cfg(test)]
2002mod test {
2003    use std::str::FromStr;
2004
2005    use maplit::hashmap;
2006    use rstest::rstest;
2007
2008    use super::*;
2009
2010    #[test]
2011    fn test_compression_backward_compatibility() {
2012        // Test old format (without compression field) - should default to false
2013        let json_without_compression = r#"{
2014            "method": "subscribe",
2015            "extractor_id": {
2016                "chain": "ethereum",
2017                "name": "test"
2018            },
2019            "include_state": true
2020        }"#;
2021
2022        let command: Command = serde_json::from_str(json_without_compression)
2023            .expect("Failed to deserialize Subscribe without compression field");
2024
2025        if let Command::Subscribe { compression, .. } = command {
2026            assert_eq!(
2027                compression, false,
2028                "compression should default to false when not specified"
2029            );
2030        } else {
2031            panic!("Expected Subscribe command");
2032        }
2033
2034        // Test new format (with compression field)
2035        let json_with_compression = r#"{
2036            "method": "subscribe",
2037            "extractor_id": {
2038                "chain": "ethereum",
2039                "name": "test"
2040            },
2041            "include_state": true,
2042            "compression": true
2043        }"#;
2044
2045        let command_with_compression: Command = serde_json::from_str(json_with_compression)
2046            .expect("Failed to deserialize Subscribe with compression field");
2047
2048        if let Command::Subscribe { compression, .. } = command_with_compression {
2049            assert_eq!(compression, true, "compression should be true as specified in the JSON");
2050        } else {
2051            panic!("Expected Subscribe command");
2052        }
2053    }
2054
2055    #[test]
2056    fn test_tracing_result_backward_compatibility() {
2057        use serde_json::json;
2058
2059        // Test old format (string storage locations)
2060        let old_format_json = json!({
2061            "retriggers": [
2062                ["0x01", "0x02"],
2063                ["0x03", "0x04"]
2064            ],
2065            "accessed_slots": {
2066                "0x05": ["0x06", "0x07"]
2067            }
2068        });
2069
2070        let result: TracingResult = serde_json::from_value(old_format_json).unwrap();
2071
2072        // Check that retriggers were deserialized correctly with offset 0
2073        assert_eq!(result.retriggers.len(), 2);
2074        let retriggers_vec: Vec<_> = result.retriggers.iter().collect();
2075        assert!(retriggers_vec.iter().any(|(k, v)| {
2076            k == &Bytes::from("0x01") && v.key == Bytes::from("0x02") && v.offset == 12
2077        }));
2078        assert!(retriggers_vec.iter().any(|(k, v)| {
2079            k == &Bytes::from("0x03") && v.key == Bytes::from("0x04") && v.offset == 12
2080        }));
2081
2082        // Test new format (AddressStorageLocation objects)
2083        let new_format_json = json!({
2084            "retriggers": [
2085                ["0x01", {"key": "0x02", "offset": 12}],
2086                ["0x03", {"key": "0x04", "offset": 5}]
2087            ],
2088            "accessed_slots": {
2089                "0x05": ["0x06", "0x07"]
2090            }
2091        });
2092
2093        let result2: TracingResult = serde_json::from_value(new_format_json).unwrap();
2094
2095        // Check that new format retriggers were deserialized correctly with proper offsets
2096        assert_eq!(result2.retriggers.len(), 2);
2097        let retriggers_vec2: Vec<_> = result2.retriggers.iter().collect();
2098        assert!(retriggers_vec2.iter().any(|(k, v)| {
2099            k == &Bytes::from("0x01") && v.key == Bytes::from("0x02") && v.offset == 12
2100        }));
2101        assert!(retriggers_vec2.iter().any(|(k, v)| {
2102            k == &Bytes::from("0x03") && v.key == Bytes::from("0x04") && v.offset == 5
2103        }));
2104    }
2105
2106    #[test]
2107    fn test_protocol_components_equality() {
2108        let body1 = ProtocolComponentsRequestBody {
2109            protocol_system: "protocol1".to_string(),
2110            component_ids: Some(vec!["component1".to_string(), "component2".to_string()]),
2111            tvl_gt: Some(1000.0),
2112            chain: Chain::Ethereum,
2113            pagination: PaginationParams::default(),
2114        };
2115
2116        let body2 = ProtocolComponentsRequestBody {
2117            protocol_system: "protocol1".to_string(),
2118            component_ids: Some(vec!["component1".to_string(), "component2".to_string()]),
2119            tvl_gt: Some(1000.0 + 1e-7), // Within the tolerance ±1e-6
2120            chain: Chain::Ethereum,
2121            pagination: PaginationParams::default(),
2122        };
2123
2124        // These should be considered equal due to the tolerance in tvl_gt
2125        assert_eq!(body1, body2);
2126    }
2127
2128    #[test]
2129    fn test_protocol_components_inequality() {
2130        let body1 = ProtocolComponentsRequestBody {
2131            protocol_system: "protocol1".to_string(),
2132            component_ids: Some(vec!["component1".to_string(), "component2".to_string()]),
2133            tvl_gt: Some(1000.0),
2134            chain: Chain::Ethereum,
2135            pagination: PaginationParams::default(),
2136        };
2137
2138        let body2 = ProtocolComponentsRequestBody {
2139            protocol_system: "protocol1".to_string(),
2140            component_ids: Some(vec!["component1".to_string(), "component2".to_string()]),
2141            tvl_gt: Some(1000.0 + 1e-5), // Outside the tolerance ±1e-6
2142            chain: Chain::Ethereum,
2143            pagination: PaginationParams::default(),
2144        };
2145
2146        // These should not be equal due to the difference in tvl_gt
2147        assert_ne!(body1, body2);
2148    }
2149
2150    #[test]
2151    fn test_parse_state_request() {
2152        let json_str = r#"
2153    {
2154        "contractIds": [
2155            "0xb4eccE46b8D4e4abFd03C9B806276A6735C9c092"
2156        ],
2157        "protocol_system": "uniswap_v2",
2158        "version": {
2159            "timestamp": "2069-01-01T04:20:00",
2160            "block": {
2161                "hash": "0x24101f9cb26cd09425b52da10e8c2f56ede94089a8bbe0f31f1cda5f4daa52c4",
2162                "number": 213,
2163                "chain": "ethereum"
2164            }
2165        }
2166    }
2167    "#;
2168
2169        let result: StateRequestBody = serde_json::from_str(json_str).unwrap();
2170
2171        let contract0 = "b4eccE46b8D4e4abFd03C9B806276A6735C9c092"
2172            .parse()
2173            .unwrap();
2174        let block_hash = "24101f9cb26cd09425b52da10e8c2f56ede94089a8bbe0f31f1cda5f4daa52c4"
2175            .parse()
2176            .unwrap();
2177        let block_number = 213;
2178
2179        let expected_timestamp =
2180            NaiveDateTime::parse_from_str("2069-01-01T04:20:00", "%Y-%m-%dT%H:%M:%S").unwrap();
2181
2182        let expected = StateRequestBody {
2183            contract_ids: Some(vec![contract0]),
2184            protocol_system: "uniswap_v2".to_string(),
2185            version: VersionParam {
2186                timestamp: Some(expected_timestamp),
2187                block: Some(BlockParam {
2188                    hash: Some(block_hash),
2189                    chain: Some(Chain::Ethereum),
2190                    number: Some(block_number),
2191                }),
2192            },
2193            chain: Chain::Ethereum,
2194            pagination: PaginationParams::default(),
2195        };
2196
2197        assert_eq!(result, expected);
2198    }
2199
2200    #[test]
2201    fn test_parse_state_request_dual_interface() {
2202        let json_common = r#"
2203    {
2204        "__CONTRACT_IDS__": [
2205            "0xb4eccE46b8D4e4abFd03C9B806276A6735C9c092"
2206        ],
2207        "version": {
2208            "timestamp": "2069-01-01T04:20:00",
2209            "block": {
2210                "hash": "0x24101f9cb26cd09425b52da10e8c2f56ede94089a8bbe0f31f1cda5f4daa52c4",
2211                "number": 213,
2212                "chain": "ethereum"
2213            }
2214        }
2215    }
2216    "#;
2217
2218        let json_str_snake = json_common.replace("\"__CONTRACT_IDS__\"", "\"contract_ids\"");
2219        let json_str_camel = json_common.replace("\"__CONTRACT_IDS__\"", "\"contractIds\"");
2220
2221        let snake: StateRequestBody = serde_json::from_str(&json_str_snake).unwrap();
2222        let camel: StateRequestBody = serde_json::from_str(&json_str_camel).unwrap();
2223
2224        assert_eq!(snake, camel);
2225    }
2226
2227    #[test]
2228    fn test_parse_state_request_unknown_field() {
2229        let body = r#"
2230    {
2231        "contract_ids_with_typo_error": [
2232            {
2233                "address": "0xb4eccE46b8D4e4abFd03C9B806276A6735C9c092",
2234                "chain": "ethereum"
2235            }
2236        ],
2237        "version": {
2238            "timestamp": "2069-01-01T04:20:00",
2239            "block": {
2240                "hash": "0x24101f9cb26cd09425b52da10e8c2f56ede94089a8bbe0f31f1cda5f4daa52c4",
2241                "parentHash": "0x8d75152454e60413efe758cc424bfd339897062d7e658f302765eb7b50971815",
2242                "number": 213,
2243                "chain": "ethereum"
2244            }
2245        }
2246    }
2247    "#;
2248
2249        let decoded = serde_json::from_str::<StateRequestBody>(body);
2250
2251        assert!(decoded.is_err(), "Expected an error due to unknown field");
2252
2253        if let Err(e) = decoded {
2254            assert!(
2255                e.to_string()
2256                    .contains("unknown field `contract_ids_with_typo_error`"),
2257                "Error message does not contain expected unknown field information"
2258            );
2259        }
2260    }
2261
2262    #[test]
2263    fn test_parse_state_request_no_contract_specified() {
2264        let json_str = r#"
2265    {
2266        "protocol_system": "uniswap_v2",
2267        "version": {
2268            "timestamp": "2069-01-01T04:20:00",
2269            "block": {
2270                "hash": "0x24101f9cb26cd09425b52da10e8c2f56ede94089a8bbe0f31f1cda5f4daa52c4",
2271                "number": 213,
2272                "chain": "ethereum"
2273            }
2274        }
2275    }
2276    "#;
2277
2278        let result: StateRequestBody = serde_json::from_str(json_str).unwrap();
2279
2280        let block_hash = "24101f9cb26cd09425b52da10e8c2f56ede94089a8bbe0f31f1cda5f4daa52c4".into();
2281        let block_number = 213;
2282        let expected_timestamp =
2283            NaiveDateTime::parse_from_str("2069-01-01T04:20:00", "%Y-%m-%dT%H:%M:%S").unwrap();
2284
2285        let expected = StateRequestBody {
2286            contract_ids: None,
2287            protocol_system: "uniswap_v2".to_string(),
2288            version: VersionParam {
2289                timestamp: Some(expected_timestamp),
2290                block: Some(BlockParam {
2291                    hash: Some(block_hash),
2292                    chain: Some(Chain::Ethereum),
2293                    number: Some(block_number),
2294                }),
2295            },
2296            chain: Chain::Ethereum,
2297            pagination: PaginationParams { page: 0, page_size: 20 },
2298        };
2299
2300        assert_eq!(result, expected);
2301    }
2302
2303    #[rstest]
2304    #[case::deprecated_ids(
2305        r#"
2306    {
2307        "protocol_ids": [
2308            {
2309                "id": "0xb4eccE46b8D4e4abFd03C9B806276A6735C9c092",
2310                "chain": "ethereum"
2311            }
2312        ],
2313        "protocol_system": "uniswap_v2",
2314        "include_balances": false,
2315        "version": {
2316            "timestamp": "2069-01-01T04:20:00",
2317            "block": {
2318                "hash": "0x24101f9cb26cd09425b52da10e8c2f56ede94089a8bbe0f31f1cda5f4daa52c4",
2319                "number": 213,
2320                "chain": "ethereum"
2321            }
2322        }
2323    }
2324    "#
2325    )]
2326    #[case(
2327        r#"
2328    {
2329        "protocolIds": [
2330            "0xb4eccE46b8D4e4abFd03C9B806276A6735C9c092"
2331        ],
2332        "protocol_system": "uniswap_v2",
2333        "include_balances": false,
2334        "version": {
2335            "timestamp": "2069-01-01T04:20:00",
2336            "block": {
2337                "hash": "0x24101f9cb26cd09425b52da10e8c2f56ede94089a8bbe0f31f1cda5f4daa52c4",
2338                "number": 213,
2339                "chain": "ethereum"
2340            }
2341        }
2342    }
2343    "#
2344    )]
2345    fn test_parse_protocol_state_request(#[case] json_str: &str) {
2346        let result: ProtocolStateRequestBody = serde_json::from_str(json_str).unwrap();
2347
2348        let block_hash = "24101f9cb26cd09425b52da10e8c2f56ede94089a8bbe0f31f1cda5f4daa52c4"
2349            .parse()
2350            .unwrap();
2351        let block_number = 213;
2352
2353        let expected_timestamp =
2354            NaiveDateTime::parse_from_str("2069-01-01T04:20:00", "%Y-%m-%dT%H:%M:%S").unwrap();
2355
2356        let expected = ProtocolStateRequestBody {
2357            protocol_ids: Some(vec!["0xb4eccE46b8D4e4abFd03C9B806276A6735C9c092".to_string()]),
2358            protocol_system: "uniswap_v2".to_string(),
2359            version: VersionParam {
2360                timestamp: Some(expected_timestamp),
2361                block: Some(BlockParam {
2362                    hash: Some(block_hash),
2363                    chain: Some(Chain::Ethereum),
2364                    number: Some(block_number),
2365                }),
2366            },
2367            chain: Chain::Ethereum,
2368            include_balances: false,
2369            pagination: PaginationParams::default(),
2370        };
2371
2372        assert_eq!(result, expected);
2373    }
2374
2375    #[rstest]
2376    #[case::with_protocol_ids(vec![ProtocolId { id: "id1".to_string(), chain: Chain::Ethereum }, ProtocolId { id: "id2".to_string(), chain: Chain::Ethereum }], vec!["id1".to_string(), "id2".to_string()])]
2377    #[case::with_strings(vec!["id1".to_string(), "id2".to_string()], vec!["id1".to_string(), "id2".to_string()])]
2378    fn test_id_filtered<T>(#[case] input_ids: Vec<T>, #[case] expected_ids: Vec<String>)
2379    where
2380        T: Into<String> + Clone,
2381    {
2382        let request_body = ProtocolStateRequestBody::id_filtered(input_ids);
2383
2384        assert_eq!(request_body.protocol_ids, Some(expected_ids));
2385    }
2386
2387    fn create_models_block_changes() -> crate::models::blockchain::BlockAggregatedChanges {
2388        let base_ts = 1694534400; // Example base timestamp for 2023-09-14T00:00:00
2389
2390        crate::models::blockchain::BlockAggregatedChanges {
2391            extractor: "native_name".to_string(),
2392            block: models::blockchain::Block::new(
2393                3,
2394                models::Chain::Ethereum,
2395                Bytes::from_str("0x0000000000000000000000000000000000000000000000000000000000000003").unwrap(),
2396                Bytes::from_str("0x0000000000000000000000000000000000000000000000000000000000000002").unwrap(),
2397                chrono::DateTime::from_timestamp(base_ts + 3000, 0).unwrap().naive_utc(),
2398            ),
2399            db_committed_block_height: Some(1),
2400            finalized_block_height: 1,
2401            revert: true,
2402            state_deltas: HashMap::from([
2403                ("pc_1".to_string(), models::protocol::ProtocolComponentStateDelta {
2404                    component_id: "pc_1".to_string(),
2405                    updated_attributes: HashMap::from([
2406                        ("attr_2".to_string(), Bytes::from("0x0000000000000002")),
2407                        ("attr_1".to_string(), Bytes::from("0x00000000000003e8")),
2408                    ]),
2409                    deleted_attributes: HashSet::new(),
2410                }),
2411            ]),
2412            new_protocol_components: HashMap::from([
2413                ("pc_2".to_string(), crate::models::protocol::ProtocolComponent {
2414                    id: "pc_2".to_string(),
2415                    protocol_system: "native_protocol_system".to_string(),
2416                    protocol_type_name: "pt_1".to_string(),
2417                    chain: models::Chain::Ethereum,
2418                    tokens: vec![
2419                        Bytes::from_str("0xdac17f958d2ee523a2206206994597c13d831ec7").unwrap(),
2420                        Bytes::from_str("0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48").unwrap(),
2421                    ],
2422                    contract_addresses: vec![],
2423                    static_attributes: HashMap::new(),
2424                    change: models::ChangeType::Creation,
2425                    creation_tx: Bytes::from_str("0x000000000000000000000000000000000000000000000000000000000000c351").unwrap(),
2426                    created_at: chrono::DateTime::from_timestamp(base_ts + 5000, 0).unwrap().naive_utc(),
2427                }),
2428            ]),
2429            deleted_protocol_components: HashMap::from([
2430                ("pc_3".to_string(), crate::models::protocol::ProtocolComponent {
2431                    id: "pc_3".to_string(),
2432                    protocol_system: "native_protocol_system".to_string(),
2433                    protocol_type_name: "pt_2".to_string(),
2434                    chain: models::Chain::Ethereum,
2435                    tokens: vec![
2436                        Bytes::from_str("0x6b175474e89094c44da98b954eedeac495271d0f").unwrap(),
2437                        Bytes::from_str("0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2").unwrap(),
2438                    ],
2439                    contract_addresses: vec![],
2440                    static_attributes: HashMap::new(),
2441                    change: models::ChangeType::Deletion,
2442                    creation_tx: Bytes::from_str("0x0000000000000000000000000000000000000000000000000000000000009c41").unwrap(),
2443                    created_at: chrono::DateTime::from_timestamp(base_ts + 4000, 0).unwrap().naive_utc(),
2444                }),
2445            ]),
2446            component_balances: HashMap::from([
2447                ("pc_1".to_string(), HashMap::from([
2448                    (Bytes::from_str("0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48").unwrap(), models::protocol::ComponentBalance {
2449                        token: Bytes::from_str("0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48").unwrap(),
2450                        balance: Bytes::from("0x00000001"),
2451                        balance_float: 1.0,
2452                        modify_tx: Bytes::from_str("0x0000000000000000000000000000000000000000000000000000000000000000").unwrap(),
2453                        component_id: "pc_1".to_string(),
2454                    }),
2455                    (Bytes::from_str("0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2").unwrap(), models::protocol::ComponentBalance {
2456                        token: Bytes::from_str("0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2").unwrap(),
2457                        balance: Bytes::from("0x000003e8"),
2458                        balance_float: 1000.0,
2459                        modify_tx: Bytes::from_str("0x0000000000000000000000000000000000000000000000000000000000007531").unwrap(),
2460                        component_id: "pc_1".to_string(),
2461                    }),
2462                ])),
2463            ]),
2464            account_balances: HashMap::from([
2465                (Bytes::from_str("0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2").unwrap(), HashMap::from([
2466                    (Bytes::from_str("0x7a250d5630b4cf539739df2c5dacb4c659f2488d").unwrap(), models::contract::AccountBalance {
2467                        account: Bytes::from_str("0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2").unwrap(),
2468                        token: Bytes::from_str("0x7a250d5630b4cf539739df2c5dacb4c659f2488d").unwrap(),
2469                        balance: Bytes::from("0x000003e8"),
2470                        modify_tx: Bytes::from_str("0x0000000000000000000000000000000000000000000000000000000000007531").unwrap(),
2471                    }),
2472                    ])),
2473            ]),
2474            ..Default::default()
2475        }
2476    }
2477
2478    #[test]
2479    fn test_serialize_deserialize_block_changes() {
2480        // Test that models::BlockAggregatedChanges serialized as json can be deserialized as
2481        // dto::BlockChanges.
2482
2483        // Create a models::BlockAggregatedChanges instance
2484        let block_entity_changes = create_models_block_changes();
2485
2486        // Serialize the struct into JSON
2487        let json_data = serde_json::to_string(&block_entity_changes).expect("Failed to serialize");
2488
2489        // Deserialize the JSON back into a dto::BlockChanges struct
2490        serde_json::from_str::<BlockChanges>(&json_data).expect("parsing failed");
2491    }
2492
2493    #[test]
2494    fn test_parse_block_changes() {
2495        let json_data = r#"
2496        {
2497            "extractor": "vm:ambient",
2498            "chain": "ethereum",
2499            "block": {
2500                "number": 123,
2501                "hash": "0x0000000000000000000000000000000000000000000000000000000000000000",
2502                "parent_hash": "0x0000000000000000000000000000000000000000000000000000000000000000",
2503                "chain": "ethereum",
2504                "ts": "2023-09-14T00:00:00"
2505            },
2506            "finalized_block_height": 0,
2507            "revert": false,
2508            "new_tokens": {},
2509            "account_updates": {
2510                "0x7a250d5630b4cf539739df2c5dacb4c659f2488d": {
2511                    "address": "0x7a250d5630b4cf539739df2c5dacb4c659f2488d",
2512                    "chain": "ethereum",
2513                    "slots": {},
2514                    "balance": "0x01f4",
2515                    "code": "",
2516                    "change": "Update"
2517                }
2518            },
2519            "state_updates": {
2520                "component_1": {
2521                    "component_id": "component_1",
2522                    "updated_attributes": {"attr1": "0x01"},
2523                    "deleted_attributes": ["attr2"]
2524                }
2525            },
2526            "new_protocol_components":
2527                { "protocol_1": {
2528                        "id": "protocol_1",
2529                        "protocol_system": "system_1",
2530                        "protocol_type_name": "type_1",
2531                        "chain": "ethereum",
2532                        "tokens": ["0x01", "0x02"],
2533                        "contract_ids": ["0x01", "0x02"],
2534                        "static_attributes": {"attr1": "0x01f4"},
2535                        "change": "Update",
2536                        "creation_tx": "0x01",
2537                        "created_at": "2023-09-14T00:00:00"
2538                    }
2539                },
2540            "deleted_protocol_components": {},
2541            "component_balances": {
2542                "protocol_1":
2543                    {
2544                        "0x01": {
2545                            "token": "0x01",
2546                            "balance": "0xb77831d23691653a01",
2547                            "balance_float": 3.3844151001790677e21,
2548                            "modify_tx": "0x01",
2549                            "component_id": "protocol_1"
2550                        }
2551                    }
2552            },
2553            "account_balances": {
2554                "0x7a250d5630b4cf539739df2c5dacb4c659f2488d": {
2555                    "0x7a250d5630b4cf539739df2c5dacb4c659f2488d": {
2556                        "account": "0x7a250d5630b4cf539739df2c5dacb4c659f2488d",
2557                        "token": "0x7a250d5630b4cf539739df2c5dacb4c659f2488d",
2558                        "balance": "0x01f4",
2559                        "modify_tx": "0x01"
2560                    }
2561                }
2562            },
2563            "component_tvl": {
2564                "protocol_1": 1000.0
2565            },
2566            "dci_update": {
2567                "new_entrypoints": {
2568                    "component_1": [
2569                        {
2570                            "external_id": "0x01:sig()",
2571                            "target": "0x01",
2572                            "signature": "sig()"
2573                        }
2574                    ]
2575                },
2576                "new_entrypoint_params": {
2577                    "0x01:sig()": [
2578                        [
2579                            {
2580                                "method": "rpctracer",
2581                                "caller": "0x01",
2582                                "calldata": "0x02"
2583                            },
2584                            "component_1"
2585                        ]
2586                    ]
2587                },
2588                "trace_results": {
2589                    "0x01:sig()": {
2590                        "retriggers": [
2591                            ["0x01", {"key": "0x02", "offset": 12}]
2592                        ],
2593                        "accessed_slots": {
2594                            "0x03": ["0x03", "0x04"]
2595                        }
2596                    }
2597                }
2598            }
2599        }
2600        "#;
2601
2602        serde_json::from_str::<BlockChanges>(json_data).expect("parsing failed");
2603    }
2604
2605    #[test]
2606    fn test_parse_websocket_message() {
2607        let json_data = r#"
2608        {
2609            "subscription_id": "5d23bfbe-89ad-4ea3-8672-dc9e973ac9dc",
2610            "deltas": {
2611                "type": "BlockChanges",
2612                "extractor": "uniswap_v2",
2613                "chain": "ethereum",
2614                "block": {
2615                    "number": 19291517,
2616                    "hash": "0xbc3ea4896c0be8da6229387a8571b72818aa258daf4fab46471003ad74c4ee83",
2617                    "parent_hash": "0x89ca5b8d593574cf6c886f41ef8208bf6bdc1a90ef36046cb8c84bc880b9af8f",
2618                    "chain": "ethereum",
2619                    "ts": "2024-02-23T16:35:35"
2620                },
2621                "finalized_block_height": 0,
2622                "revert": false,
2623                "new_tokens": {},
2624                "account_updates": {
2625                    "0x7a250d5630b4cf539739df2c5dacb4c659f2488d": {
2626                        "address": "0x7a250d5630b4cf539739df2c5dacb4c659f2488d",
2627                        "chain": "ethereum",
2628                        "slots": {},
2629                        "balance": "0x01f4",
2630                        "code": "",
2631                        "change": "Update"
2632                    }
2633                },
2634                "state_updates": {
2635                    "0xde6faedbcae38eec6d33ad61473a04a6dd7f6e28": {
2636                        "component_id": "0xde6faedbcae38eec6d33ad61473a04a6dd7f6e28",
2637                        "updated_attributes": {
2638                            "reserve0": "0x87f7b5973a7f28a8b32404",
2639                            "reserve1": "0x09e9564b11"
2640                        },
2641                        "deleted_attributes": []
2642                    },
2643                    "0x99c59000f5a76c54c4fd7d82720c045bdcf1450d": {
2644                        "component_id": "0x99c59000f5a76c54c4fd7d82720c045bdcf1450d",
2645                        "updated_attributes": {
2646                            "reserve1": "0x44d9a8fd662c2f4d03",
2647                            "reserve0": "0x500b1261f811d5bf423e"
2648                        },
2649                        "deleted_attributes": []
2650                    }
2651                },
2652                "new_protocol_components": {},
2653                "deleted_protocol_components": {},
2654                "component_balances": {
2655                    "0x99c59000f5a76c54c4fd7d82720c045bdcf1450d": {
2656                        "0x9012744b7a564623b6c3e40b144fc196bdedf1a9": {
2657                            "token": "0x9012744b7a564623b6c3e40b144fc196bdedf1a9",
2658                            "balance": "0x500b1261f811d5bf423e",
2659                            "balance_float": 3.779935574269033E23,
2660                            "modify_tx": "0xe46c4db085fb6c6f3408a65524555797adb264e1d5cf3b66ad154598f85ac4bf",
2661                            "component_id": "0x99c59000f5a76c54c4fd7d82720c045bdcf1450d"
2662                        },
2663                        "0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2": {
2664                            "token": "0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2",
2665                            "balance": "0x44d9a8fd662c2f4d03",
2666                            "balance_float": 1.270062661329837E21,
2667                            "modify_tx": "0xe46c4db085fb6c6f3408a65524555797adb264e1d5cf3b66ad154598f85ac4bf",
2668                            "component_id": "0x99c59000f5a76c54c4fd7d82720c045bdcf1450d"
2669                        }
2670                    }
2671                },
2672                "account_balances": {
2673                    "0x7a250d5630b4cf539739df2c5dacb4c659f2488d": {
2674                        "0x7a250d5630b4cf539739df2c5dacb4c659f2488d": {
2675                            "account": "0x7a250d5630b4cf539739df2c5dacb4c659f2488d",
2676                            "token": "0x7a250d5630b4cf539739df2c5dacb4c659f2488d",
2677                            "balance": "0x01f4",
2678                            "modify_tx": "0x01"
2679                        }
2680                    }
2681                },
2682                "component_tvl": {},
2683                "dci_update": {
2684                    "new_entrypoints": {
2685                        "0xde6faedbcae38eec6d33ad61473a04a6dd7f6e28": [
2686                            {
2687                                "external_id": "0x01:sig()",
2688                                "target": "0x01",
2689                                "signature": "sig()"
2690                            }
2691                        ]
2692                    },
2693                    "new_entrypoint_params": {
2694                        "0x01:sig()": [
2695                            [
2696                                {
2697                                    "method": "rpctracer",
2698                                    "caller": "0x01",
2699                                    "calldata": "0x02"
2700                                },
2701                                "0xde6faedbcae38eec6d33ad61473a04a6dd7f6e28"
2702                            ]
2703                        ]
2704                    },
2705                    "trace_results": {
2706                        "0x01:sig()": {
2707                            "retriggers": [
2708                                ["0x01", {"key": "0x02", "offset": 12}]
2709                            ],
2710                            "accessed_slots": {
2711                                "0x03": ["0x03", "0x04"]
2712                            }
2713                        }
2714                    }
2715                }
2716            }
2717        }
2718        "#;
2719        serde_json::from_str::<WebSocketMessage>(json_data).expect("parsing failed");
2720    }
2721
2722    #[test]
2723    fn test_protocol_state_delta_merge_update_delete() {
2724        // Initialize ProtocolStateDelta instances
2725        let mut delta1 = ProtocolStateDelta {
2726            component_id: "Component1".to_string(),
2727            updated_attributes: HashMap::from([(
2728                "Attribute1".to_string(),
2729                Bytes::from("0xbadbabe420"),
2730            )]),
2731            deleted_attributes: HashSet::new(),
2732        };
2733        let delta2 = ProtocolStateDelta {
2734            component_id: "Component1".to_string(),
2735            updated_attributes: HashMap::from([(
2736                "Attribute2".to_string(),
2737                Bytes::from("0x0badbabe"),
2738            )]),
2739            deleted_attributes: HashSet::from(["Attribute1".to_string()]),
2740        };
2741        let exp = ProtocolStateDelta {
2742            component_id: "Component1".to_string(),
2743            updated_attributes: HashMap::from([(
2744                "Attribute2".to_string(),
2745                Bytes::from("0x0badbabe"),
2746            )]),
2747            deleted_attributes: HashSet::from(["Attribute1".to_string()]),
2748        };
2749
2750        delta1.merge(&delta2);
2751
2752        assert_eq!(delta1, exp);
2753    }
2754
2755    #[test]
2756    fn test_protocol_state_delta_merge_delete_update() {
2757        // Initialize ProtocolStateDelta instances
2758        let mut delta1 = ProtocolStateDelta {
2759            component_id: "Component1".to_string(),
2760            updated_attributes: HashMap::new(),
2761            deleted_attributes: HashSet::from(["Attribute1".to_string()]),
2762        };
2763        let delta2 = ProtocolStateDelta {
2764            component_id: "Component1".to_string(),
2765            updated_attributes: HashMap::from([(
2766                "Attribute1".to_string(),
2767                Bytes::from("0x0badbabe"),
2768            )]),
2769            deleted_attributes: HashSet::new(),
2770        };
2771        let exp = ProtocolStateDelta {
2772            component_id: "Component1".to_string(),
2773            updated_attributes: HashMap::from([(
2774                "Attribute1".to_string(),
2775                Bytes::from("0x0badbabe"),
2776            )]),
2777            deleted_attributes: HashSet::new(),
2778        };
2779
2780        delta1.merge(&delta2);
2781
2782        assert_eq!(delta1, exp);
2783    }
2784
2785    #[test]
2786    fn test_account_update_merge() {
2787        // Initialize AccountUpdate instances with same address and valid hex strings for Bytes
2788        let mut account1 = AccountUpdate::new(
2789            Bytes::from(b"0x1234"),
2790            Chain::Ethereum,
2791            HashMap::from([(Bytes::from("0xaabb"), Bytes::from("0xccdd"))]),
2792            Some(Bytes::from("0x1000")),
2793            Some(Bytes::from("0xdeadbeaf")),
2794            ChangeType::Creation,
2795        );
2796
2797        let account2 = AccountUpdate::new(
2798            Bytes::from(b"0x1234"), // Same id as account1
2799            Chain::Ethereum,
2800            HashMap::from([(Bytes::from("0xeeff"), Bytes::from("0x11223344"))]),
2801            Some(Bytes::from("0x2000")),
2802            Some(Bytes::from("0xcafebabe")),
2803            ChangeType::Update,
2804        );
2805
2806        // Merge account2 into account1
2807        account1.merge(&account2);
2808
2809        // Define the expected state after merge
2810        let expected = AccountUpdate::new(
2811            Bytes::from(b"0x1234"), // Same id as before the merge
2812            Chain::Ethereum,
2813            HashMap::from([
2814                (Bytes::from("0xaabb"), Bytes::from("0xccdd")), // Original slot from account1
2815                (Bytes::from("0xeeff"), Bytes::from("0x11223344")), // New slot from account2
2816            ]),
2817            Some(Bytes::from("0x2000")),     // Updated balance
2818            Some(Bytes::from("0xcafebabe")), // Updated code
2819            ChangeType::Creation,            // Updated change type
2820        );
2821
2822        // Assert the new account1 equals to the expected state
2823        assert_eq!(account1, expected);
2824    }
2825
2826    #[test]
2827    fn test_block_account_changes_merge() {
2828        // Prepare account updates
2829        let old_account_updates: HashMap<Bytes, AccountUpdate> = [(
2830            Bytes::from("0x0011"),
2831            AccountUpdate {
2832                address: Bytes::from("0x00"),
2833                chain: Chain::Ethereum,
2834                slots: HashMap::from([(Bytes::from("0x0022"), Bytes::from("0x0033"))]),
2835                balance: Some(Bytes::from("0x01")),
2836                code: Some(Bytes::from("0x02")),
2837                change: ChangeType::Creation,
2838            },
2839        )]
2840        .into_iter()
2841        .collect();
2842        let new_account_updates: HashMap<Bytes, AccountUpdate> = [(
2843            Bytes::from("0x0011"),
2844            AccountUpdate {
2845                address: Bytes::from("0x00"),
2846                chain: Chain::Ethereum,
2847                slots: HashMap::from([(Bytes::from("0x0044"), Bytes::from("0x0055"))]),
2848                balance: Some(Bytes::from("0x03")),
2849                code: Some(Bytes::from("0x04")),
2850                change: ChangeType::Update,
2851            },
2852        )]
2853        .into_iter()
2854        .collect();
2855        // Create initial and new BlockAccountChanges instances
2856        let block_account_changes_initial = BlockChanges {
2857            extractor: "extractor1".to_string(),
2858            revert: false,
2859            account_updates: old_account_updates,
2860            ..Default::default()
2861        };
2862
2863        let block_account_changes_new = BlockChanges {
2864            extractor: "extractor2".to_string(),
2865            revert: true,
2866            account_updates: new_account_updates,
2867            ..Default::default()
2868        };
2869
2870        // Merge the new BlockChanges into the initial one
2871        let res = block_account_changes_initial.merge(block_account_changes_new);
2872
2873        // Create the expected result of the merge operation
2874        let expected_account_updates: HashMap<Bytes, AccountUpdate> = [(
2875            Bytes::from("0x0011"),
2876            AccountUpdate {
2877                address: Bytes::from("0x00"),
2878                chain: Chain::Ethereum,
2879                slots: HashMap::from([
2880                    (Bytes::from("0x0044"), Bytes::from("0x0055")),
2881                    (Bytes::from("0x0022"), Bytes::from("0x0033")),
2882                ]),
2883                balance: Some(Bytes::from("0x03")),
2884                code: Some(Bytes::from("0x04")),
2885                change: ChangeType::Creation,
2886            },
2887        )]
2888        .into_iter()
2889        .collect();
2890        let block_account_changes_expected = BlockChanges {
2891            extractor: "extractor1".to_string(),
2892            revert: true,
2893            account_updates: expected_account_updates,
2894            ..Default::default()
2895        };
2896        assert_eq!(res, block_account_changes_expected);
2897    }
2898
2899    #[test]
2900    fn test_block_entity_changes_merge() {
2901        // Initialize two BlockChanges instances with different details
2902        let block_entity_changes_result1 = BlockChanges {
2903            extractor: String::from("extractor1"),
2904            revert: false,
2905            state_updates: hashmap! { "state1".to_string() => ProtocolStateDelta::default() },
2906            new_protocol_components: hashmap! { "component1".to_string() => ProtocolComponent::default() },
2907            deleted_protocol_components: HashMap::new(),
2908            component_balances: hashmap! {
2909                "component1".to_string() => TokenBalances(hashmap! {
2910                    Bytes::from("0x01") => ComponentBalance {
2911                            token: Bytes::from("0x01"),
2912                            balance: Bytes::from("0x01"),
2913                            balance_float: 1.0,
2914                            modify_tx: Bytes::from("0x00"),
2915                            component_id: "component1".to_string()
2916                        },
2917                    Bytes::from("0x02") => ComponentBalance {
2918                        token: Bytes::from("0x02"),
2919                        balance: Bytes::from("0x02"),
2920                        balance_float: 2.0,
2921                        modify_tx: Bytes::from("0x00"),
2922                        component_id: "component1".to_string()
2923                    },
2924                })
2925
2926            },
2927            component_tvl: hashmap! { "tvl1".to_string() => 1000.0 },
2928            ..Default::default()
2929        };
2930        let block_entity_changes_result2 = BlockChanges {
2931            extractor: String::from("extractor2"),
2932            revert: true,
2933            state_updates: hashmap! { "state2".to_string() => ProtocolStateDelta::default() },
2934            new_protocol_components: hashmap! { "component2".to_string() => ProtocolComponent::default() },
2935            deleted_protocol_components: hashmap! { "component3".to_string() => ProtocolComponent::default() },
2936            component_balances: hashmap! {
2937                "component1".to_string() => TokenBalances::default(),
2938                "component2".to_string() => TokenBalances::default()
2939            },
2940            component_tvl: hashmap! { "tvl2".to_string() => 2000.0 },
2941            ..Default::default()
2942        };
2943
2944        let res = block_entity_changes_result1.merge(block_entity_changes_result2);
2945
2946        let expected_block_entity_changes_result = BlockChanges {
2947            extractor: String::from("extractor1"),
2948            revert: true,
2949            state_updates: hashmap! {
2950                "state1".to_string() => ProtocolStateDelta::default(),
2951                "state2".to_string() => ProtocolStateDelta::default(),
2952            },
2953            new_protocol_components: hashmap! {
2954                "component1".to_string() => ProtocolComponent::default(),
2955                "component2".to_string() => ProtocolComponent::default(),
2956            },
2957            deleted_protocol_components: hashmap! {
2958                "component3".to_string() => ProtocolComponent::default(),
2959            },
2960            component_balances: hashmap! {
2961                "component1".to_string() => TokenBalances(hashmap! {
2962                    Bytes::from("0x01") => ComponentBalance {
2963                            token: Bytes::from("0x01"),
2964                            balance: Bytes::from("0x01"),
2965                            balance_float: 1.0,
2966                            modify_tx: Bytes::from("0x00"),
2967                            component_id: "component1".to_string()
2968                        },
2969                    Bytes::from("0x02") => ComponentBalance {
2970                        token: Bytes::from("0x02"),
2971                        balance: Bytes::from("0x02"),
2972                        balance_float: 2.0,
2973                        modify_tx: Bytes::from("0x00"),
2974                        component_id: "component1".to_string()
2975                        },
2976                    }),
2977                "component2".to_string() => TokenBalances::default(),
2978            },
2979            component_tvl: hashmap! {
2980                "tvl1".to_string() => 1000.0,
2981                "tvl2".to_string() => 2000.0
2982            },
2983            ..Default::default()
2984        };
2985
2986        assert_eq!(res, expected_block_entity_changes_result);
2987    }
2988
2989    #[test]
2990    fn test_websocket_error_serialization() {
2991        let extractor_id = ExtractorIdentity::new(Chain::Ethereum, "test_extractor");
2992        let subscription_id = Uuid::new_v4();
2993
2994        // Test ExtractorNotFound serialization
2995        let error = WebsocketError::ExtractorNotFound(extractor_id.clone());
2996        let json = serde_json::to_string(&error).unwrap();
2997        let deserialized: WebsocketError = serde_json::from_str(&json).unwrap();
2998        assert_eq!(error, deserialized);
2999
3000        // Test SubscriptionNotFound serialization
3001        let error = WebsocketError::SubscriptionNotFound(subscription_id);
3002        let json = serde_json::to_string(&error).unwrap();
3003        let deserialized: WebsocketError = serde_json::from_str(&json).unwrap();
3004        assert_eq!(error, deserialized);
3005
3006        // Test ParseError serialization
3007        let error = WebsocketError::ParseError("{asd".to_string(), "invalid json".to_string());
3008        let json = serde_json::to_string(&error).unwrap();
3009        let deserialized: WebsocketError = serde_json::from_str(&json).unwrap();
3010        assert_eq!(error, deserialized);
3011
3012        // Test SubscribeError serialization
3013        let error = WebsocketError::SubscribeError(extractor_id.clone());
3014        let json = serde_json::to_string(&error).unwrap();
3015        let deserialized: WebsocketError = serde_json::from_str(&json).unwrap();
3016        assert_eq!(error, deserialized);
3017
3018        // Test CompressionError serialization
3019        let error =
3020            WebsocketError::CompressionError(subscription_id, "Compression failed".to_string());
3021        let json = serde_json::to_string(&error).unwrap();
3022        let deserialized: WebsocketError = serde_json::from_str(&json).unwrap();
3023        assert_eq!(error, deserialized);
3024    }
3025
3026    #[test]
3027    fn test_websocket_message_with_error_response() {
3028        let error =
3029            WebsocketError::ParseError("}asdfas".to_string(), "malformed request".to_string());
3030        let response = Response::Error(error.clone());
3031        let message = WebSocketMessage::Response(response);
3032
3033        let json = serde_json::to_string(&message).unwrap();
3034        let deserialized: WebSocketMessage = serde_json::from_str(&json).unwrap();
3035
3036        if let WebSocketMessage::Response(Response::Error(deserialized_error)) = deserialized {
3037            assert_eq!(error, deserialized_error);
3038        } else {
3039            panic!("Expected WebSocketMessage::Response(Response::Error)");
3040        }
3041    }
3042
3043    #[test]
3044    fn test_websocket_error_conversion_from_models() {
3045        use crate::models::error::WebsocketError as ModelsError;
3046
3047        let extractor_id =
3048            crate::models::ExtractorIdentity::new(crate::models::Chain::Ethereum, "test");
3049        let subscription_id = Uuid::new_v4();
3050
3051        // Test ExtractorNotFound conversion
3052        let models_error = ModelsError::ExtractorNotFound(extractor_id.clone());
3053        let dto_error: WebsocketError = models_error.into();
3054        assert_eq!(dto_error, WebsocketError::ExtractorNotFound(extractor_id.clone().into()));
3055
3056        // Test SubscriptionNotFound conversion
3057        let models_error = ModelsError::SubscriptionNotFound(subscription_id);
3058        let dto_error: WebsocketError = models_error.into();
3059        assert_eq!(dto_error, WebsocketError::SubscriptionNotFound(subscription_id));
3060
3061        // Test ParseError conversion - create a real JSON parse error
3062        let json_result: Result<serde_json::Value, _> = serde_json::from_str("{invalid json");
3063        let json_error = json_result.unwrap_err();
3064        let models_error = ModelsError::ParseError("{invalid json".to_string(), json_error);
3065        let dto_error: WebsocketError = models_error.into();
3066        if let WebsocketError::ParseError(msg, error) = dto_error {
3067            // Just check that we have a non-empty error message
3068            assert!(!error.is_empty(), "Error message should not be empty, got: '{}'", msg);
3069        } else {
3070            panic!("Expected ParseError variant");
3071        }
3072
3073        // Test SubscribeError conversion
3074        let models_error = ModelsError::SubscribeError(extractor_id.clone());
3075        let dto_error: WebsocketError = models_error.into();
3076        assert_eq!(dto_error, WebsocketError::SubscribeError(extractor_id.into()));
3077
3078        // Test CompressionError conversion
3079        let io_error = std::io::Error::other("Compression failed");
3080        let models_error = ModelsError::CompressionError(subscription_id, io_error);
3081        let dto_error: WebsocketError = models_error.into();
3082        if let WebsocketError::CompressionError(sub_id, msg) = &dto_error {
3083            assert_eq!(*sub_id, subscription_id);
3084            assert!(msg.contains("Compression failed"));
3085        } else {
3086            panic!("Expected CompressionError variant");
3087        }
3088    }
3089}
3090
3091#[cfg(test)]
3092mod memory_size_tests {
3093    use std::collections::HashMap;
3094
3095    use super::*;
3096
3097    #[test]
3098    fn test_state_request_response_memory_size_empty() {
3099        let response = StateRequestResponse {
3100            accounts: vec![],
3101            pagination: PaginationResponse::new(1, 10, 0),
3102        };
3103
3104        let size = response.deep_size_of();
3105
3106        // Should at least include base struct sizes
3107        assert!(size >= 48, "Empty response should have minimum size of 48 bytes, got {}", size);
3108        assert!(size < 200, "Empty response should not be too large, got {}", size);
3109    }
3110
3111    #[test]
3112    fn test_state_request_response_memory_size_scales_with_slots() {
3113        let create_response_with_slots = |slot_count: usize| {
3114            let mut slots = HashMap::new();
3115            for i in 0..slot_count {
3116                let key = vec![i as u8; 32]; // 32-byte key
3117                let value = vec![(i + 100) as u8; 32]; // 32-byte value
3118                slots.insert(key.into(), value.into());
3119            }
3120
3121            let account = ResponseAccount::new(
3122                Chain::Ethereum,
3123                vec![1; 20].into(),
3124                "Pool".to_string(),
3125                slots,
3126                vec![1; 32].into(),
3127                HashMap::new(),
3128                vec![].into(), // empty code
3129                vec![1; 32].into(),
3130                vec![1; 32].into(),
3131                vec![1; 32].into(),
3132                None,
3133            );
3134
3135            StateRequestResponse {
3136                accounts: vec![account],
3137                pagination: PaginationResponse::new(1, 10, 1),
3138            }
3139        };
3140
3141        let small_response = create_response_with_slots(10);
3142        let large_response = create_response_with_slots(100);
3143
3144        let small_size = small_response.deep_size_of();
3145        let large_size = large_response.deep_size_of();
3146
3147        // Large response should be significantly bigger
3148        assert!(
3149            large_size > small_size * 5,
3150            "Large response ({} bytes) should be much larger than small response ({} bytes)",
3151            large_size,
3152            small_size
3153        );
3154
3155        // Each slot should contribute at least 64 bytes (32 + 32 + overhead)
3156        let size_diff = large_size - small_size;
3157        let expected_min_diff = 90 * 64; // 90 additional slots * 64 bytes each
3158        assert!(
3159            size_diff > expected_min_diff,
3160            "Size difference ({} bytes) should reflect the additional slot data",
3161            size_diff
3162        );
3163    }
3164}