Liquidity.Party/tycho-protocol-sdk
2
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: Add state_registry.rs (#285)

Browse Source 
This it to make it easier for users to add new protocol states

#time 34m


#time 0m
This commit is contained in:
dianacarvalho1
2025-09-26 17:10:44 +01:00
committed by GitHub
parent 9a7e6a1cf7
commit ea10bfa99a
3 changed files with 54 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
protocol-testing/src/main.rs
View File

@@ -3,6 +3,7 @@ mod config;
mod encoding;
mod execution;
mod rpc;
mod state_registry;
mod test_runner;
mod traces;
mod tycho_rpc;

50
protocol-testing/src/state_registry.rs Normal file
View File

@@ -0,0 +1,50 @@
use tycho_simulation::{
evm::{
decoder::TychoStreamDecoder,
engine_db::tycho_db::PreCachedDB,
protocol::{
ekubo::state::EkuboState, pancakeswap_v2::state::PancakeswapV2State,
uniswap_v2::state::UniswapV2State, uniswap_v3::state::UniswapV3State,
vm::state::EVMPoolState,
},
},
protocol::models::DecoderContext,
tycho_client::feed::BlockHeader,
};
/// Register decoder based on protocol system. Defaults to EVMPoolState.
/// To add a new protocol, just add a case to the match statement.
pub fn register_decoder_for_protocol(
decoder: &mut TychoStreamDecoder<BlockHeader>,
protocol_system: &str,
decoder_context: DecoderContext,
) -> miette::Result<()> {
match protocol_system {
"uniswap_v2" | "sushiswap_v2" => {
decoder
.register_decoder_with_context::<UniswapV2State>(protocol_system, decoder_context);
}
"pancakeswap_v2" => {
decoder.register_decoder_with_context::<PancakeswapV2State>(
protocol_system,
decoder_context,
);
}
"uniswap_v3" | "pancakeswap_v3" => {
decoder
.register_decoder_with_context::<UniswapV3State>(protocol_system, decoder_context);
}
"ekubo_v2" => {
decoder.register_decoder_with_context::<EkuboState>(protocol_system, decoder_context);
}
// Default to EVMPoolState for all other protocols
_ => {
decoder.register_decoder_with_context::<EVMPoolState<PreCachedDB>>(
protocol_system,
decoder_context,
);
}
}
Ok(())
}

42
protocol-testing/src/test_runner.rs
View File

@@ -22,15 +22,7 @@ use postgres::{Client, Error, NoTls};
use tokio::runtime::Runtime;
use tracing::{debug, error, info, warn};
use tycho_simulation::{
evm::{
decoder::TychoStreamDecoder,
engine_db::tycho_db::PreCachedDB,
protocol::{
ekubo::state::EkuboState, pancakeswap_v2::state::PancakeswapV2State,
u256_num::bytes_to_u256, uniswap_v2::state::UniswapV2State,
uniswap_v3::state::UniswapV3State, vm::state::EVMPoolState,
},
},
evm::{decoder::TychoStreamDecoder, protocol::u256_num::bytes_to_u256},
protocol::models::{DecoderContext, Update},
tycho_client::feed::{
synchronizer::{ComponentWithState, Snapshot, StateSyncMessage},
@@ -50,6 +42,7 @@ use crate::{
encoding::encode_swap,
execution,
rpc::RPCProvider,
state_registry::register_decoder_for_protocol,
tycho_rpc::TychoClient,
tycho_runner::TychoRunner,
utils::build_spkg,
@@ -436,36 +429,7 @@ impl TestRunner {
if let Some(vm_adapter_path) = adapter_contract_path_str {
decoder_context = decoder_context.vm_adapter_path(vm_adapter_path);
}
match protocol_system {
"uniswap_v2" | "sushiswap_v2" => {
decoder.register_decoder_with_context::<UniswapV2State>(
protocol_system,
decoder_context,
);
}
"pancakeswap_v2" => {
decoder.register_decoder_with_context::<PancakeswapV2State>(
protocol_system,
decoder_context,
);
}
"uniswap_v3" | "pancakeswap_v3" => {
decoder.register_decoder_with_context::<UniswapV3State>(
protocol_system,
decoder_context,
);
}
"ekubo_v2" => {
decoder
.register_decoder_with_context::<EkuboState>(protocol_system, decoder_context);
}
_ => {
decoder.register_decoder_with_context::<EVMPoolState<PreCachedDB>>(
protocol_system,
decoder_context,
);
}
}
register_decoder_for_protocol(&mut decoder, protocol_system, decoder_context)?;
// Mock a stream message, with only a Snapshot and no deltas
let mut states: HashMap<String, ComponentWithState> = HashMap::new();