1use std::{
2 any::Any,
3 collections::{hash_map::Entry, HashMap, HashSet},
4 future::Future,
5 pin::Pin,
6 sync::Arc,
7};
8
9use alloy::primitives::{Address, U256};
10use thiserror::Error;
11use tokio::sync::{watch, RwLock, RwLockReadGuard};
12use tracing::{debug, error, info, warn};
13use tycho_client::feed::{synchronizer::ComponentWithState, BlockHeader, FeedMessage, HeaderLike};
14use tycho_common::{
15 dto::{ChangeType, ProtocolStateDelta},
16 models::{blockchain::BlockAggregatedChanges, token::Token, Chain},
17 simulation::protocol_sim::{Balances, BlockContext, ProtocolSim},
18 Bytes,
19};
20#[cfg(test)]
21use {
22 mockall::mock,
23 num_bigint::BigUint,
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 simulation::{BlockEnvOverrides, PendingOverrides},
39 tycho_models::{AccountUpdate, ResponseAccount},
40 },
41 protocol::{
42 errors::InvalidSnapshotError,
43 models::{DecoderContext, ProtocolComponent, TryFromWithBlock, Update},
44 },
45};
46
47#[derive(Error, Debug)]
48pub enum StreamDecodeError {
49 #[error("{0}")]
50 Fatal(String),
51}
52
53#[derive(Default)]
54struct DecoderState {
55 tokens: HashMap<Bytes, Token>,
56 states: HashMap<String, Box<dyn ProtocolSim>>,
57 components: HashMap<String, ProtocolComponent>,
58 contracts_map: HashMap<Bytes, HashSet<String>>,
60 proxy_token_addresses: HashMap<Address, Address>,
62 failed_components: HashSet<String>,
66 current_block_number: u64,
68}
69
70type DecodeFut =
71 Pin<Box<dyn Future<Output = Result<Box<dyn ProtocolSim>, InvalidSnapshotError>> + Send + Sync>>;
72type AccountBalances = HashMap<Bytes, HashMap<Bytes, Bytes>>;
73type RegistryFn<H> = dyn Fn(
74 ComponentWithState,
75 H,
76 AccountBalances,
77 Arc<RwLock<DecoderState>>,
78 Option<watch::Receiver<OverrideSnapshot>>,
79 ) -> DecodeFut
80 + Send
81 + Sync;
82type FilterFn = fn(&ComponentWithState) -> bool;
83
84pub struct TychoStreamDecoder<H>
97where
98 H: HeaderLike,
99{
100 state: Arc<RwLock<DecoderState>>,
101 skip_state_decode_failures: bool,
102 min_token_quality: u32,
103 registry: HashMap<String, Box<RegistryFn<H>>>,
104 inclusion_filters: HashMap<String, Vec<FilterFn>>,
105 protocols_skip_unknown_tokens: HashSet<String>,
107 override_providers: HashMap<String, Arc<dyn StateOverrideProvider>>,
110 block_time_secs: u64,
113 chain: Chain,
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(chain: Chain) -> Self {
137 Self {
138 state: Arc::new(RwLock::new(DecoderState::default())),
139 skip_state_decode_failures: false,
140 min_token_quality: 100,
141 registry: HashMap::new(),
142 inclusion_filters: HashMap::new(),
143 protocols_skip_unknown_tokens: HashSet::new(),
144 override_providers: HashMap::new(),
145 block_time_secs: chain.block_time_secs(),
146 chain,
147 }
148 }
149
150 fn execution_block(&self, header: &BlockHeader) -> BlockContext {
155 if header.partial_block_index.is_some() {
156 BlockContext::new(header.number, header.timestamp)
157 } else {
158 BlockContext::new(header.number + 1, header.timestamp + self.block_time_secs)
159 }
160 }
161
162 fn refresh_execution_block<C>(
171 updated_states: &mut HashMap<String, Box<dyn ProtocolSim>>,
172 stored_states: &mut HashMap<String, Box<dyn ProtocolSim>>,
173 failed_components: &HashSet<String>,
174 removed_components: &HashMap<String, C>,
175 execution_block: &BlockContext,
176 ) {
177 for state in updated_states.values_mut() {
178 state.apply_block(execution_block);
179 }
180 for (id, state) in stored_states.iter_mut() {
181 if failed_components.contains(id) ||
182 removed_components.contains_key(id) ||
183 updated_states.contains_key(id)
184 {
185 continue;
186 }
187 if state.apply_block(execution_block) {
188 updated_states.insert(id.clone(), state.clone_box());
189 }
190 }
191 }
192
193 pub fn set_override_provider(
199 &mut self,
200 protocol_system: String,
201 provider: Arc<dyn StateOverrideProvider>,
202 ) {
203 self.override_providers
204 .insert(protocol_system, provider);
205 }
206
207 pub async fn set_tokens(&self, tokens: HashMap<Bytes, Token>) {
212 let mut guard = self.state.write().await;
213 guard.tokens = tokens;
214 }
215
216 pub fn skip_state_decode_failures(&mut self, skip: bool) {
217 self.skip_state_decode_failures = skip;
218 }
219
220 pub fn min_token_quality(&mut self, quality: u32) {
226 self.min_token_quality = quality;
227 }
228
229 pub fn register_decoder_with_context<T>(&mut self, exchange: &str, mut context: DecoderContext)
244 where
245 T: ProtocolSim
246 + TryFromWithBlock<ComponentWithState, H, Error = InvalidSnapshotError>
247 + Send
248 + 'static,
249 {
250 if let Some(requested) = context
251 .chain
252 .filter(|requested| *requested != self.chain)
253 {
254 warn!(
255 exchange,
256 requested_chain = %requested,
257 decoder_chain = %self.chain,
258 "DecoderContext declares a different chain than the decoder; using the decoder's \
259 chain"
260 );
261 }
262 context.chain = Some(self.chain);
263
264 if is_deprecated_curve_registration::<T>(exchange) {
265 warn!(
266 registered_type = std::any::type_name::<T>(),
267 "Registering \"vm:curve\" with the generic VM adapter is deprecated; register the \
268 native `CurveState` decoder instead (`exchange::<CurveState>(\"vm:curve\", ...)`). \
269 The VM-adapter path still works but will be removed in a future release."
270 );
271 }
272 let decoder = Box::new(
273 move |component: ComponentWithState,
274 header: H,
275 account_balances: AccountBalances,
276 state: Arc<RwLock<DecoderState>>,
277 live_override: Option<watch::Receiver<OverrideSnapshot>>| {
278 let mut context = context.clone();
279 context.live_override = live_override;
280 Box::pin(async move {
281 let guard = state.read().await;
282 T::try_from_with_header(
283 component,
284 header,
285 &account_balances,
286 &guard.tokens,
287 &context,
288 )
289 .await
290 .map(|c| Box::new(c) as Box<dyn ProtocolSim>)
291 }) as DecodeFut
292 },
293 );
294 self.registry
295 .insert(exchange.to_string(), decoder);
296 }
297
298 pub fn register_decoder<T>(&mut self, exchange: &str)
310 where
311 T: ProtocolSim
312 + TryFromWithBlock<ComponentWithState, H, Error = InvalidSnapshotError>
313 + Send
314 + 'static,
315 {
316 let context = DecoderContext::new();
317 self.register_decoder_with_context::<T>(exchange, context);
318 }
319
320 pub fn register_decoder_dropping_unknown_tokens<T>(&mut self, exchange: &str)
324 where
325 T: ProtocolSim
326 + TryFromWithBlock<ComponentWithState, H, Error = InvalidSnapshotError>
327 + Send
328 + 'static,
329 {
330 self.register_decoder::<T>(exchange);
331 self.protocols_skip_unknown_tokens
332 .insert(exchange.to_string());
333 }
334
335 pub fn register_filter(&mut self, exchange: &str, predicate: FilterFn) {
354 self.inclusion_filters
355 .entry(exchange.to_string())
356 .or_default()
357 .push(predicate);
358 }
359
360 fn admits(&self, exchange: &str, snapshot: &ComponentWithState) -> bool {
363 let Some(predicates) = self.inclusion_filters.get(exchange) else { return true };
364 predicates
365 .iter()
366 .all(|predicate| predicate(snapshot))
367 }
368
369 pub async fn decode(&self, msg: &FeedMessage<H>) -> Result<Update, StreamDecodeError> {
372 let mut updated_states = HashMap::new();
374 let mut new_pairs = HashMap::new();
375 let mut removed_pairs = HashMap::new();
376 let mut contracts_map = HashMap::new();
377 let mut msg_failed_components = HashSet::new();
378
379 let header = msg
380 .state_msgs
381 .values()
382 .next()
383 .ok_or_else(|| StreamDecodeError::Fatal("Missing block!".into()))?
384 .header
385 .clone();
386
387 let block_number_or_timestamp = header
388 .clone()
389 .block_number_or_timestamp();
390 let current_block = header.clone().block();
391 let is_partial = current_block
392 .as_ref()
393 .map(|h| h.partial_block_index.is_some())
394 .unwrap_or(false);
395
396 for (protocol, protocol_msg) in msg.state_msgs.iter() {
397 let drops_unknown_tokens = self
398 .protocols_skip_unknown_tokens
399 .contains(protocol.as_str());
400 if let Some(deltas) = protocol_msg.deltas.as_ref() {
402 let mut state_guard = self.state.write().await;
403
404 let new_tokens = deltas
405 .new_tokens
406 .iter()
407 .filter(|(addr, t)| {
408 t.quality >= self.min_token_quality &&
409 !state_guard.tokens.contains_key(*addr)
410 })
411 .map(|(addr, t)| (addr.clone(), t.clone()))
412 .collect::<HashMap<Bytes, Token>>();
413
414 if !new_tokens.is_empty() {
415 debug!(n = new_tokens.len(), "NewTokens");
416 state_guard.tokens.extend(new_tokens);
417 }
418 }
419
420 {
422 let mut state_guard = self.state.write().await;
423 let removed_components: Vec<(String, ProtocolComponent)> = protocol_msg
424 .removed_components
425 .iter()
426 .map(|(id, comp)| {
427 if *id != comp.id {
428 error!(
429 "Component id mismatch in removed components {id} != {}",
430 comp.id
431 );
432 return Err(StreamDecodeError::Fatal("Component id mismatch".into()));
433 }
434
435 let tokens = comp
436 .tokens
437 .iter()
438 .flat_map(|addr| state_guard.tokens.get(addr).cloned())
439 .collect::<Vec<_>>();
440
441 if tokens.len() == comp.tokens.len() {
442 Ok(Some((
443 id.clone(),
444 ProtocolComponent::from_with_tokens(comp.clone(), tokens),
445 )))
446 } else {
447 Ok(None)
448 }
449 })
450 .collect::<Result<Vec<Option<(String, ProtocolComponent)>>, StreamDecodeError>>(
451 )?
452 .into_iter()
453 .flatten()
454 .collect();
455
456 for (id, component) in removed_components {
458 state_guard.components.remove(&id);
459 state_guard.states.remove(&id);
460 removed_pairs.insert(id, component);
461 }
462
463 info!(
465 "Processing {} contracts from snapshots",
466 protocol_msg
467 .snapshots
468 .get_vm_storage()
469 .len()
470 );
471
472 let mut proxy_token_accounts: HashMap<Address, AccountUpdate> = HashMap::new();
473 let mut storage_by_address: HashMap<Address, ResponseAccount> = HashMap::new();
474 for (key, value) in protocol_msg
475 .snapshots
476 .get_vm_storage()
477 .iter()
478 {
479 let account: ResponseAccount = value.clone().into();
480
481 if state_guard.tokens.contains_key(key) {
482 let original_address = account.address;
483 let (impl_addr, proxy_state) = match state_guard
492 .proxy_token_addresses
493 .get(&original_address)
494 {
495 Some(impl_addr) => {
496 let proxy_state = AccountUpdate::new(
503 original_address,
504 value.chain,
505 account.slots.clone(),
506 Some(account.native_balance),
507 None,
508 ChangeType::Update,
509 );
510 (*impl_addr, proxy_state)
511 }
512 None => {
513 let impl_addr = generate_proxy_token_address(
517 state_guard.proxy_token_addresses.len() as u32,
518 )?;
519 state_guard
520 .proxy_token_addresses
521 .insert(original_address, impl_addr);
522
523 let proxy_state = create_proxy_token_account(
525 original_address,
526 Some(impl_addr),
527 &account.slots,
528 value.chain,
529 Some(account.native_balance),
530 );
531
532 (impl_addr, proxy_state)
533 }
534 };
535
536 proxy_token_accounts.insert(original_address, proxy_state);
537
538 let impl_update = ResponseAccount {
540 address: impl_addr,
541 slots: HashMap::new(),
542 ..account.clone()
543 };
544 storage_by_address.insert(impl_addr, impl_update);
545 } else {
546 storage_by_address.insert(account.address, account);
548 }
549 }
550
551 let mut proxy_creates: Vec<AccountUpdate> = Vec::new();
555 let mut proxy_updates: HashMap<Address, AccountUpdate> = HashMap::new();
556 for (addr, update) in proxy_token_accounts {
557 if matches!(update.change, ChangeType::Creation) {
558 proxy_creates.push(update);
559 } else {
560 proxy_updates.insert(addr, update);
561 }
562 }
563
564 info!("Updating engine with {} contracts from snapshots", storage_by_address.len());
565 update_engine(
566 SHARED_TYCHO_DB.clone(),
567 header.clone().block(),
568 Some(storage_by_address),
569 proxy_updates,
570 )
571 .map_err(|e| StreamDecodeError::Fatal(e.to_string()))?;
572
573 if !proxy_creates.is_empty() {
577 SHARED_TYCHO_DB
578 .force_update_accounts(proxy_creates)
579 .map_err(|e| StreamDecodeError::Fatal(e.to_string()))?;
580 }
581 info!("Engine updated");
582 drop(state_guard);
583 }
584
585 let account_balances = protocol_msg
588 .clone()
589 .snapshots
590 .get_vm_storage()
591 .iter()
592 .filter_map(|(addr, acc)| {
593 if acc.token_balances.is_empty() {
594 return None;
595 }
596 let balances = acc
597 .token_balances
598 .iter()
599 .map(|(token_addr, ab)| (token_addr.clone(), ab.balance.clone()))
600 .collect::<HashMap<Bytes, Bytes>>();
601 Some((addr.clone(), balances))
602 })
603 .collect::<AccountBalances>();
604
605 let mut new_components = HashMap::new();
606 let mut count_token_skips = 0;
607 let mut components_to_store = HashMap::new();
608 {
609 let state_guard = self.state.read().await;
610
611 'snapshot_loop: for (id, snapshot) in protocol_msg
613 .snapshots
614 .get_states()
615 .clone()
616 {
617 if !self.admits(protocol.as_str(), &snapshot) {
619 continue;
620 }
621
622 let mut component_tokens = Vec::new();
624 let mut new_tokens_accounts = HashMap::new();
625 for token in snapshot.component.tokens.clone() {
626 match state_guard.tokens.get(&token) {
627 Some(token) => {
628 let token_address = match bytes_to_address(&token.address) {
631 Ok(addr) => addr,
632 Err(_) if drops_unknown_tokens => {
633 warn!(
634 "Token address could not be decoded {}, ignoring the token in component {:x?}",
635 token.address, id
636 );
637 continue;
638 }
639 Err(_) => {
640 count_token_skips += 1;
641 msg_failed_components.insert(id.clone());
642 warn!(
643 "Token address could not be decoded {}, ignoring pool {:x?}",
644 token.address, id
645 );
646 continue 'snapshot_loop;
647 }
648 };
649 component_tokens.push(token.clone());
650 if !state_guard
652 .proxy_token_addresses
653 .contains_key(&token_address)
654 {
655 new_tokens_accounts.insert(
656 token_address,
657 create_proxy_token_account(
658 token_address,
659 None,
660 &HashMap::new(),
661 snapshot.component.chain,
662 None,
663 ),
664 );
665 }
666 }
667 None if drops_unknown_tokens => {
668 warn!(
669 "Token not found {}, ignoring the token in component {:x?}",
670 token, id
671 );
672 }
673 None => {
674 count_token_skips += 1;
675 msg_failed_components.insert(id.clone());
676 debug!("Token not found {}, ignoring pool {:x?}", token, id);
677 continue 'snapshot_loop;
678 }
679 }
680 }
681 let component = ProtocolComponent::from_with_tokens(
682 snapshot.component.clone(),
683 component_tokens,
684 );
685
686 if !new_tokens_accounts.is_empty() {
688 update_engine(
689 SHARED_TYCHO_DB.clone(),
690 header.clone().block(),
691 None,
692 new_tokens_accounts,
693 )
694 .map_err(|e| StreamDecodeError::Fatal(e.to_string()))?;
695 }
696
697 if !component
700 .static_attributes
701 .contains_key("manual_updates")
702 {
703 for contract in &component.contract_ids {
704 contracts_map
705 .entry(contract.clone())
706 .or_insert_with(HashSet::new)
707 .insert(id.clone());
708 }
709 for (_, tracing) in snapshot.entrypoints.iter() {
712 for contract in tracing.accessed_slots.keys().cloned() {
713 contracts_map
714 .entry(contract)
715 .or_insert_with(HashSet::new)
716 .insert(id.clone());
717 }
718 }
719 }
720
721 new_pairs.insert(id.clone(), component.clone());
723
724 components_to_store.insert(id.clone(), component);
726
727 if let Some(state_decode_f) = self.registry.get(protocol.as_str()) {
729 let live_override = self
730 .override_providers
731 .get(protocol.as_str())
732 .and_then(|provider| provider.subscribe(protocol.as_str()));
733 match state_decode_f(
734 snapshot,
735 header.clone(),
736 account_balances.clone(),
737 self.state.clone(),
738 live_override,
739 )
740 .await
741 {
742 Ok(state) => {
743 new_components.insert(id.clone(), state);
744 }
745 Err(e) => {
746 if self.skip_state_decode_failures {
747 warn!(pool = id, error = %e, "StateDecodingFailure");
748 msg_failed_components.insert(id.clone());
749 continue 'snapshot_loop;
750 } else {
751 error!(pool = id, error = %e, "StateDecodingFailure");
752 return Err(StreamDecodeError::Fatal(format!("{e}")));
753 }
754 }
755 }
756 } else if self.skip_state_decode_failures {
757 warn!(pool = id, "MissingDecoderRegistration");
758 msg_failed_components.insert(id.clone());
759 continue 'snapshot_loop;
760 } else {
761 error!(pool = id, "MissingDecoderRegistration");
762 return Err(StreamDecodeError::Fatal(format!(
763 "Missing decoder registration for: {id}"
764 )));
765 }
766 }
767 }
768
769 if !components_to_store.is_empty() {
771 let mut state_guard = self.state.write().await;
772 for (id, component) in components_to_store {
773 state_guard
774 .components
775 .insert(id, component);
776 }
777 }
778
779 if !protocol_msg.snapshots.states.is_empty() {
780 info!("Decoded {} snapshots for protocol {protocol}", new_components.len());
781 }
782 if count_token_skips > 0 {
783 info!("Skipped {count_token_skips} pools due to missing tokens");
784 }
785
786 updated_states.extend(new_components);
788
789 if let Some(deltas) = protocol_msg.deltas.clone() {
791 let mut state_guard = self.state.write().await;
793
794 let mut account_update_by_address: HashMap<Address, AccountUpdate> = HashMap::new();
795 let mut new_proxy_accounts: Vec<AccountUpdate> = Vec::new();
797 for (key, value) in deltas.account_deltas.iter() {
798 let mut update: AccountUpdate = value.clone().into();
799
800 if update.code.is_none() && matches!(update.change, ChangeType::Creation) {
806 error!(
807 update = ?update,
808 "FaultyCreationDelta"
809 );
810 update.code = Some(vec![]);
811 }
812
813 if state_guard.tokens.contains_key(key) {
814 let original_address = update.address;
815 let impl_addr = match state_guard
822 .proxy_token_addresses
823 .get(&original_address)
824 {
825 Some(impl_addr) => {
826 let proxy_update = AccountUpdate {
831 code: None,
832 change: ChangeType::Update,
833 ..update.clone()
834 };
835 account_update_by_address.insert(original_address, proxy_update);
836
837 *impl_addr
838 }
839 None => {
840 let impl_addr = generate_proxy_token_address(
845 state_guard.proxy_token_addresses.len() as u32,
846 )?;
847 state_guard
848 .proxy_token_addresses
849 .insert(original_address, impl_addr);
850
851 let proxy_state = create_proxy_token_account(
856 original_address,
857 Some(impl_addr),
858 &update.slots,
859 update.chain,
860 update.balance,
861 );
862 new_proxy_accounts.push(proxy_state);
863
864 impl_addr
865 }
866 };
867
868 if update.code.is_some() {
870 let impl_update = AccountUpdate {
871 address: impl_addr,
872 slots: HashMap::new(),
873 ..update.clone()
874 };
875 account_update_by_address.insert(impl_addr, impl_update);
876 }
877 } else {
878 account_update_by_address.insert(update.address, update);
880 }
881 }
882 drop(state_guard);
883
884 let state_guard = self.state.read().await;
885 info!("Updating engine with {} contract deltas", deltas.account_deltas.len());
886 update_engine(
887 SHARED_TYCHO_DB.clone(),
888 header.clone().block(),
889 None,
890 account_update_by_address,
891 )
892 .map_err(|e| StreamDecodeError::Fatal(e.to_string()))?;
893
894 if !new_proxy_accounts.is_empty() {
897 SHARED_TYCHO_DB
898 .force_update_accounts(new_proxy_accounts)
899 .map_err(|e| StreamDecodeError::Fatal(e.to_string()))?;
900 }
901 info!("Engine updated");
902
903 let mut pools_to_update = HashSet::new();
905 for (account, _update) in deltas.account_deltas {
906 pools_to_update.extend(
908 contracts_map
909 .get(&account)
910 .cloned()
911 .unwrap_or_default(),
912 );
913 pools_to_update.extend(
915 state_guard
916 .contracts_map
917 .get(&account)
918 .cloned()
919 .unwrap_or_default(),
920 );
921 }
922
923 let all_balances = Balances {
925 component_balances: deltas
926 .component_balances
927 .iter()
928 .map(|(pool_id, bals)| {
929 let mut balances = HashMap::new();
930 for (t, b) in bals {
931 balances.insert(t.clone(), b.balance.clone());
932 }
933 pools_to_update.insert(pool_id.clone());
934 (pool_id.clone(), balances)
935 })
936 .collect(),
937 account_balances: deltas
938 .account_balances
939 .iter()
940 .map(|(account, bals)| {
941 let mut balances = HashMap::new();
942 for (t, b) in bals {
943 balances.insert(t.clone(), b.balance.clone());
944 }
945 pools_to_update.extend(
946 contracts_map
947 .get(account)
948 .cloned()
949 .unwrap_or_default(),
950 );
951 (account.clone(), balances)
952 })
953 .collect(),
954 };
955
956 for (id, update) in deltas.state_deltas {
958 let update_with_block = Self::add_block_info_to_delta(
960 ProtocolStateDelta::from(update),
961 current_block.clone(),
962 );
963 match Self::apply_update(
964 &id,
965 update_with_block,
966 &mut updated_states,
967 &state_guard,
968 &all_balances,
969 ) {
970 Ok(_) => {
971 pools_to_update.remove(&id);
972 }
973 Err(e) => {
974 if self.skip_state_decode_failures {
975 warn!(pool = id, error = %e, "Failed to apply state update, marking component as removed");
976 updated_states.remove(&id);
978 if let Some(component) = new_pairs.remove(&id) {
980 removed_pairs.insert(id.clone(), component);
981 } else if let Some(component) = state_guard.components.get(&id) {
982 removed_pairs.insert(id.clone(), component.clone());
983 } else {
984 warn!(pool = id, "Component not found in new_pairs or state, cannot add to removed_pairs");
987 }
988 pools_to_update.remove(&id);
989
990 msg_failed_components.insert(id.clone());
992 } else {
993 return Err(e);
994 }
995 }
996 }
997 }
998
999 for pool in pools_to_update {
1001 let default_delta_with_block = Self::add_block_info_to_delta(
1003 ProtocolStateDelta::default(),
1004 current_block.clone(),
1005 );
1006 match Self::apply_update(
1007 &pool,
1008 default_delta_with_block,
1009 &mut updated_states,
1010 &state_guard,
1011 &all_balances,
1012 ) {
1013 Ok(_) => {}
1014 Err(e) => {
1015 if self.skip_state_decode_failures {
1016 warn!(pool = pool, error = %e, "Failed to apply contract/balance update, marking component as removed");
1017 updated_states.remove(&pool);
1019 if let Some(component) = new_pairs.remove(&pool) {
1021 removed_pairs.insert(pool.clone(), component);
1022 } else if let Some(component) = state_guard.components.get(&pool) {
1023 removed_pairs.insert(pool.clone(), component.clone());
1024 } else {
1025 warn!(pool = pool, "Component not found in new_pairs or state, cannot add to removed_pairs");
1028 }
1029
1030 msg_failed_components.insert(pool.clone());
1032 } else {
1033 return Err(e);
1034 }
1035 }
1036 }
1037 }
1038 };
1039 }
1040
1041 let mut state_guard = self.state.write().await;
1043
1044 state_guard
1046 .failed_components
1047 .extend(msg_failed_components);
1048
1049 updated_states.retain(|id, _| {
1053 !state_guard
1054 .failed_components
1055 .contains(id)
1056 });
1057 new_pairs.retain(|id, _| {
1058 !state_guard
1059 .failed_components
1060 .contains(id)
1061 });
1062
1063 if let Some(header) = current_block.as_ref() {
1064 let execution_block = self.execution_block(header);
1065 let decoder_state = &mut *state_guard;
1066 Self::refresh_execution_block(
1067 &mut updated_states,
1068 &mut decoder_state.states,
1069 &decoder_state.failed_components,
1070 &removed_pairs,
1071 &execution_block,
1072 );
1073 }
1074
1075 state_guard
1076 .states
1077 .extend(updated_states.clone());
1078
1079 state_guard.current_block_number = block_number_or_timestamp;
1080
1081 for (id, component) in new_pairs.iter() {
1083 state_guard
1084 .components
1085 .insert(id.clone(), component.clone());
1086 }
1087
1088 for id in removed_pairs.keys() {
1090 state_guard.components.remove(id);
1091 }
1092
1093 for (key, values) in contracts_map {
1094 state_guard
1095 .contracts_map
1096 .entry(key)
1097 .or_insert_with(HashSet::new)
1098 .extend(values);
1099 }
1100
1101 Ok(Update::new(block_number_or_timestamp, updated_states, new_pairs)
1103 .set_is_partial(is_partial)
1104 .set_removed_pairs(removed_pairs)
1105 .set_sync_states(msg.sync_states.clone()))
1106 }
1107
1108 pub async fn apply_deltas_ephemeral(
1128 &self,
1129 pending_deltas: &HashMap<String, BlockAggregatedChanges>,
1130 header: H,
1131 ) -> Result<Update, StreamDecodeError> {
1132 let block_number_or_timestamp = header
1133 .clone()
1134 .block_number_or_timestamp();
1135 let current_block = header.block();
1136 let state_guard = self.state.read().await;
1137
1138 let mut updated_states: HashMap<String, Box<dyn ProtocolSim>> = HashMap::new();
1139
1140 for deltas in pending_deltas.values() {
1141 let all_balances = Balances {
1142 component_balances: deltas
1143 .component_balances
1144 .iter()
1145 .map(|(pool_id, bals)| {
1146 let balances = bals
1147 .iter()
1148 .map(|(t, b)| (t.clone(), b.balance.clone()))
1149 .collect();
1150 (pool_id.clone(), balances)
1151 })
1152 .collect(),
1153 account_balances: HashMap::new(),
1154 };
1155
1156 for (id, state_delta) in &deltas.state_deltas {
1157 let dto_delta = Self::add_block_info_to_delta(
1158 ProtocolStateDelta::from(state_delta.clone()),
1159 current_block.clone(),
1160 );
1161 if let Err(e) = Self::apply_update(
1162 id,
1163 dto_delta,
1164 &mut updated_states,
1165 &state_guard,
1166 &all_balances,
1167 ) {
1168 warn!(pool = id, error = %e, "EphemeralDeltaTransitionError");
1169 }
1170 }
1171
1172 if !deltas.account_deltas.is_empty() {
1173 Self::apply_pending_overrides(
1174 deltas,
1175 ¤t_block,
1176 &mut updated_states,
1177 &state_guard,
1178 &all_balances,
1179 )?;
1180 }
1181 }
1182
1183 if let Some(header) = current_block.as_ref() {
1188 let execution_block = BlockContext::new(header.number, header.timestamp);
1189 for state in updated_states.values_mut() {
1190 state.apply_block(&execution_block);
1191 }
1192 }
1193
1194 Ok(Update::new(block_number_or_timestamp, updated_states, HashMap::new()))
1195 }
1196
1197 fn apply_pending_overrides(
1207 deltas: &BlockAggregatedChanges,
1208 current_block: &Option<BlockHeader>,
1209 updated_states: &mut HashMap<String, Box<dyn ProtocolSim>>,
1210 state_guard: &RwLockReadGuard<'_, DecoderState>,
1211 all_balances: &Balances,
1212 ) -> Result<(), StreamDecodeError> {
1213 let mut linked_pools: HashSet<&String> = HashSet::new();
1214 for account in deltas.account_deltas.keys() {
1215 if let Some(pools) = state_guard.contracts_map.get(account) {
1216 linked_pools.extend(pools);
1217 }
1218 }
1219
1220 let block = current_block
1221 .as_ref()
1222 .map(|h| BlockEnvOverrides { number: Some(h.number), timestamp: Some(h.timestamp) });
1223 let overrides: Arc<dyn Any + Send + Sync> =
1224 Arc::new(PendingOverrides::from_account_deltas(&deltas.account_deltas, block));
1225 for id in linked_pools {
1226 if !updated_states.contains_key(id) {
1230 let empty_delta = Self::add_block_info_to_delta(
1231 ProtocolStateDelta { component_id: id.clone(), ..Default::default() },
1232 current_block.clone(),
1233 );
1234 if let Err(e) =
1235 Self::apply_update(id, empty_delta, updated_states, state_guard, all_balances)
1236 {
1237 warn!(pool = id, error = %e, "EphemeralDeltaTransitionError");
1238 continue;
1239 }
1240 }
1241 if let Some(state) = updated_states.get_mut(id) {
1242 state
1243 .set_pending_overrides(Arc::clone(&overrides))
1244 .map_err(|e| {
1245 StreamDecodeError::Fatal(format!("Pending overrides for pool {id}: {e}"))
1246 })?;
1247 }
1248 }
1249 Ok(())
1250 }
1251
1252 fn add_block_info_to_delta(
1254 mut delta: ProtocolStateDelta,
1255 block_header_opt: Option<BlockHeader>,
1256 ) -> ProtocolStateDelta {
1257 if let Some(header) = block_header_opt {
1258 delta.updated_attributes.insert(
1261 "block_number".to_string(),
1262 Bytes::from(header.number.to_be_bytes().to_vec()),
1263 );
1264 delta.updated_attributes.insert(
1265 "block_timestamp".to_string(),
1266 Bytes::from(header.timestamp.to_be_bytes().to_vec()),
1267 );
1268 }
1269 delta
1270 }
1271
1272 fn apply_update(
1273 id: &String,
1274 update: ProtocolStateDelta,
1275 updated_states: &mut HashMap<String, Box<dyn ProtocolSim>>,
1276 state_guard: &RwLockReadGuard<'_, DecoderState>,
1277 all_balances: &Balances,
1278 ) -> Result<(), StreamDecodeError> {
1279 match updated_states.entry(id.clone()) {
1280 Entry::Occupied(mut entry) => {
1281 let state: &mut Box<dyn ProtocolSim> = entry.get_mut();
1283 state
1284 .delta_transition(update, &state_guard.tokens, all_balances)
1285 .map_err(|e| {
1286 error!(pool = id, error = ?e, "DeltaTransitionError");
1287 StreamDecodeError::Fatal(format!("TransitionFailure: {e:?}"))
1288 })?;
1289 }
1290 Entry::Vacant(_) => {
1291 match state_guard.states.get(id) {
1292 Some(stored_state) => {
1295 let mut state = stored_state.clone();
1296 state
1297 .delta_transition(update, &state_guard.tokens, all_balances)
1298 .map_err(|e| {
1299 error!(pool = id, error = ?e, "DeltaTransitionError");
1300 StreamDecodeError::Fatal(format!("TransitionFailure: {e:?}"))
1301 })?;
1302 updated_states.insert(id.clone(), state);
1303 }
1304 None => debug!(pool = id, reason = "MissingState", "DeltaTransitionError"),
1305 }
1306 }
1307 }
1308 Ok(())
1309 }
1310}
1311
1312fn generate_proxy_token_address(idx: u32) -> Result<Address, StreamDecodeError> {
1314 let padded_idx = format!("{idx:x}");
1315 let padded_zeroes = "0".repeat(33 - padded_idx.len());
1316 let proxy_token_address = format!("{padded_zeroes}{padded_idx}BAdbaBe");
1317 let decoded = hex::decode(proxy_token_address).map_err(|e| {
1318 StreamDecodeError::Fatal(format!("Invalid proxy token address encoding: {e}"))
1319 })?;
1320
1321 const ADDRESS_LENGTH: usize = 20;
1322 if decoded.len() != ADDRESS_LENGTH {
1323 return Err(StreamDecodeError::Fatal(format!(
1324 "Invalid proxy token address length: expected {}, got {}",
1325 ADDRESS_LENGTH,
1326 decoded.len(),
1327 )));
1328 }
1329
1330 Ok(Address::from_slice(&decoded))
1331}
1332
1333fn create_proxy_token_account(
1338 addr: Address,
1339 new_address: Option<Address>,
1340 storage: &HashMap<U256, U256>,
1341 chain: Chain,
1342 balance: Option<U256>,
1343) -> AccountUpdate {
1344 let mut slots = storage.clone();
1345 if let Some(new_address) = new_address {
1346 slots.insert(*IMPLEMENTATION_SLOT, U256::from_be_slice(new_address.as_slice()));
1347 }
1348
1349 AccountUpdate {
1350 address: addr,
1351 chain,
1352 slots,
1353 balance,
1354 code: Some(ERC20_PROXY_BYTECODE.to_vec()),
1355 change: ChangeType::Creation,
1356 }
1357}
1358
1359#[cfg(test)]
1360mock! {
1361 #[derive(Debug)]
1362 pub ProtocolSim {
1363 pub fn fee(&self) -> f64;
1364 pub fn spot_price(&self, base: &Token, quote: &Token) -> Result<f64, SimulationError>;
1365 pub fn get_amount_out(
1366 &self,
1367 amount_in: BigUint,
1368 token_in: &Token,
1369 token_out: &Token,
1370 ) -> Result<GetAmountOutResult, SimulationError>;
1371 pub fn get_limits(
1372 &self,
1373 sell_token: Bytes,
1374 buy_token: Bytes,
1375 ) -> Result<(BigUint, BigUint), SimulationError>;
1376 pub fn delta_transition(
1377 &mut self,
1378 delta: ProtocolStateDelta,
1379 tokens: &HashMap<Bytes, Token>,
1380 balances: &Balances,
1381 ) -> Result<(), TransitionError>;
1382 pub fn clone_box(&self) -> Box<dyn ProtocolSim>;
1383 pub fn eq(&self, other: &dyn ProtocolSim) -> bool;
1384 }
1385}
1386
1387#[cfg(test)]
1388crate::impl_non_serializable_protocol!(MockProtocolSim, "test protocol");
1389
1390#[cfg(test)]
1391impl ProtocolSim for MockProtocolSim {
1392 fn fee(&self) -> f64 {
1393 self.fee()
1394 }
1395
1396 fn spot_price(&self, base: &Token, quote: &Token) -> Result<f64, SimulationError> {
1397 self.spot_price(base, quote)
1398 }
1399
1400 fn get_amount_out(
1401 &self,
1402 amount_in: BigUint,
1403 token_in: &Token,
1404 token_out: &Token,
1405 ) -> Result<GetAmountOutResult, SimulationError> {
1406 self.get_amount_out(amount_in, token_in, token_out)
1407 }
1408
1409 fn get_limits(
1410 &self,
1411 sell_token: Bytes,
1412 buy_token: Bytes,
1413 ) -> Result<(BigUint, BigUint), SimulationError> {
1414 self.get_limits(sell_token, buy_token)
1415 }
1416
1417 fn delta_transition(
1418 &mut self,
1419 delta: ProtocolStateDelta,
1420 tokens: &HashMap<Bytes, Token>,
1421 balances: &Balances,
1422 ) -> Result<(), TransitionError> {
1423 self.delta_transition(delta, tokens, balances)
1424 }
1425
1426 fn clone_box(&self) -> Box<dyn ProtocolSim> {
1427 self.clone_box()
1428 }
1429
1430 fn as_any(&self) -> &dyn Any {
1431 panic!("MockProtocolSim does not support as_any")
1432 }
1433
1434 fn as_any_mut(&mut self) -> &mut dyn Any {
1435 panic!("MockProtocolSim does not support as_any_mut")
1436 }
1437
1438 fn eq(&self, other: &dyn ProtocolSim) -> bool {
1439 self.eq(other)
1440 }
1441
1442 fn typetag_name(&self) -> &'static str {
1443 "MockProtocolSim"
1445 }
1446
1447 fn typetag_deserialize(&self) {
1448 unreachable!()
1449 }
1450}
1451
1452#[cfg(test)]
1453mod tests {
1454 use std::str::FromStr;
1455
1456 use alloy::primitives::address;
1457 use mockall::predicate::*;
1458 use rstest::*;
1459 use tycho_client::feed::BlockHeader;
1460 use tycho_common::{models::Chain, Bytes};
1461
1462 use super::*;
1463
1464 fn header_at(number: u64, timestamp: u64, partial: Option<u32>) -> BlockHeader {
1465 BlockHeader {
1466 hash: Bytes::from([0u8; 32]),
1467 number,
1468 parent_hash: Bytes::from([0u8; 32]),
1469 revert: false,
1470 timestamp,
1471 partial_block_index: partial,
1472 }
1473 }
1474
1475 fn block_sensitive_state() -> Box<dyn ProtocolSim> {
1477 use crate::evm::protocol::{
1478 aerodrome_slipstreams::state::AerodromeSlipstreamsState,
1479 utils::{
1480 slipstreams::{dynamic_fee_module::DynamicFeeConfig, observations::Observation},
1481 uniswap::{tick_list::TickInfo, tick_math::get_sqrt_ratio_at_tick},
1482 },
1483 };
1484
1485 Box::new(
1486 AerodromeSlipstreamsState::new(
1487 "block-sensitive".to_string(),
1488 0,
1489 1_000_000_000_000_000_000,
1490 get_sqrt_ratio_at_tick(0).unwrap(),
1491 0,
1492 1,
1493 3000,
1494 1,
1495 0,
1496 vec![TickInfo::new(-120, 0).unwrap(), TickInfo::new(120, 0).unwrap()],
1497 vec![Observation { block_timestamp: 500, initialized: true, ..Default::default() }],
1498 DynamicFeeConfig::new(2700, 30_000, 0, true, 750),
1499 )
1500 .expect("state should build")
1501 .with_position_assumption(crate::protocol::models::BlockPositionAssumption::First),
1504 )
1505 }
1506
1507 #[test]
1508 fn confirmed_header_targets_the_next_block() {
1509 let decoder = TychoStreamDecoder::<BlockHeader>::new(Chain::Base);
1510
1511 let execution_block = decoder.execution_block(&header_at(100, 1_000, None));
1512
1513 assert_eq!(execution_block.number(), 101);
1514 assert_eq!(execution_block.timestamp(), 1_000 + Chain::Base.block_time_secs());
1515 }
1516
1517 #[test]
1518 fn partial_header_targets_the_block_that_is_still_open() {
1519 let decoder = TychoStreamDecoder::<BlockHeader>::new(Chain::Ethereum);
1520
1521 let execution_block = decoder.execution_block(&header_at(100, 1_000, Some(3)));
1522
1523 assert_eq!(execution_block.number(), 100);
1524 assert_eq!(execution_block.timestamp(), 1_000);
1525 }
1526
1527 #[test]
1528 fn refresh_re_emits_a_state_whose_fee_flipped_without_a_delta() {
1529 let mut stored = HashMap::from([("block-sensitive".to_string(), {
1533 let mut state = block_sensitive_state();
1534 state.apply_block(&BlockContext::new(100, 500));
1535 state
1536 })]);
1537 let mut updated: HashMap<String, Box<dyn ProtocolSim>> = HashMap::new();
1538
1539 TychoStreamDecoder::<BlockHeader>::refresh_execution_block(
1540 &mut updated,
1541 &mut stored,
1542 &HashSet::new(),
1543 &HashMap::<String, ()>::new(),
1544 &BlockContext::new(101, 502),
1545 );
1546
1547 let emitted = updated
1548 .get("block-sensitive")
1549 .expect("a fee flip must be emitted even without a delta");
1550 assert_eq!(emitted.fee(), 750.0 / 1_000_000.0);
1551 assert_eq!(stored["block-sensitive"].fee(), 750.0 / 1_000_000.0);
1553 }
1554
1555 #[test]
1556 fn refresh_never_re_emits_failed_components() {
1557 let mut stored = HashMap::from([("zombie".to_string(), {
1560 let mut state = block_sensitive_state();
1561 state.apply_block(&BlockContext::new(100, 500));
1562 state
1563 })]);
1564 let failed = HashSet::from(["zombie".to_string()]);
1565 let mut updated: HashMap<String, Box<dyn ProtocolSim>> = HashMap::new();
1566
1567 TychoStreamDecoder::<BlockHeader>::refresh_execution_block(
1568 &mut updated,
1569 &mut stored,
1570 &failed,
1571 &HashMap::<String, ()>::new(),
1572 &BlockContext::new(101, 502),
1573 );
1574
1575 assert!(updated.is_empty());
1576 }
1577
1578 #[test]
1579 fn refresh_never_re_emits_removed_components() {
1580 let mut stored = HashMap::from([("gone".to_string(), {
1583 let mut state = block_sensitive_state();
1584 state.apply_block(&BlockContext::new(100, 500));
1585 state
1586 })]);
1587 let removed = HashMap::from([("gone".to_string(), ())]);
1588 let mut updated: HashMap<String, Box<dyn ProtocolSim>> = HashMap::new();
1589
1590 TychoStreamDecoder::<BlockHeader>::refresh_execution_block(
1591 &mut updated,
1592 &mut stored,
1593 &HashSet::new(),
1594 &removed,
1595 &BlockContext::new(101, 502),
1596 );
1597
1598 assert!(updated.is_empty());
1599 }
1600
1601 #[test]
1602 fn refresh_stays_quiet_when_no_fee_changed() {
1603 let mut stored: HashMap<String, Box<dyn ProtocolSim>> = HashMap::from([
1606 ("idle-sensitive".to_string(), {
1607 let mut state = block_sensitive_state();
1608 state.apply_block(&BlockContext::new(101, 502));
1609 state
1610 }),
1611 (
1612 "univ2".to_string(),
1613 Box::new(crate::evm::protocol::uniswap_v2::state::UniswapV2State::new(
1614 U256::from(1_000_000u64),
1615 U256::from(1_000_000u64),
1616 )) as Box<dyn ProtocolSim>,
1617 ),
1618 ]);
1619 let mut updated: HashMap<String, Box<dyn ProtocolSim>> = HashMap::new();
1620
1621 TychoStreamDecoder::<BlockHeader>::refresh_execution_block(
1622 &mut updated,
1623 &mut stored,
1624 &HashSet::new(),
1625 &HashMap::<String, ()>::new(),
1626 &BlockContext::new(102, 504),
1627 );
1628
1629 assert!(updated.is_empty());
1630 }
1631 use crate::evm::protocol::{curve::CurveState, uniswap_v2::state::UniswapV2State};
1632
1633 #[test]
1634 fn curve_vm_adapter_registration_flagged_deprecated() {
1635 assert!(!is_deprecated_curve_registration::<CurveState>("vm:curve"));
1637 assert!(is_deprecated_curve_registration::<UniswapV2State>("vm:curve"));
1639 assert!(!is_deprecated_curve_registration::<UniswapV2State>("uniswap_v2"));
1641 }
1642
1643 async fn setup_decoder(set_tokens: bool) -> TychoStreamDecoder<BlockHeader> {
1644 let mut decoder = TychoStreamDecoder::new(Chain::Ethereum);
1645 decoder.register_decoder::<UniswapV2State>("uniswap_v2");
1646 if set_tokens {
1647 let tokens = [
1648 Bytes::from("0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2").lpad(20, 0),
1649 Bytes::from("0xdac17f958d2ee523a2206206994597c13d831ec7").lpad(20, 0),
1650 ]
1651 .iter()
1652 .map(|addr| {
1653 let addr_str = format!("{addr:x}");
1654 (
1655 addr.clone(),
1656 Token::new(addr, &addr_str, 18, 100, &[Some(100_000)], Chain::Ethereum, 100),
1657 )
1658 })
1659 .collect();
1660 decoder.set_tokens(tokens).await;
1661 }
1662 decoder
1663 }
1664
1665 fn load_test_msg(name: &str) -> FeedMessage<BlockHeader> {
1666 use std::{fs, path::Path};
1667
1668 use tycho_client::feed::dto;
1669 let project_root = env!("CARGO_MANIFEST_DIR");
1670 let asset_path = Path::new(project_root).join(format!("tests/assets/decoder/{name}.json"));
1671 let json_data = fs::read_to_string(asset_path).expect("Failed to read test asset");
1672 let feed_msg: dto::FeedMessage<BlockHeader> =
1673 serde_json::from_str(&json_data).expect("Failed to deserialize FeedMsg json!");
1674 FeedMessage::from(feed_msg)
1675 }
1676
1677 #[tokio::test]
1678 async fn test_decode() {
1679 let decoder = setup_decoder(true).await;
1680
1681 let msg = load_test_msg("uniswap_v2_snapshot");
1682 let res1 = decoder
1683 .decode(&msg)
1684 .await
1685 .expect("decode failure");
1686 let msg = load_test_msg("uniswap_v2_delta");
1687 let res2 = decoder
1688 .decode(&msg)
1689 .await
1690 .expect("decode failure");
1691
1692 assert_eq!(res1.states.len(), 1);
1693 assert_eq!(res2.states.len(), 1);
1694 assert_eq!(res1.sync_states.len(), 1);
1695 assert_eq!(res2.sync_states.len(), 1);
1696 }
1697
1698 #[tokio::test]
1699 async fn test_decode_token_creation_delta_with_existing_proxy() {
1700 let decoder = setup_decoder(true).await;
1701 let msg = load_test_msg("uniswap_v2_delta_token_creation");
1702
1703 decoder
1706 .decode(&msg)
1707 .await
1708 .expect("first decode (proxy creation) failed");
1709
1710 decoder
1714 .decode(&msg)
1715 .await
1716 .expect("decode of a token Creation delta with an existing proxy failed");
1717 }
1718
1719 #[tokio::test]
1720 async fn test_decode_component_missing_token() {
1721 let decoder = setup_decoder(false).await;
1722 let tokens = [Bytes::from("0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2").lpad(20, 0)]
1723 .iter()
1724 .map(|addr| {
1725 let addr_str = format!("{addr:x}");
1726 (
1727 addr.clone(),
1728 Token::new(addr, &addr_str, 18, 100, &[Some(100_000)], Chain::Ethereum, 100),
1729 )
1730 })
1731 .collect();
1732 decoder.set_tokens(tokens).await;
1733
1734 let msg = load_test_msg("uniswap_v2_snapshot");
1735 let res1 = decoder
1736 .decode(&msg)
1737 .await
1738 .expect("decode failure");
1739
1740 assert_eq!(res1.states.len(), 0);
1741 }
1742
1743 #[rstest]
1744 #[case::unknown_token(false)]
1745 #[case::undecodable_address(true)]
1746 #[tokio::test]
1747 async fn test_dropping_decoder_keeps_component(#[case] bad_address: bool) {
1748 let weth = Bytes::from("0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2").lpad(20, 0);
1749 let other = match bad_address {
1751 true => Bytes::from([1u8; 21].to_vec()),
1752 false => Bytes::from("0xdac17f958d2ee523a2206206994597c13d831ec7").lpad(20, 0),
1753 };
1754 let mut decoder = TychoStreamDecoder::new(Chain::Ethereum);
1755 decoder.register_decoder_dropping_unknown_tokens::<UniswapV2State>("uniswap_v2");
1756 let token = |address: &Bytes, symbol| {
1757 let token =
1758 Token::new(address, symbol, 18, 100, &[Some(100_000)], Chain::Ethereum, 100);
1759 (address.clone(), token)
1760 };
1761 let mut known = HashMap::from([token(&weth, "WETH")]);
1762 if bad_address {
1763 known.extend([token(&other, "OTHER")]);
1764 }
1765 decoder.set_tokens(known).await;
1766 let mut msg = load_test_msg("uniswap_v2_snapshot");
1767 let (id, snapshot) = msg
1768 .state_msgs
1769 .get_mut("uniswap_v2")
1770 .unwrap()
1771 .snapshots
1772 .states
1773 .iter_mut()
1774 .next()
1775 .unwrap();
1776 snapshot.component.tokens = vec![weth, other];
1777 let id = id.clone();
1778
1779 let update = decoder.decode(&msg).await.unwrap();
1780
1781 let symbols: Vec<_> = update.new_pairs[&id]
1782 .tokens
1783 .iter()
1784 .map(|token| token.symbol.as_str())
1785 .collect();
1786 assert_eq!(symbols, ["WETH"]);
1787 assert_eq!(update.states.len(), 1);
1788 }
1789
1790 #[tokio::test]
1791 async fn test_decode_component_bad_id() {
1792 let decoder = setup_decoder(true).await;
1793 let msg = load_test_msg("uniswap_v2_snapshot_broken_id");
1794
1795 match decoder.decode(&msg).await {
1796 Err(StreamDecodeError::Fatal(msg)) => {
1797 assert_eq!(msg, "Component id mismatch");
1798 }
1799 Ok(_) => {
1800 panic!("Expected failures to be raised")
1801 }
1802 }
1803 }
1804
1805 #[rstest]
1806 #[case(true)]
1807 #[case(false)]
1808 #[tokio::test]
1809 async fn test_decode_component_bad_state(#[case] skip_failures: bool) {
1810 let mut decoder = setup_decoder(true).await;
1811 decoder.skip_state_decode_failures = skip_failures;
1812
1813 let msg = load_test_msg("uniswap_v2_snapshot_broken_state");
1814 match decoder.decode(&msg).await {
1815 Err(StreamDecodeError::Fatal(msg)) => {
1816 if !skip_failures {
1817 assert_eq!(msg, "Missing attributes reserve0");
1818 } else {
1819 panic!("Expected failures to be ignored. Err: {msg}")
1820 }
1821 }
1822 Ok(res) => {
1823 if !skip_failures {
1824 panic!("Expected failures to be raised")
1825 } else {
1826 assert_eq!(res.states.len(), 0);
1827 }
1828 }
1829 }
1830 }
1831
1832 #[tokio::test]
1833 async fn test_decode_updates_state_on_contract_change() {
1834 let decoder = setup_decoder(true).await;
1835
1836 let mut mock_state = MockProtocolSim::new();
1838
1839 mock_state
1840 .expect_clone_box()
1841 .times(1)
1842 .returning(|| {
1843 let mut cloned_mock_state = MockProtocolSim::new();
1844 cloned_mock_state
1846 .expect_delta_transition()
1847 .times(1)
1848 .returning(|_, _, _| Ok(()));
1849 cloned_mock_state
1850 .expect_clone_box()
1851 .times(1)
1852 .returning(|| Box::new(MockProtocolSim::new()));
1853 Box::new(cloned_mock_state)
1854 });
1855
1856 let pool_id =
1858 "0x93d199263632a4ef4bb438f1feb99e57b4b5f0bd0000000000000000000005c2".to_string();
1859 decoder
1860 .state
1861 .write()
1862 .await
1863 .states
1864 .insert(pool_id.clone(), Box::new(mock_state) as Box<dyn ProtocolSim>);
1865 decoder
1866 .state
1867 .write()
1868 .await
1869 .contracts_map
1870 .insert(
1871 Bytes::from("0xba12222222228d8ba445958a75a0704d566bf2c8").lpad(20, 0),
1872 HashSet::from([pool_id.clone()]),
1873 );
1874
1875 let msg = load_test_msg("balancer_v2_delta");
1877
1878 let _ = decoder
1880 .decode(&msg)
1881 .await
1882 .expect("decode failure");
1883
1884 }
1886
1887 #[tokio::test]
1890 async fn test_apply_deltas_ephemeral_sets_pending_overrides_on_the_clone_only() {
1891 use tycho_common::models::{
1892 blockchain::BlockAggregatedChanges, contract::AccountDelta,
1893 protocol::ProtocolComponentStateDelta, ChangeType as ModelChangeType,
1894 };
1895
1896 use crate::evm::protocol::uniswap_v4::state::{UniswapV4Fees, UniswapV4State};
1897
1898 let decoder = TychoStreamDecoder::<BlockHeader>::new(Chain::Ethereum);
1899 let pool_id = "0xhooked".to_string();
1900 let quiet_pool_id = "0xhooked-no-log".to_string();
1901 let pool = UniswapV4State::new(
1902 1000,
1903 U256::from(1u8) << 96,
1904 UniswapV4Fees::new(0, 0, 3000),
1905 0,
1906 60,
1907 vec![],
1908 )
1909 .unwrap();
1910 let hook = Bytes::from(vec![7u8; 20]);
1911 {
1912 let mut state = decoder.state.write().await;
1913 state
1914 .states
1915 .insert(pool_id.clone(), Box::new(pool.clone()));
1916 state
1917 .states
1918 .insert(quiet_pool_id.clone(), Box::new(pool));
1919 state
1920 .contracts_map
1921 .insert(hook.clone(), HashSet::from([pool_id.clone(), quiet_pool_id.clone()]));
1922 }
1923
1924 let deltas = BlockAggregatedChanges {
1925 extractor: "uniswap_v4_hooks".to_string(),
1926 state_deltas: HashMap::from([(
1927 pool_id.clone(),
1928 ProtocolComponentStateDelta {
1929 component_id: pool_id.clone(),
1930 updated_attributes: HashMap::from([(
1931 "liquidity".to_string(),
1932 Bytes::from(2000_u64.to_be_bytes().to_vec()),
1933 )]),
1934 deleted_attributes: HashSet::new(),
1935 created_attributes: HashSet::new(),
1936 },
1937 )]),
1938 account_deltas: HashMap::from([(
1939 hook.clone(),
1940 AccountDelta::new(
1941 Chain::Ethereum,
1942 hook.clone(),
1943 HashMap::from([(Bytes::from(vec![0u8]), Some(Bytes::from(vec![9u8])))]),
1944 None,
1945 None,
1946 ModelChangeType::Update,
1947 ),
1948 )]),
1949 ..Default::default()
1950 };
1951 let header = BlockHeader { number: 10, timestamp: 20, ..Default::default() };
1952
1953 let update = decoder
1954 .apply_deltas_ephemeral(
1955 &HashMap::from([("uniswap_v4_hooks".to_string(), deltas)]),
1956 header,
1957 )
1958 .await
1959 .unwrap();
1960
1961 let clone = update.states[&pool_id]
1962 .as_any()
1963 .downcast_ref::<UniswapV4State>()
1964 .unwrap();
1965 let pending = clone
1966 .pending_overrides()
1967 .expect("the clone carries the pending block's overrides");
1968 assert_eq!(
1969 pending.block,
1970 Some(BlockEnvOverrides { number: Some(10), timestamp: Some(20) })
1971 );
1972 assert_eq!(
1973 pending.storage.as_ref().unwrap()[&Address::from_slice(&hook)][&U256::ZERO],
1974 U256::from(9)
1975 );
1976 let quiet = update.states[&quiet_pool_id]
1977 .as_any()
1978 .downcast_ref::<UniswapV4State>()
1979 .expect("a pool whose hook was written is cloned without a delta of its own");
1980 assert!(quiet.pending_overrides().is_some());
1981
1982 let stored = decoder.state.read().await;
1983 for id in [&pool_id, &quiet_pool_id] {
1984 let stored = stored.states[id]
1985 .as_any()
1986 .downcast_ref::<UniswapV4State>()
1987 .unwrap();
1988 assert!(stored.pending_overrides().is_none(), "confirmed state is never overlaid");
1989 }
1990 }
1991
1992 #[test]
1993 fn test_generate_proxy_token_address() {
1994 let idx = 1;
1995 let generated_address =
1996 generate_proxy_token_address(idx).expect("proxy token address should be valid");
1997 assert_eq!(generated_address, address!("000000000000000000000000000000001badbabe"));
1998
1999 let idx = 123456;
2000 let generated_address =
2001 generate_proxy_token_address(idx).expect("proxy token address should be valid");
2002 assert_eq!(generated_address, address!("00000000000000000000000000001e240badbabe"));
2003 }
2004
2005 fn hooked_v4_msg() -> FeedMessage<BlockHeader> {
2011 use tycho_client::feed::synchronizer::{Snapshot, StateSyncMessage};
2012 use tycho_common::models::protocol::{ProtocolComponent, ProtocolComponentState};
2013
2014 let id = "0x00000000000000000000000000000000000000000000000000000000000000c4";
2015 let static_attributes = HashMap::from([
2016 ("key_lp_fee".to_string(), Bytes::from(500_i32.to_be_bytes().to_vec())),
2017 ("tick_spacing".to_string(), Bytes::from(60_i32.to_be_bytes().to_vec())),
2018 (
2019 "hooks".to_string(),
2020 Bytes::from_str("0x00000000000000000000000000000000000000c4").unwrap(),
2021 ),
2022 ]);
2023 let attributes = HashMap::from([
2024 ("liquidity".to_string(), Bytes::from(100_u64.to_be_bytes().to_vec())),
2025 ("tick".to_string(), Bytes::from(300_i32.to_be_bytes().to_vec())),
2026 (
2027 "sqrt_price_x96".to_string(),
2028 Bytes::from(
2029 79228162514264337593543950336_u128
2030 .to_be_bytes()
2031 .to_vec(),
2032 ),
2033 ),
2034 ("protocol_fees/zero2one".to_string(), Bytes::from(0_u32.to_be_bytes().to_vec())),
2035 ("protocol_fees/one2zero".to_string(), Bytes::from(0_u32.to_be_bytes().to_vec())),
2036 ("ticks/60/net_liquidity".to_string(), Bytes::from(400_i128.to_be_bytes().to_vec())),
2037 ]);
2038
2039 let snapshot = ComponentWithState {
2040 state: ProtocolComponentState::new(id, attributes, HashMap::new()),
2041 component: ProtocolComponent {
2042 id: id.to_string(),
2043 static_attributes,
2044 ..Default::default()
2045 },
2046 component_tvl: None,
2047 entrypoints: Vec::new(),
2048 };
2049
2050 FeedMessage {
2051 state_msgs: HashMap::from([(
2052 "uniswap_v4_hooks".to_string(),
2053 StateSyncMessage {
2054 header: header_at(1, 1_000, None),
2055 snapshots: Snapshot {
2056 states: HashMap::from([(id.to_string(), snapshot)]),
2057 vm_storage: HashMap::new(),
2058 },
2059 deltas: None,
2060 removed_components: HashMap::new(),
2061 },
2062 )]),
2063 sync_states: HashMap::new(),
2064 }
2065 }
2066
2067 fn hooked_v4_decoder(chain: Chain, context: DecoderContext) -> TychoStreamDecoder<BlockHeader> {
2069 let mut decoder = TychoStreamDecoder::new(chain);
2070 decoder.register_decoder_with_context::<crate::evm::protocol::uniswap_v4::state::UniswapV4State>(
2071 "uniswap_v4_hooks", context
2072 );
2073 decoder
2074 }
2075
2076 #[tokio::test]
2077 async fn test_unregistered_hook_is_rejected_off_the_generic_vm_chains() {
2078 let decoder = hooked_v4_decoder(Chain::Robinhood, DecoderContext::new());
2079
2080 let Err(StreamDecodeError::Fatal(message)) = decoder.decode(&hooked_v4_msg()).await else {
2081 panic!("a hooked pool on robinhood must not decode");
2082 };
2083
2084 assert!(message.contains("unsupported uniswap v4 hook"), "{message}");
2085 assert!(message.contains("robinhood"), "{message}");
2086 }
2087
2088 #[tokio::test]
2089 async fn test_decoder_chain_overrides_the_registered_context_chain() {
2090 let decoder =
2091 hooked_v4_decoder(Chain::Robinhood, DecoderContext::new().chain(Chain::Ethereum));
2092
2093 let Err(StreamDecodeError::Fatal(message)) = decoder.decode(&hooked_v4_msg()).await else {
2094 panic!("the decoder's chain must win over the context's");
2095 };
2096
2097 assert!(message.contains("robinhood"), "{message}");
2098 }
2099
2100 #[tokio::test]
2101 async fn test_unregistered_hook_takes_the_generic_path_on_ethereum() {
2102 let decoder = hooked_v4_decoder(Chain::Ethereum, DecoderContext::new());
2103
2104 let Err(StreamDecodeError::Fatal(message)) = decoder.decode(&hooked_v4_msg()).await else {
2105 panic!("the generic VM creator has no balance_owner attribute to work from");
2106 };
2107
2108 assert!(!message.contains("unsupported uniswap v4 hook"), "{message}");
2109 assert!(message.contains("balance_owner"), "{message}");
2110 }
2111
2112 #[tokio::test(flavor = "multi_thread")]
2113 async fn test_euler_hook_low_pool_manager_balance() {
2114 let mut decoder = TychoStreamDecoder::new(Chain::Ethereum);
2115
2116 decoder.register_decoder_with_context::<crate::evm::protocol::uniswap_v4::state::UniswapV4State>(
2117 "uniswap_v4_hooks", DecoderContext::new().vm_traces(true)
2118 );
2119
2120 let weth = Bytes::from_str("0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2").unwrap();
2121 let teth = Bytes::from_str("0xd11c452fc99cf405034ee446803b6f6c1f6d5ed8").unwrap();
2122 let tokens = HashMap::from([
2123 (
2124 weth.clone(),
2125 Token::new(&weth, "WETH", 18, 100, &[Some(100_000)], Chain::Ethereum, 100),
2126 ),
2127 (
2128 teth.clone(),
2129 Token::new(&teth, "tETH", 18, 100, &[Some(100_000)], Chain::Ethereum, 100),
2130 ),
2131 ]);
2132
2133 decoder.set_tokens(tokens.clone()).await;
2134
2135 let msg = load_test_msg("euler_hook_snapshot");
2136 let res = decoder
2137 .decode(&msg)
2138 .await
2139 .expect("decode failure");
2140
2141 let pool_state = res
2142 .states
2143 .get("0xc70d7fbd7fcccdf726e02fed78548b40dc52502b097c7a1ee7d995f4d4396134")
2144 .expect("Couldn't find target pool");
2145 let amount_out = pool_state
2146 .get_amount_out(
2147 BigUint::from_str("1000000000000000000").unwrap(),
2148 tokens.get(&teth).unwrap(),
2149 tokens.get(&weth).unwrap(),
2150 )
2151 .expect("Get amount out failed");
2152
2153 assert_eq!(amount_out.amount, BigUint::from_str("1216190190361759119").unwrap());
2154 }
2155
2156 fn pons_feed_message(snapshot: ComponentWithState) -> FeedMessage<BlockHeader> {
2158 use tycho_client::feed::synchronizer::{Snapshot, StateSyncMessage};
2159
2160 use crate::evm::protocol::uniswap_v4::pons_fixture;
2161
2162 FeedMessage {
2163 state_msgs: HashMap::from([(
2164 "uniswap_v4_hooks".to_string(),
2165 StateSyncMessage {
2166 header: pons_fixture::header(),
2167 snapshots: Snapshot {
2168 states: HashMap::from([(pons_fixture::POOL_ID.to_string(), snapshot)]),
2169 vm_storage: HashMap::new(),
2170 },
2171 deltas: None,
2172 removed_components: HashMap::new(),
2173 },
2174 )]),
2175 sync_states: HashMap::new(),
2176 }
2177 }
2178
2179 async fn pons_stream_decoder() -> TychoStreamDecoder<BlockHeader> {
2181 use crate::evm::protocol::uniswap_v4::{
2182 hooks::hook_handler_creator::initialize_hook_handlers, pons_fixture,
2183 state::UniswapV4State,
2184 };
2185
2186 initialize_hook_handlers().expect("hook handler registration should succeed");
2187 let mut decoder = TychoStreamDecoder::new(Chain::Robinhood);
2188 decoder.register_decoder::<UniswapV4State>("uniswap_v4_hooks");
2189 decoder
2190 .set_tokens(pons_fixture::tokens())
2191 .await;
2192 decoder
2193 }
2194
2195 #[tokio::test]
2199 async fn test_pons_pool_decodes_through_the_stream_decoder_on_robinhood() {
2200 use crate::evm::protocol::uniswap_v4::{
2201 hooks::pons_v2::hook_handler::PonsV2HookHandler, pons_fixture, state::UniswapV4State,
2202 };
2203
2204 let decoder = pons_stream_decoder().await;
2205
2206 let result = decoder
2207 .decode(&pons_feed_message(pons_fixture::snapshot()))
2208 .await
2209 .expect("decode failure");
2210
2211 let emitted = result
2212 .states
2213 .get(pons_fixture::POOL_ID)
2214 .expect("the Pons pool must reach consumers");
2215 let pool = emitted
2216 .as_any()
2217 .downcast_ref::<UniswapV4State>()
2218 .expect("a uniswap_v4_hooks component decodes into a UniswapV4State");
2219 let handler = pool
2220 .hook
2221 .as_ref()
2222 .expect("a Pons pool must carry a hook handler");
2223 let pons = handler
2224 .as_any()
2225 .downcast_ref::<PonsV2HookHandler>()
2226 .expect("the registry must hand back the native Pons handler");
2227 assert_eq!(u32::from(pons.hook_fee_bps()), pons_fixture::HOOK_FEE_BPS);
2228 assert_eq!(u32::from(pons.creator_tax_bps()), pons_fixture::CREATOR_TAX_BPS);
2229
2230 let core = pons_fixture::decode_on(pons_fixture::hookless_snapshot(), Chain::Robinhood)
2231 .await
2232 .expect("the same pool with no hook must decode too");
2233 let (xlg, nvda) = (pons_fixture::xlg(), pons_fixture::nvda());
2234 let amount_in = BigUint::from(10u64).pow(18);
2235 for (token_in, token_out) in [(&xlg, &nvda), (&nvda, &xlg)] {
2236 let core_out = core
2237 .get_amount_out(amount_in.clone(), token_in, token_out)
2238 .expect("the reference pool quotes one whole token")
2239 .amount;
2240 let hooked_out = emitted
2241 .get_amount_out(amount_in.clone(), token_in, token_out)
2242 .expect("the emitted pool quotes one whole token")
2243 .amount;
2244
2245 let expected = pons_fixture::net_of_hook_take(&core_out);
2246 assert!(expected < core_out, "the hook must take something out of {core_out}");
2247 assert_eq!(hooked_out, expected, "{} -> {}", token_in.symbol, token_out.symbol);
2248 }
2249 }
2250
2251 #[tokio::test]
2254 async fn test_unknown_hook_component_is_not_emitted_on_robinhood() {
2255 use crate::evm::protocol::uniswap_v4::pons_fixture;
2256
2257 let mut decoder = pons_stream_decoder().await;
2258 decoder.skip_state_decode_failures(true);
2259
2260 let result = decoder
2261 .decode(&pons_feed_message(pons_fixture::unknown_hook_snapshot()))
2262 .await
2263 .expect("a skipped decode failure is not fatal");
2264
2265 assert!(
2266 !result
2267 .states
2268 .contains_key(pons_fixture::POOL_ID),
2269 "a pool whose hook cannot be modelled must not be quoted"
2270 );
2271 assert!(result.states.is_empty());
2272 }
2273
2274 fn component_with_id(id: &str) -> ComponentWithState {
2275 use tycho_common::models::protocol::{ProtocolComponent, ProtocolComponentState};
2276
2277 ComponentWithState {
2278 state: ProtocolComponentState::new(id, HashMap::new(), HashMap::new()),
2279 component: ProtocolComponent { id: id.to_string(), ..Default::default() },
2280 component_tvl: None,
2281 entrypoints: Vec::new(),
2282 }
2283 }
2284
2285 fn rejects_a(component: &ComponentWithState) -> bool {
2286 component.component.id != "a"
2287 }
2288
2289 fn rejects_b(component: &ComponentWithState) -> bool {
2290 component.component.id != "b"
2291 }
2292
2293 #[test]
2294 fn test_admits_requires_every_registered_filter() {
2295 let mut decoder = TychoStreamDecoder::<BlockHeader>::new(Chain::Ethereum);
2298 decoder.register_filter("x", rejects_a);
2299 decoder.register_filter("x", rejects_b);
2300
2301 assert!(!decoder.admits("x", &component_with_id("a")));
2302 assert!(!decoder.admits("x", &component_with_id("b")));
2303 assert!(decoder.admits("x", &component_with_id("c")));
2304 assert!(decoder.admits("y", &component_with_id("a")));
2305 }
2306}