From ea10bfa99a6524eec4725f2b37aab31de7299d63 Mon Sep 17 00:00:00 2001
From: dianacarvalho1 <diana@propellerheads.xyz>
Date: Fri, 26 Sep 2025 17:10:44 +0100
Subject: [PATCH] feat: Add state_registry.rs (#285)

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

#time 34m


#time 0m
---
 protocol-testing/src/main.rs           |  1 +
 protocol-testing/src/state_registry.rs | 50 ++++++++++++++++++++++++++
 protocol-testing/src/test_runner.rs    | 42 ++--------------------
 3 files changed, 54 insertions(+), 39 deletions(-)
 create mode 100644 protocol-testing/src/state_registry.rs

diff --git a/protocol-testing/src/main.rs b/protocol-testing/src/main.rs
index 1ae8748..50a3293 100644
--- a/protocol-testing/src/main.rs
+++ b/protocol-testing/src/main.rs
@@ -3,6 +3,7 @@ mod config;
 mod encoding;
 mod execution;
 mod rpc;
+mod state_registry;
 mod test_runner;
 mod traces;
 mod tycho_rpc;
diff --git a/protocol-testing/src/state_registry.rs b/protocol-testing/src/state_registry.rs
new file mode 100644
index 0000000..2fed5b5
--- /dev/null
+++ b/protocol-testing/src/state_registry.rs
@@ -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(())
+}
diff --git a/protocol-testing/src/test_runner.rs b/protocol-testing/src/test_runner.rs
index 12a09b1..faec808 100644
--- a/protocol-testing/src/test_runner.rs
+++ b/protocol-testing/src/test_runner.rs
@@ -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();