1use std::{
2 collections::{hash_map::Entry, HashMap, HashSet},
3 future::Future,
4 pin::Pin,
5 sync::Arc,
6};
7
8use alloy::primitives::{Address, U256};
9use thiserror::Error;
10use tokio::sync::{watch, RwLock, RwLockReadGuard};
11use tracing::{debug, error, info, warn};
12use tycho_client::feed::{synchronizer::ComponentWithState, BlockHeader, FeedMessage, HeaderLike};
13use tycho_common::{
14 dto::{ChangeType, ProtocolStateDelta},
15 models::{blockchain::BlockAggregatedChanges, token::Token, Chain},
16 simulation::protocol_sim::{Balances, ProtocolSim},
17 Bytes,
18};
19#[cfg(test)]
20use {
21 mockall::mock,
22 num_bigint::BigUint,
23 std::any::Any,
24 tycho_common::simulation::{
25 errors::{SimulationError, TransitionError},
26 protocol_sim::GetAmountOutResult,
27 },
28};
29
30use crate::{
31 evm::{
32 engine_db::{update_engine, SHARED_TYCHO_DB},
33 override_stream::{OverrideSnapshot, StateOverrideProvider},
34 protocol::{
35 utils::bytes_to_address,
36 vm::{constants::ERC20_PROXY_BYTECODE, erc20_token::IMPLEMENTATION_SLOT},
37 },
38 tycho_models::{AccountUpdate, ResponseAccount},
39 },
40 protocol::{
41 errors::InvalidSnapshotError,
42 models::{DecoderContext, ProtocolComponent, TryFromWithBlock, Update},
43 },
44};
45
46#[derive(Error, Debug)]
47pub enum StreamDecodeError {
48 #[error("{0}")]
49 Fatal(String),
50}
51
52#[derive(Default)]
53struct DecoderState {
54 tokens: HashMap<Bytes, Token>,
55 states: HashMap<String, Box<dyn ProtocolSim>>,
56 components: HashMap<String, ProtocolComponent>,
57 contracts_map: HashMap<Bytes, HashSet<String>>,
59 proxy_token_addresses: HashMap<Address, Address>,
61 failed_components: HashSet<String>,
65 current_block_number: u64,
67}
68
69type DecodeFut =
70 Pin<Box<dyn Future<Output = Result<Box<dyn ProtocolSim>, InvalidSnapshotError>> + Send + Sync>>;
71type AccountBalances = HashMap<Bytes, HashMap<Bytes, Bytes>>;
72type RegistryFn<H> = dyn Fn(
73 ComponentWithState,
74 H,
75 AccountBalances,
76 Arc<RwLock<DecoderState>>,
77 Option<watch::Receiver<OverrideSnapshot>>,
78 ) -> DecodeFut
79 + Send
80 + Sync;
81type FilterFn = fn(&ComponentWithState) -> bool;
82
83pub struct TychoStreamDecoder<H>
96where
97 H: HeaderLike,
98{
99 state: Arc<RwLock<DecoderState>>,
100 skip_state_decode_failures: bool,
101 min_token_quality: u32,
102 registry: HashMap<String, Box<RegistryFn<H>>>,
103 inclusion_filters: HashMap<String, Vec<FilterFn>>,
104 override_providers: HashMap<String, Arc<dyn StateOverrideProvider>>,
107}
108
109impl<H> Default for TychoStreamDecoder<H>
110where
111 H: HeaderLike + Clone + Sync + Send + 'static + std::fmt::Debug,
112{
113 fn default() -> Self {
114 Self::new()
115 }
116}
117
118fn is_deprecated_curve_registration<T: 'static>(exchange: &str) -> bool {
122 exchange == "vm:curve" &&
123 std::any::type_name::<T>() !=
124 std::any::type_name::<crate::evm::protocol::curve::CurveState>()
125}
126
127impl<H> TychoStreamDecoder<H>
128where
129 H: HeaderLike + Clone + Sync + Send + 'static + std::fmt::Debug,
130{
131 pub fn new() -> Self {
132 Self {
133 state: Arc::new(RwLock::new(DecoderState::default())),
134 skip_state_decode_failures: false,
135 min_token_quality: 100,
136 registry: HashMap::new(),
137 inclusion_filters: HashMap::new(),
138 override_providers: HashMap::new(),
139 }
140 }
141
142 pub fn set_override_provider(
148 &mut self,
149 protocol_system: String,
150 provider: Arc<dyn StateOverrideProvider>,
151 ) {
152 self.override_providers
153 .insert(protocol_system, provider);
154 }
155
156 pub async fn set_tokens(&self, tokens: HashMap<Bytes, Token>) {
161 let mut guard = self.state.write().await;
162 guard.tokens = tokens;
163 }
164
165 pub fn skip_state_decode_failures(&mut self, skip: bool) {
166 self.skip_state_decode_failures = skip;
167 }
168
169 pub fn min_token_quality(&mut self, quality: u32) {
175 self.min_token_quality = quality;
176 }
177
178 pub fn register_decoder_with_context<T>(&mut self, exchange: &str, context: DecoderContext)
191 where
192 T: ProtocolSim
193 + TryFromWithBlock<ComponentWithState, H, Error = InvalidSnapshotError>
194 + Send
195 + 'static,
196 {
197 if is_deprecated_curve_registration::<T>(exchange) {
198 warn!(
199 registered_type = std::any::type_name::<T>(),
200 "Registering \"vm:curve\" with the generic VM adapter is deprecated; register the \
201 native `CurveState` decoder instead (`exchange::<CurveState>(\"vm:curve\", ...)`). \
202 The VM-adapter path still works but will be removed in a future release."
203 );
204 }
205 let decoder = Box::new(
206 move |component: ComponentWithState,
207 header: H,
208 account_balances: AccountBalances,
209 state: Arc<RwLock<DecoderState>>,
210 live_override: Option<watch::Receiver<OverrideSnapshot>>| {
211 let mut context = context.clone();
212 context.live_override = live_override;
213 Box::pin(async move {
214 let guard = state.read().await;
215 T::try_from_with_header(
216 component,
217 header,
218 &account_balances,
219 &guard.tokens,
220 &context,
221 )
222 .await
223 .map(|c| Box::new(c) as Box<dyn ProtocolSim>)
224 }) as DecodeFut
225 },
226 );
227 self.registry
228 .insert(exchange.to_string(), decoder);
229 }
230
231 pub fn register_decoder<T>(&mut self, exchange: &str)
243 where
244 T: ProtocolSim
245 + TryFromWithBlock<ComponentWithState, H, Error = InvalidSnapshotError>
246 + Send
247 + 'static,
248 {
249 let context = DecoderContext::new();
250 self.register_decoder_with_context::<T>(exchange, context);
251 }
252
253 pub fn register_filter(&mut self, exchange: &str, predicate: FilterFn) {
272 self.inclusion_filters
273 .entry(exchange.to_string())
274 .or_default()
275 .push(predicate);
276 }
277
278 fn admits(&self, exchange: &str, snapshot: &ComponentWithState) -> bool {
281 let Some(predicates) = self.inclusion_filters.get(exchange) else { return true };
282 predicates
283 .iter()
284 .all(|predicate| predicate(snapshot))
285 }
286
287 pub async fn decode(&self, msg: &FeedMessage<H>) -> Result<Update, StreamDecodeError> {
290 let mut updated_states = HashMap::new();
292 let mut new_pairs = HashMap::new();
293 let mut removed_pairs = HashMap::new();
294 let mut contracts_map = HashMap::new();
295 let mut msg_failed_components = HashSet::new();
296
297 let header = msg
298 .state_msgs
299 .values()
300 .next()
301 .ok_or_else(|| StreamDecodeError::Fatal("Missing block!".into()))?
302 .header
303 .clone();
304
305 let block_number_or_timestamp = header
306 .clone()
307 .block_number_or_timestamp();
308 let current_block = header.clone().block();
309 let is_partial = current_block
310 .as_ref()
311 .map(|h| h.partial_block_index.is_some())
312 .unwrap_or(false);
313
314 for (protocol, protocol_msg) in msg.state_msgs.iter() {
315 if let Some(deltas) = protocol_msg.deltas.as_ref() {
317 let mut state_guard = self.state.write().await;
318
319 let new_tokens = deltas
320 .new_tokens
321 .iter()
322 .filter(|(addr, t)| {
323 t.quality >= self.min_token_quality &&
324 !state_guard.tokens.contains_key(*addr)
325 })
326 .map(|(addr, t)| (addr.clone(), t.clone()))
327 .collect::<HashMap<Bytes, Token>>();
328
329 if !new_tokens.is_empty() {
330 debug!(n = new_tokens.len(), "NewTokens");
331 state_guard.tokens.extend(new_tokens);
332 }
333 }
334
335 {
337 let mut state_guard = self.state.write().await;
338 let removed_components: Vec<(String, ProtocolComponent)> = protocol_msg
339 .removed_components
340 .iter()
341 .map(|(id, comp)| {
342 if *id != comp.id {
343 error!(
344 "Component id mismatch in removed components {id} != {}",
345 comp.id
346 );
347 return Err(StreamDecodeError::Fatal("Component id mismatch".into()));
348 }
349
350 let tokens = comp
351 .tokens
352 .iter()
353 .flat_map(|addr| state_guard.tokens.get(addr).cloned())
354 .collect::<Vec<_>>();
355
356 if tokens.len() == comp.tokens.len() {
357 Ok(Some((
358 id.clone(),
359 ProtocolComponent::from_with_tokens(comp.clone(), tokens),
360 )))
361 } else {
362 Ok(None)
363 }
364 })
365 .collect::<Result<Vec<Option<(String, ProtocolComponent)>>, StreamDecodeError>>(
366 )?
367 .into_iter()
368 .flatten()
369 .collect();
370
371 for (id, component) in removed_components {
373 state_guard.components.remove(&id);
374 state_guard.states.remove(&id);
375 removed_pairs.insert(id, component);
376 }
377
378 info!(
380 "Processing {} contracts from snapshots",
381 protocol_msg
382 .snapshots
383 .get_vm_storage()
384 .len()
385 );
386
387 let mut proxy_token_accounts: HashMap<Address, AccountUpdate> = HashMap::new();
388 let mut storage_by_address: HashMap<Address, ResponseAccount> = HashMap::new();
389 for (key, value) in protocol_msg
390 .snapshots
391 .get_vm_storage()
392 .iter()
393 {
394 let account: ResponseAccount = value.clone().into();
395
396 if state_guard.tokens.contains_key(key) {
397 let original_address = account.address;
398 let (impl_addr, proxy_state) = match state_guard
407 .proxy_token_addresses
408 .get(&original_address)
409 {
410 Some(impl_addr) => {
411 let proxy_state = AccountUpdate::new(
418 original_address,
419 value.chain,
420 account.slots.clone(),
421 Some(account.native_balance),
422 None,
423 ChangeType::Update,
424 );
425 (*impl_addr, proxy_state)
426 }
427 None => {
428 let impl_addr = generate_proxy_token_address(
432 state_guard.proxy_token_addresses.len() as u32,
433 )?;
434 state_guard
435 .proxy_token_addresses
436 .insert(original_address, impl_addr);
437
438 let proxy_state = create_proxy_token_account(
440 original_address,
441 Some(impl_addr),
442 &account.slots,
443 value.chain,
444 Some(account.native_balance),
445 );
446
447 (impl_addr, proxy_state)
448 }
449 };
450
451 proxy_token_accounts.insert(original_address, proxy_state);
452
453 let impl_update = ResponseAccount {
455 address: impl_addr,
456 slots: HashMap::new(),
457 ..account.clone()
458 };
459 storage_by_address.insert(impl_addr, impl_update);
460 } else {
461 storage_by_address.insert(account.address, account);
463 }
464 }
465
466 let mut proxy_creates: Vec<AccountUpdate> = Vec::new();
470 let mut proxy_updates: HashMap<Address, AccountUpdate> = HashMap::new();
471 for (addr, update) in proxy_token_accounts {
472 if matches!(update.change, ChangeType::Creation) {
473 proxy_creates.push(update);
474 } else {
475 proxy_updates.insert(addr, update);
476 }
477 }
478
479 info!("Updating engine with {} contracts from snapshots", storage_by_address.len());
480 update_engine(
481 SHARED_TYCHO_DB.clone(),
482 header.clone().block(),
483 Some(storage_by_address),
484 proxy_updates,
485 )
486 .map_err(|e| StreamDecodeError::Fatal(e.to_string()))?;
487
488 if !proxy_creates.is_empty() {
492 SHARED_TYCHO_DB
493 .force_update_accounts(proxy_creates)
494 .map_err(|e| StreamDecodeError::Fatal(e.to_string()))?;
495 }
496 info!("Engine updated");
497 drop(state_guard);
498 }
499
500 let account_balances = protocol_msg
503 .clone()
504 .snapshots
505 .get_vm_storage()
506 .iter()
507 .filter_map(|(addr, acc)| {
508 if acc.token_balances.is_empty() {
509 return None;
510 }
511 let balances = acc
512 .token_balances
513 .iter()
514 .map(|(token_addr, ab)| (token_addr.clone(), ab.balance.clone()))
515 .collect::<HashMap<Bytes, Bytes>>();
516 Some((addr.clone(), balances))
517 })
518 .collect::<AccountBalances>();
519
520 let mut new_components = HashMap::new();
521 let mut count_token_skips = 0;
522 let mut components_to_store = HashMap::new();
523 {
524 let state_guard = self.state.read().await;
525
526 'snapshot_loop: for (id, snapshot) in protocol_msg
528 .snapshots
529 .get_states()
530 .clone()
531 {
532 if !self.admits(protocol.as_str(), &snapshot) {
534 continue;
535 }
536
537 let mut component_tokens = Vec::new();
539 let mut new_tokens_accounts = HashMap::new();
540 for token in snapshot.component.tokens.clone() {
541 match state_guard.tokens.get(&token) {
542 Some(token) => {
543 component_tokens.push(token.clone());
544
545 let token_address = match bytes_to_address(&token.address) {
548 Ok(addr) => addr,
549 Err(_) => {
550 count_token_skips += 1;
551 msg_failed_components.insert(id.clone());
552 warn!(
553 "Token address could not be decoded {}, ignoring pool {:x?}",
554 token.address, id
555 );
556 continue 'snapshot_loop;
557 }
558 };
559 if !state_guard
561 .proxy_token_addresses
562 .contains_key(&token_address)
563 {
564 new_tokens_accounts.insert(
565 token_address,
566 create_proxy_token_account(
567 token_address,
568 None,
569 &HashMap::new(),
570 snapshot.component.chain,
571 None,
572 ),
573 );
574 }
575 }
576 None => {
577 count_token_skips += 1;
578 msg_failed_components.insert(id.clone());
579 debug!("Token not found {}, ignoring pool {:x?}", token, id);
580 continue 'snapshot_loop;
581 }
582 }
583 }
584 let component = ProtocolComponent::from_with_tokens(
585 snapshot.component.clone(),
586 component_tokens,
587 );
588
589 if !new_tokens_accounts.is_empty() {
591 update_engine(
592 SHARED_TYCHO_DB.clone(),
593 header.clone().block(),
594 None,
595 new_tokens_accounts,
596 )
597 .map_err(|e| StreamDecodeError::Fatal(e.to_string()))?;
598 }
599
600 if !component
603 .static_attributes
604 .contains_key("manual_updates")
605 {
606 for contract in &component.contract_ids {
607 contracts_map
608 .entry(contract.clone())
609 .or_insert_with(HashSet::new)
610 .insert(id.clone());
611 }
612 for (_, tracing) in snapshot.entrypoints.iter() {
615 for contract in tracing.accessed_slots.keys().cloned() {
616 contracts_map
617 .entry(contract)
618 .or_insert_with(HashSet::new)
619 .insert(id.clone());
620 }
621 }
622 }
623
624 new_pairs.insert(id.clone(), component.clone());
626
627 components_to_store.insert(id.clone(), component);
629
630 if let Some(state_decode_f) = self.registry.get(protocol.as_str()) {
632 let live_override = self
633 .override_providers
634 .get(protocol.as_str())
635 .and_then(|provider| provider.subscribe(protocol.as_str()));
636 match state_decode_f(
637 snapshot,
638 header.clone(),
639 account_balances.clone(),
640 self.state.clone(),
641 live_override,
642 )
643 .await
644 {
645 Ok(state) => {
646 new_components.insert(id.clone(), state);
647 }
648 Err(e) => {
649 if self.skip_state_decode_failures {
650 warn!(pool = id, error = %e, "StateDecodingFailure");
651 msg_failed_components.insert(id.clone());
652 continue 'snapshot_loop;
653 } else {
654 error!(pool = id, error = %e, "StateDecodingFailure");
655 return Err(StreamDecodeError::Fatal(format!("{e}")));
656 }
657 }
658 }
659 } else if self.skip_state_decode_failures {
660 warn!(pool = id, "MissingDecoderRegistration");
661 msg_failed_components.insert(id.clone());
662 continue 'snapshot_loop;
663 } else {
664 error!(pool = id, "MissingDecoderRegistration");
665 return Err(StreamDecodeError::Fatal(format!(
666 "Missing decoder registration for: {id}"
667 )));
668 }
669 }
670 }
671
672 if !components_to_store.is_empty() {
674 let mut state_guard = self.state.write().await;
675 for (id, component) in components_to_store {
676 state_guard
677 .components
678 .insert(id, component);
679 }
680 }
681
682 if !protocol_msg.snapshots.states.is_empty() {
683 info!("Decoded {} snapshots for protocol {protocol}", new_components.len());
684 }
685 if count_token_skips > 0 {
686 info!("Skipped {count_token_skips} pools due to missing tokens");
687 }
688
689 updated_states.extend(new_components);
691
692 if let Some(deltas) = protocol_msg.deltas.clone() {
694 let mut state_guard = self.state.write().await;
696
697 let mut account_update_by_address: HashMap<Address, AccountUpdate> = HashMap::new();
698 let mut new_proxy_accounts: Vec<AccountUpdate> = Vec::new();
700 for (key, value) in deltas.account_deltas.iter() {
701 let mut update: AccountUpdate = value.clone().into();
702
703 if update.code.is_none() && matches!(update.change, ChangeType::Creation) {
709 error!(
710 update = ?update,
711 "FaultyCreationDelta"
712 );
713 update.code = Some(vec![]);
714 }
715
716 if state_guard.tokens.contains_key(key) {
717 let original_address = update.address;
718 let impl_addr = match state_guard
725 .proxy_token_addresses
726 .get(&original_address)
727 {
728 Some(impl_addr) => {
729 let proxy_update = AccountUpdate {
734 code: None,
735 change: ChangeType::Update,
736 ..update.clone()
737 };
738 account_update_by_address.insert(original_address, proxy_update);
739
740 *impl_addr
741 }
742 None => {
743 let impl_addr = generate_proxy_token_address(
748 state_guard.proxy_token_addresses.len() as u32,
749 )?;
750 state_guard
751 .proxy_token_addresses
752 .insert(original_address, impl_addr);
753
754 let proxy_state = create_proxy_token_account(
759 original_address,
760 Some(impl_addr),
761 &update.slots,
762 update.chain,
763 update.balance,
764 );
765 new_proxy_accounts.push(proxy_state);
766
767 impl_addr
768 }
769 };
770
771 if update.code.is_some() {
773 let impl_update = AccountUpdate {
774 address: impl_addr,
775 slots: HashMap::new(),
776 ..update.clone()
777 };
778 account_update_by_address.insert(impl_addr, impl_update);
779 }
780 } else {
781 account_update_by_address.insert(update.address, update);
783 }
784 }
785 drop(state_guard);
786
787 let state_guard = self.state.read().await;
788 info!("Updating engine with {} contract deltas", deltas.account_deltas.len());
789 update_engine(
790 SHARED_TYCHO_DB.clone(),
791 header.clone().block(),
792 None,
793 account_update_by_address,
794 )
795 .map_err(|e| StreamDecodeError::Fatal(e.to_string()))?;
796
797 if !new_proxy_accounts.is_empty() {
800 SHARED_TYCHO_DB
801 .force_update_accounts(new_proxy_accounts)
802 .map_err(|e| StreamDecodeError::Fatal(e.to_string()))?;
803 }
804 info!("Engine updated");
805
806 let mut pools_to_update = HashSet::new();
808 for (account, _update) in deltas.account_deltas {
809 pools_to_update.extend(
811 contracts_map
812 .get(&account)
813 .cloned()
814 .unwrap_or_default(),
815 );
816 pools_to_update.extend(
818 state_guard
819 .contracts_map
820 .get(&account)
821 .cloned()
822 .unwrap_or_default(),
823 );
824 }
825
826 let all_balances = Balances {
828 component_balances: deltas
829 .component_balances
830 .iter()
831 .map(|(pool_id, bals)| {
832 let mut balances = HashMap::new();
833 for (t, b) in bals {
834 balances.insert(t.clone(), b.balance.clone());
835 }
836 pools_to_update.insert(pool_id.clone());
837 (pool_id.clone(), balances)
838 })
839 .collect(),
840 account_balances: deltas
841 .account_balances
842 .iter()
843 .map(|(account, bals)| {
844 let mut balances = HashMap::new();
845 for (t, b) in bals {
846 balances.insert(t.clone(), b.balance.clone());
847 }
848 pools_to_update.extend(
849 contracts_map
850 .get(account)
851 .cloned()
852 .unwrap_or_default(),
853 );
854 (account.clone(), balances)
855 })
856 .collect(),
857 };
858
859 for (id, update) in deltas.state_deltas {
861 let update_with_block = Self::add_block_info_to_delta(
863 ProtocolStateDelta::from(update),
864 current_block.clone(),
865 );
866 match Self::apply_update(
867 &id,
868 update_with_block,
869 &mut updated_states,
870 &state_guard,
871 &all_balances,
872 ) {
873 Ok(_) => {
874 pools_to_update.remove(&id);
875 }
876 Err(e) => {
877 if self.skip_state_decode_failures {
878 warn!(pool = id, error = %e, "Failed to apply state update, marking component as removed");
879 updated_states.remove(&id);
881 if let Some(component) = new_pairs.remove(&id) {
883 removed_pairs.insert(id.clone(), component);
884 } else if let Some(component) = state_guard.components.get(&id) {
885 removed_pairs.insert(id.clone(), component.clone());
886 } else {
887 warn!(pool = id, "Component not found in new_pairs or state, cannot add to removed_pairs");
890 }
891 pools_to_update.remove(&id);
892
893 msg_failed_components.insert(id.clone());
895 } else {
896 return Err(e);
897 }
898 }
899 }
900 }
901
902 for pool in pools_to_update {
904 let default_delta_with_block = Self::add_block_info_to_delta(
906 ProtocolStateDelta::default(),
907 current_block.clone(),
908 );
909 match Self::apply_update(
910 &pool,
911 default_delta_with_block,
912 &mut updated_states,
913 &state_guard,
914 &all_balances,
915 ) {
916 Ok(_) => {}
917 Err(e) => {
918 if self.skip_state_decode_failures {
919 warn!(pool = pool, error = %e, "Failed to apply contract/balance update, marking component as removed");
920 updated_states.remove(&pool);
922 if let Some(component) = new_pairs.remove(&pool) {
924 removed_pairs.insert(pool.clone(), component);
925 } else if let Some(component) = state_guard.components.get(&pool) {
926 removed_pairs.insert(pool.clone(), component.clone());
927 } else {
928 warn!(pool = pool, "Component not found in new_pairs or state, cannot add to removed_pairs");
931 }
932
933 msg_failed_components.insert(pool.clone());
935 } else {
936 return Err(e);
937 }
938 }
939 }
940 }
941 };
942 }
943
944 let mut state_guard = self.state.write().await;
946
947 state_guard
949 .failed_components
950 .extend(msg_failed_components);
951
952 updated_states.retain(|id, _| {
956 !state_guard
957 .failed_components
958 .contains(id)
959 });
960 new_pairs.retain(|id, _| {
961 !state_guard
962 .failed_components
963 .contains(id)
964 });
965
966 state_guard
967 .states
968 .extend(updated_states.clone());
969
970 state_guard.current_block_number = block_number_or_timestamp;
971
972 for (id, component) in new_pairs.iter() {
974 state_guard
975 .components
976 .insert(id.clone(), component.clone());
977 }
978
979 for id in removed_pairs.keys() {
981 state_guard.components.remove(id);
982 }
983
984 for (key, values) in contracts_map {
985 state_guard
986 .contracts_map
987 .entry(key)
988 .or_insert_with(HashSet::new)
989 .extend(values);
990 }
991
992 Ok(Update::new(block_number_or_timestamp, updated_states, new_pairs)
994 .set_is_partial(is_partial)
995 .set_removed_pairs(removed_pairs)
996 .set_sync_states(msg.sync_states.clone()))
997 }
998
999 pub async fn apply_deltas_ephemeral(
1018 &self,
1019 pending_deltas: &HashMap<String, BlockAggregatedChanges>,
1020 header: H,
1021 ) -> Result<Update, StreamDecodeError> {
1022 let block_number_or_timestamp = header
1023 .clone()
1024 .block_number_or_timestamp();
1025 let current_block = header.block();
1026 let state_guard = self.state.read().await;
1027
1028 let mut updated_states: HashMap<String, Box<dyn ProtocolSim>> = HashMap::new();
1029
1030 for deltas in pending_deltas.values() {
1031 let all_balances = Balances {
1032 component_balances: deltas
1033 .component_balances
1034 .iter()
1035 .map(|(pool_id, bals)| {
1036 let balances = bals
1037 .iter()
1038 .map(|(t, b)| (t.clone(), b.balance.clone()))
1039 .collect();
1040 (pool_id.clone(), balances)
1041 })
1042 .collect(),
1043 account_balances: HashMap::new(),
1044 };
1045
1046 for (id, state_delta) in &deltas.state_deltas {
1047 let dto_delta = Self::add_block_info_to_delta(
1048 ProtocolStateDelta::from(state_delta.clone()),
1049 current_block.clone(),
1050 );
1051 if let Err(e) = Self::apply_update(
1052 id,
1053 dto_delta,
1054 &mut updated_states,
1055 &state_guard,
1056 &all_balances,
1057 ) {
1058 warn!(pool = id, error = %e, "EphemeralDeltaTransitionError");
1059 }
1060 }
1061 }
1062
1063 Ok(Update::new(block_number_or_timestamp, updated_states, HashMap::new()))
1064 }
1065
1066 fn add_block_info_to_delta(
1068 mut delta: ProtocolStateDelta,
1069 block_header_opt: Option<BlockHeader>,
1070 ) -> ProtocolStateDelta {
1071 if let Some(header) = block_header_opt {
1072 delta.updated_attributes.insert(
1075 "block_number".to_string(),
1076 Bytes::from(header.number.to_be_bytes().to_vec()),
1077 );
1078 delta.updated_attributes.insert(
1079 "block_timestamp".to_string(),
1080 Bytes::from(header.timestamp.to_be_bytes().to_vec()),
1081 );
1082 }
1083 delta
1084 }
1085
1086 fn apply_update(
1087 id: &String,
1088 update: ProtocolStateDelta,
1089 updated_states: &mut HashMap<String, Box<dyn ProtocolSim>>,
1090 state_guard: &RwLockReadGuard<'_, DecoderState>,
1091 all_balances: &Balances,
1092 ) -> Result<(), StreamDecodeError> {
1093 match updated_states.entry(id.clone()) {
1094 Entry::Occupied(mut entry) => {
1095 let state: &mut Box<dyn ProtocolSim> = entry.get_mut();
1097 state
1098 .delta_transition(update, &state_guard.tokens, all_balances)
1099 .map_err(|e| {
1100 error!(pool = id, error = ?e, "DeltaTransitionError");
1101 StreamDecodeError::Fatal(format!("TransitionFailure: {e:?}"))
1102 })?;
1103 }
1104 Entry::Vacant(_) => {
1105 match state_guard.states.get(id) {
1106 Some(stored_state) => {
1109 let mut state = stored_state.clone();
1110 state
1111 .delta_transition(update, &state_guard.tokens, all_balances)
1112 .map_err(|e| {
1113 error!(pool = id, error = ?e, "DeltaTransitionError");
1114 StreamDecodeError::Fatal(format!("TransitionFailure: {e:?}"))
1115 })?;
1116 updated_states.insert(id.clone(), state);
1117 }
1118 None => debug!(pool = id, reason = "MissingState", "DeltaTransitionError"),
1119 }
1120 }
1121 }
1122 Ok(())
1123 }
1124}
1125
1126fn generate_proxy_token_address(idx: u32) -> Result<Address, StreamDecodeError> {
1128 let padded_idx = format!("{idx:x}");
1129 let padded_zeroes = "0".repeat(33 - padded_idx.len());
1130 let proxy_token_address = format!("{padded_zeroes}{padded_idx}BAdbaBe");
1131 let decoded = hex::decode(proxy_token_address).map_err(|e| {
1132 StreamDecodeError::Fatal(format!("Invalid proxy token address encoding: {e}"))
1133 })?;
1134
1135 const ADDRESS_LENGTH: usize = 20;
1136 if decoded.len() != ADDRESS_LENGTH {
1137 return Err(StreamDecodeError::Fatal(format!(
1138 "Invalid proxy token address length: expected {}, got {}",
1139 ADDRESS_LENGTH,
1140 decoded.len(),
1141 )));
1142 }
1143
1144 Ok(Address::from_slice(&decoded))
1145}
1146
1147fn create_proxy_token_account(
1152 addr: Address,
1153 new_address: Option<Address>,
1154 storage: &HashMap<U256, U256>,
1155 chain: Chain,
1156 balance: Option<U256>,
1157) -> AccountUpdate {
1158 let mut slots = storage.clone();
1159 if let Some(new_address) = new_address {
1160 slots.insert(*IMPLEMENTATION_SLOT, U256::from_be_slice(new_address.as_slice()));
1161 }
1162
1163 AccountUpdate {
1164 address: addr,
1165 chain,
1166 slots,
1167 balance,
1168 code: Some(ERC20_PROXY_BYTECODE.to_vec()),
1169 change: ChangeType::Creation,
1170 }
1171}
1172
1173#[cfg(test)]
1174mock! {
1175 #[derive(Debug)]
1176 pub ProtocolSim {
1177 pub fn fee(&self) -> f64;
1178 pub fn spot_price(&self, base: &Token, quote: &Token) -> Result<f64, SimulationError>;
1179 pub fn get_amount_out(
1180 &self,
1181 amount_in: BigUint,
1182 token_in: &Token,
1183 token_out: &Token,
1184 ) -> Result<GetAmountOutResult, SimulationError>;
1185 pub fn get_limits(
1186 &self,
1187 sell_token: Bytes,
1188 buy_token: Bytes,
1189 ) -> Result<(BigUint, BigUint), SimulationError>;
1190 pub fn delta_transition(
1191 &mut self,
1192 delta: ProtocolStateDelta,
1193 tokens: &HashMap<Bytes, Token>,
1194 balances: &Balances,
1195 ) -> Result<(), TransitionError>;
1196 pub fn clone_box(&self) -> Box<dyn ProtocolSim>;
1197 pub fn eq(&self, other: &dyn ProtocolSim) -> bool;
1198 }
1199}
1200
1201#[cfg(test)]
1202crate::impl_non_serializable_protocol!(MockProtocolSim, "test protocol");
1203
1204#[cfg(test)]
1205impl ProtocolSim for MockProtocolSim {
1206 fn fee(&self) -> f64 {
1207 self.fee()
1208 }
1209
1210 fn spot_price(&self, base: &Token, quote: &Token) -> Result<f64, SimulationError> {
1211 self.spot_price(base, quote)
1212 }
1213
1214 fn get_amount_out(
1215 &self,
1216 amount_in: BigUint,
1217 token_in: &Token,
1218 token_out: &Token,
1219 ) -> Result<GetAmountOutResult, SimulationError> {
1220 self.get_amount_out(amount_in, token_in, token_out)
1221 }
1222
1223 fn get_limits(
1224 &self,
1225 sell_token: Bytes,
1226 buy_token: Bytes,
1227 ) -> Result<(BigUint, BigUint), SimulationError> {
1228 self.get_limits(sell_token, buy_token)
1229 }
1230
1231 fn delta_transition(
1232 &mut self,
1233 delta: ProtocolStateDelta,
1234 tokens: &HashMap<Bytes, Token>,
1235 balances: &Balances,
1236 ) -> Result<(), TransitionError> {
1237 self.delta_transition(delta, tokens, balances)
1238 }
1239
1240 fn clone_box(&self) -> Box<dyn ProtocolSim> {
1241 self.clone_box()
1242 }
1243
1244 fn as_any(&self) -> &dyn Any {
1245 panic!("MockProtocolSim does not support as_any")
1246 }
1247
1248 fn as_any_mut(&mut self) -> &mut dyn Any {
1249 panic!("MockProtocolSim does not support as_any_mut")
1250 }
1251
1252 fn eq(&self, other: &dyn ProtocolSim) -> bool {
1253 self.eq(other)
1254 }
1255
1256 fn typetag_name(&self) -> &'static str {
1257 unreachable!()
1258 }
1259
1260 fn typetag_deserialize(&self) {
1261 unreachable!()
1262 }
1263}
1264
1265#[cfg(test)]
1266mod tests {
1267 use std::str::FromStr;
1268
1269 use alloy::primitives::address;
1270 use mockall::predicate::*;
1271 use rstest::*;
1272 use tycho_client::feed::BlockHeader;
1273 use tycho_common::{models::Chain, Bytes};
1274
1275 use super::*;
1276 use crate::evm::protocol::{curve::CurveState, uniswap_v2::state::UniswapV2State};
1277
1278 #[test]
1279 fn curve_vm_adapter_registration_flagged_deprecated() {
1280 assert!(!is_deprecated_curve_registration::<CurveState>("vm:curve"));
1282 assert!(is_deprecated_curve_registration::<UniswapV2State>("vm:curve"));
1284 assert!(!is_deprecated_curve_registration::<UniswapV2State>("uniswap_v2"));
1286 }
1287
1288 async fn setup_decoder(set_tokens: bool) -> TychoStreamDecoder<BlockHeader> {
1289 let mut decoder = TychoStreamDecoder::new();
1290 decoder.register_decoder::<UniswapV2State>("uniswap_v2");
1291 if set_tokens {
1292 let tokens = [
1293 Bytes::from("0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2").lpad(20, 0),
1294 Bytes::from("0xdac17f958d2ee523a2206206994597c13d831ec7").lpad(20, 0),
1295 ]
1296 .iter()
1297 .map(|addr| {
1298 let addr_str = format!("{addr:x}");
1299 (
1300 addr.clone(),
1301 Token::new(addr, &addr_str, 18, 100, &[Some(100_000)], Chain::Ethereum, 100),
1302 )
1303 })
1304 .collect();
1305 decoder.set_tokens(tokens).await;
1306 }
1307 decoder
1308 }
1309
1310 fn load_test_msg(name: &str) -> FeedMessage<BlockHeader> {
1311 use std::{fs, path::Path};
1312
1313 use tycho_client::feed::dto;
1314 let project_root = env!("CARGO_MANIFEST_DIR");
1315 let asset_path = Path::new(project_root).join(format!("tests/assets/decoder/{name}.json"));
1316 let json_data = fs::read_to_string(asset_path).expect("Failed to read test asset");
1317 let feed_msg: dto::FeedMessage<BlockHeader> =
1318 serde_json::from_str(&json_data).expect("Failed to deserialize FeedMsg json!");
1319 FeedMessage::from(feed_msg)
1320 }
1321
1322 #[tokio::test]
1323 async fn test_decode() {
1324 let decoder = setup_decoder(true).await;
1325
1326 let msg = load_test_msg("uniswap_v2_snapshot");
1327 let res1 = decoder
1328 .decode(&msg)
1329 .await
1330 .expect("decode failure");
1331 let msg = load_test_msg("uniswap_v2_delta");
1332 let res2 = decoder
1333 .decode(&msg)
1334 .await
1335 .expect("decode failure");
1336
1337 assert_eq!(res1.states.len(), 1);
1338 assert_eq!(res2.states.len(), 1);
1339 assert_eq!(res1.sync_states.len(), 1);
1340 assert_eq!(res2.sync_states.len(), 1);
1341 }
1342
1343 #[tokio::test]
1344 async fn test_decode_token_creation_delta_with_existing_proxy() {
1345 let decoder = setup_decoder(true).await;
1346 let msg = load_test_msg("uniswap_v2_delta_token_creation");
1347
1348 decoder
1351 .decode(&msg)
1352 .await
1353 .expect("first decode (proxy creation) failed");
1354
1355 decoder
1359 .decode(&msg)
1360 .await
1361 .expect("decode of a token Creation delta with an existing proxy failed");
1362 }
1363
1364 #[tokio::test]
1365 async fn test_decode_component_missing_token() {
1366 let decoder = setup_decoder(false).await;
1367 let tokens = [Bytes::from("0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2").lpad(20, 0)]
1368 .iter()
1369 .map(|addr| {
1370 let addr_str = format!("{addr:x}");
1371 (
1372 addr.clone(),
1373 Token::new(addr, &addr_str, 18, 100, &[Some(100_000)], Chain::Ethereum, 100),
1374 )
1375 })
1376 .collect();
1377 decoder.set_tokens(tokens).await;
1378
1379 let msg = load_test_msg("uniswap_v2_snapshot");
1380 let res1 = decoder
1381 .decode(&msg)
1382 .await
1383 .expect("decode failure");
1384
1385 assert_eq!(res1.states.len(), 0);
1386 }
1387
1388 #[tokio::test]
1389 async fn test_decode_component_bad_id() {
1390 let decoder = setup_decoder(true).await;
1391 let msg = load_test_msg("uniswap_v2_snapshot_broken_id");
1392
1393 match decoder.decode(&msg).await {
1394 Err(StreamDecodeError::Fatal(msg)) => {
1395 assert_eq!(msg, "Component id mismatch");
1396 }
1397 Ok(_) => {
1398 panic!("Expected failures to be raised")
1399 }
1400 }
1401 }
1402
1403 #[rstest]
1404 #[case(true)]
1405 #[case(false)]
1406 #[tokio::test]
1407 async fn test_decode_component_bad_state(#[case] skip_failures: bool) {
1408 let mut decoder = setup_decoder(true).await;
1409 decoder.skip_state_decode_failures = skip_failures;
1410
1411 let msg = load_test_msg("uniswap_v2_snapshot_broken_state");
1412 match decoder.decode(&msg).await {
1413 Err(StreamDecodeError::Fatal(msg)) => {
1414 if !skip_failures {
1415 assert_eq!(msg, "Missing attributes reserve0");
1416 } else {
1417 panic!("Expected failures to be ignored. Err: {msg}")
1418 }
1419 }
1420 Ok(res) => {
1421 if !skip_failures {
1422 panic!("Expected failures to be raised")
1423 } else {
1424 assert_eq!(res.states.len(), 0);
1425 }
1426 }
1427 }
1428 }
1429
1430 #[tokio::test]
1431 async fn test_decode_updates_state_on_contract_change() {
1432 let decoder = setup_decoder(true).await;
1433
1434 let mut mock_state = MockProtocolSim::new();
1436
1437 mock_state
1438 .expect_clone_box()
1439 .times(1)
1440 .returning(|| {
1441 let mut cloned_mock_state = MockProtocolSim::new();
1442 cloned_mock_state
1444 .expect_delta_transition()
1445 .times(1)
1446 .returning(|_, _, _| Ok(()));
1447 cloned_mock_state
1448 .expect_clone_box()
1449 .times(1)
1450 .returning(|| Box::new(MockProtocolSim::new()));
1451 Box::new(cloned_mock_state)
1452 });
1453
1454 let pool_id =
1456 "0x93d199263632a4ef4bb438f1feb99e57b4b5f0bd0000000000000000000005c2".to_string();
1457 decoder
1458 .state
1459 .write()
1460 .await
1461 .states
1462 .insert(pool_id.clone(), Box::new(mock_state) as Box<dyn ProtocolSim>);
1463 decoder
1464 .state
1465 .write()
1466 .await
1467 .contracts_map
1468 .insert(
1469 Bytes::from("0xba12222222228d8ba445958a75a0704d566bf2c8").lpad(20, 0),
1470 HashSet::from([pool_id.clone()]),
1471 );
1472
1473 let msg = load_test_msg("balancer_v2_delta");
1475
1476 let _ = decoder
1478 .decode(&msg)
1479 .await
1480 .expect("decode failure");
1481
1482 }
1484
1485 #[test]
1486 fn test_generate_proxy_token_address() {
1487 let idx = 1;
1488 let generated_address =
1489 generate_proxy_token_address(idx).expect("proxy token address should be valid");
1490 assert_eq!(generated_address, address!("000000000000000000000000000000001badbabe"));
1491
1492 let idx = 123456;
1493 let generated_address =
1494 generate_proxy_token_address(idx).expect("proxy token address should be valid");
1495 assert_eq!(generated_address, address!("00000000000000000000000000001e240badbabe"));
1496 }
1497
1498 #[tokio::test(flavor = "multi_thread")]
1499 async fn test_euler_hook_low_pool_manager_balance() {
1500 let mut decoder = TychoStreamDecoder::new();
1501
1502 decoder.register_decoder_with_context::<crate::evm::protocol::uniswap_v4::state::UniswapV4State>(
1503 "uniswap_v4_hooks", DecoderContext::new().vm_traces(true)
1504 );
1505
1506 let weth = Bytes::from_str("0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2").unwrap();
1507 let teth = Bytes::from_str("0xd11c452fc99cf405034ee446803b6f6c1f6d5ed8").unwrap();
1508 let tokens = HashMap::from([
1509 (
1510 weth.clone(),
1511 Token::new(&weth, "WETH", 18, 100, &[Some(100_000)], Chain::Ethereum, 100),
1512 ),
1513 (
1514 teth.clone(),
1515 Token::new(&teth, "tETH", 18, 100, &[Some(100_000)], Chain::Ethereum, 100),
1516 ),
1517 ]);
1518
1519 decoder.set_tokens(tokens.clone()).await;
1520
1521 let msg = load_test_msg("euler_hook_snapshot");
1522 let res = decoder
1523 .decode(&msg)
1524 .await
1525 .expect("decode failure");
1526
1527 let pool_state = res
1528 .states
1529 .get("0xc70d7fbd7fcccdf726e02fed78548b40dc52502b097c7a1ee7d995f4d4396134")
1530 .expect("Couldn't find target pool");
1531 let amount_out = pool_state
1532 .get_amount_out(
1533 BigUint::from_str("1000000000000000000").unwrap(),
1534 tokens.get(&teth).unwrap(),
1535 tokens.get(&weth).unwrap(),
1536 )
1537 .expect("Get amount out failed");
1538
1539 assert_eq!(amount_out.amount, BigUint::from_str("1216190190361759119").unwrap());
1540 }
1541
1542 fn component_with_id(id: &str) -> ComponentWithState {
1543 use tycho_common::models::protocol::{ProtocolComponent, ProtocolComponentState};
1544
1545 ComponentWithState {
1546 state: ProtocolComponentState::new(id, HashMap::new(), HashMap::new()),
1547 component: ProtocolComponent { id: id.to_string(), ..Default::default() },
1548 component_tvl: None,
1549 entrypoints: Vec::new(),
1550 }
1551 }
1552
1553 fn rejects_a(component: &ComponentWithState) -> bool {
1554 component.component.id != "a"
1555 }
1556
1557 fn rejects_b(component: &ComponentWithState) -> bool {
1558 component.component.id != "b"
1559 }
1560
1561 #[test]
1562 fn test_admits_requires_every_registered_filter() {
1563 let mut decoder = TychoStreamDecoder::<BlockHeader>::new();
1566 decoder.register_filter("x", rejects_a);
1567 decoder.register_filter("x", rejects_b);
1568
1569 assert!(!decoder.admits("x", &component_with_id("a")));
1570 assert!(!decoder.admits("x", &component_with_id("b")));
1571 assert!(decoder.admits("x", &component_with_id("c")));
1572 assert!(decoder.admits("y", &component_with_id("a")));
1573 }
1574}