Merge pull request #52 from propeller-heads/encoding/tnl/configurable-chain

feat: Get native/wrapped addresses for chain

Cargo.lock

@@ -4242,6 +4242,7 @@ dependencies = [
  "lazy_static",
  "num-bigint",
  "num-traits",
+ "once_cell",
  "rstest",
  "serde",
  "serde_json",

Cargo.toml

@@ -19,6 +19,7 @@ alloy = { version = "0.9.2", features = ["providers", "rpc-types-eth", "eip712",
 alloy-sol-types = { version = "0.8.14", optional = true }
 alloy-primitives = { version = "0.8.9", optional = true }
 tycho-core = { git = "https://github.com/propeller-heads/tycho-indexer.git", package = "tycho-core", tag = "0.46.0" }
+once_cell = "1.20.2"
 
 [dev-dependencies]
 rstest = "0.24.0"

examples/quickstart/main.rs

@@ -2,7 +2,7 @@ use std::str::FromStr;
 
 use num_bigint::BigUint;
 use tycho_core::{
-    dto::{Chain, ProtocolComponent},
+    dto::{Chain as TychoCoreChain, ProtocolComponent},
     Bytes,
 };
 use tycho_execution::encoding::{
@@ -10,7 +10,7 @@ use tycho_execution::encoding::{
         strategy_encoder::strategy_encoder_registry::EVMStrategyEncoderRegistry,
         tycho_encoder::EVMTychoEncoder,
     },
-    models::{Solution, Swap},
+    models::{Chain, Solution, Swap},
     strategy_encoder::StrategyEncoderRegistry,
     tycho_encoder::TychoEncoder,
 };
@@ -24,11 +24,12 @@ fn main() {
         .expect("Failed to create user address");
     let executors_file_path = "src/encoding/config/executor_addresses.json";
 
+    let eth_chain = Chain::from(TychoCoreChain::Ethereum);
     // Initialize the encoder
     let strategy_encoder_registry =
-        EVMStrategyEncoderRegistry::new(Chain::Ethereum, executors_file_path, signer_pk.clone())
+        EVMStrategyEncoderRegistry::new(eth_chain.clone(), executors_file_path, signer_pk.clone())
             .expect("Failed to create strategy encoder registry");
-    let encoder = EVMTychoEncoder::new(strategy_encoder_registry, router_address)
+    let encoder = EVMTychoEncoder::new(strategy_encoder_registry, router_address, eth_chain)
         .expect("Failed to create encoder");
 
     // ------------------- Encode a simple swap -------------------

src/encoding/evm/approvals/permit2.rs

@@ -12,15 +12,15 @@ use alloy_sol_types::{eip712_domain, sol, SolStruct, SolValue};
 use chrono::Utc;
 use num_bigint::BigUint;
 use tokio::runtime::Runtime;
-use tycho_core::{dto::Chain, Bytes};
+use tycho_core::Bytes;
 
 use crate::encoding::{
     errors::EncodingError,
     evm::{
         approvals::protocol_approvals_manager::get_client,
-        models::ChainId,
         utils::{biguint_to_u256, bytes_to_address, encode_input},
     },
+    models::{Chain, ChainId},
 };
 
 /// Struct for managing Permit2 operations, including encoding approvals and fetching allowance
@@ -60,7 +60,6 @@ sol! {
 
 impl Permit2 {
     pub fn new(signer_pk: String, chain: Chain) -> Result<Self, EncodingError> {
-        let chain_id = ChainId::from(chain);
         let runtime = Runtime::new()
             .map_err(|_| EncodingError::FatalError("Failed to create runtime".to_string()))?;
         let client = runtime.block_on(get_client())?;
@@ -76,7 +75,7 @@ impl Permit2 {
             client,
             runtime,
             signer,
-            chain_id,
+            chain_id: chain.id,
         })
     }
 
@@ -165,6 +164,7 @@ mod tests {
 
     use alloy_primitives::Uint;
     use num_bigint::BigUint;
+    use tycho_core::dto::Chain as TychoCoreChain;
 
     use super::*;
 
@@ -199,11 +199,15 @@ mod tests {
         }
     }
 
+    fn eth_chain() -> Chain {
+        TychoCoreChain::Ethereum.into()
+    }
+
     #[test]
     fn test_get_existing_allowance() {
         let signer_pk =
             "4c0883a69102937d6231471b5dbb6204fe512961708279feb1be6ae5538da033".to_string();
-        let manager = Permit2::new(signer_pk, Chain::Ethereum).unwrap();
+        let manager = Permit2::new(signer_pk, eth_chain()).unwrap();
 
         let token = Bytes::from_str("0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48").unwrap();
         let owner = Bytes::from_str("0x2c6a3cd97c6283b95ac8c5a4459ebb0d5fd404f4").unwrap();
@@ -223,7 +227,7 @@ mod tests {
         // Set up a mock private key for signing
         let private_key =
             "4c0883a69102937d6231471b5dbb6204fe512961708279feb1be6ae5538da033".to_string();
-        let permit2 = Permit2::new(private_key, Chain::Ethereum).expect("Failed to create Permit2");
+        let permit2 = Permit2::new(private_key, eth_chain()).expect("Failed to create Permit2");
 
         let owner = Bytes::from_str("0x2c6a3cd97c6283b95ac8c5a4459ebb0d5fd404f4").unwrap();
         let spender = Bytes::from_str("0xba12222222228d8ba445958a75a0704d566bf2c8").unwrap();
@@ -264,7 +268,7 @@ mod tests {
             "0xac0974bec39a17e36ba4a6b4d238ff944bacb478cbed5efcae784d7bf4f2ff80".to_string();
 
         let permit2 =
-            Permit2::new(anvil_private_key, Chain::Ethereum).expect("Failed to create Permit2");
+            Permit2::new(anvil_private_key, eth_chain()).expect("Failed to create Permit2");
 
         let token = Bytes::from_str("0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48").unwrap();
         let amount = BigUint::from(1000u64);

src/encoding/evm/constants.rs

@@ -1,10 +0,0 @@
-use alloy_primitives::hex;
-use lazy_static::lazy_static;
-use tycho_core::Bytes;
-
-lazy_static! {
-    pub static ref NATIVE_ADDRESS: Bytes =
-        Bytes::from(hex!("0000000000000000000000000000000000000000").to_vec());
-    pub static ref WETH_ADDRESS: Bytes =
-        Bytes::from(hex!("c02aaa39b223fe8d0a0e5c4f27ead9083c756cc2").to_vec());
-}

src/encoding/evm/mod.rs

@@ -1,6 +1,4 @@
 pub mod approvals;
-mod constants;
-mod models;
 pub mod strategy_encoder;
 mod swap_encoder;
 pub mod tycho_encoder;

src/encoding/evm/models.rs

@@ -1,20 +0,0 @@
-use tycho_core::dto::Chain;
-
-pub struct ChainId(u64);
-
-impl ChainId {
-    pub fn id(&self) -> u64 {
-        self.0
-    }
-}
-
-impl From<Chain> for ChainId {
-    fn from(chain: Chain) -> Self {
-        match chain {
-            Chain::Ethereum => ChainId(1),
-            Chain::ZkSync => ChainId(324),
-            Chain::Arbitrum => ChainId(42161),
-            Chain::Starknet => ChainId(0),
-        }
-    }
-}

src/encoding/evm/strategy_encoder/strategy_encoder_registry.rs

@@ -1,14 +1,12 @@
 use std::collections::HashMap;
 
-use tycho_core::dto::Chain;
-
 use crate::encoding::{
     errors::EncodingError,
     evm::{
         strategy_encoder::strategy_encoders::{ExecutorStrategyEncoder, SplitSwapStrategyEncoder},
         swap_encoder::swap_encoder_registry::SwapEncoderRegistry,
     },
-    models::Solution,
+    models::{Chain, Solution},
     strategy_encoder::{StrategyEncoder, StrategyEncoderRegistry},
 };
 
@@ -22,7 +20,7 @@ impl StrategyEncoderRegistry for EVMStrategyEncoderRegistry {
         executors_file_path: &str,
         signer_pk: Option<String>,
     ) -> Result<Self, EncodingError> {
-        let swap_encoder_registry = SwapEncoderRegistry::new(executors_file_path, Chain::Ethereum)?;
+        let swap_encoder_registry = SwapEncoderRegistry::new(executors_file_path, chain.clone())?;
 
         let mut strategies: HashMap<String, Box<dyn StrategyEncoder>> = HashMap::new();
         strategies.insert(

src/encoding/evm/strategy_encoder/strategy_encoders.rs

@@ -7,17 +7,16 @@ use std::{
 use alloy_primitives::{aliases::U24, FixedBytes, U256, U8};
 use alloy_sol_types::SolValue;
 use num_bigint::BigUint;
-use tycho_core::{dto::Chain, keccak256, Bytes};
+use tycho_core::{keccak256, Bytes};
 
 use crate::encoding::{
     errors::EncodingError,
     evm::{
         approvals::permit2::Permit2,
-        constants::{NATIVE_ADDRESS, WETH_ADDRESS},
         swap_encoder::swap_encoder_registry::SwapEncoderRegistry,
         utils::{biguint_to_u256, bytes_to_address, encode_input, percentage_to_uint24},
     },
-    models::{EncodingContext, NativeAction, Solution, Swap},
+    models::{Chain, EncodingContext, NativeAction, Solution, Swap},
     strategy_encoder::StrategyEncoder,
     swap_encoder::SwapEncoder,
 };
@@ -62,6 +61,8 @@ pub struct SplitSwapStrategyEncoder {
     swap_encoder_registry: SwapEncoderRegistry,
     permit2: Permit2,
     selector: String,
+    native_address: Bytes,
+    wrapped_address: Bytes,
 }
 
 impl SplitSwapStrategyEncoder {
@@ -71,7 +72,13 @@ impl SplitSwapStrategyEncoder {
         swap_encoder_registry: SwapEncoderRegistry,
     ) -> Result<Self, EncodingError> {
         let selector = "swap(uint256,address,address,uint256,bool,bool,uint256,address,((address,uint160,uint48,uint48),address,uint256),bytes,bytes)".to_string();
-        Ok(Self { permit2: Permit2::new(signer_pk, chain)?, selector, swap_encoder_registry })
+        Ok(Self {
+            permit2: Permit2::new(signer_pk, chain.clone())?,
+            selector,
+            swap_encoder_registry,
+            native_address: chain.native_token()?,
+            wrapped_address: chain.wrapped_token()?,
+        })
     }
 
     fn validate_split_percentages(&self, swaps: &[Swap]) -> Result<(), EncodingError> {
@@ -152,18 +159,18 @@ impl SplitSwapStrategyEncoder {
         native_action: &Option<NativeAction>,
     ) -> Result<(), EncodingError> {
         // Convert ETH to WETH only if there's a corresponding wrap/unwrap action
-        let given_token = if *given_token == *NATIVE_ADDRESS {
+        let given_token = if *given_token == *self.native_address {
             match native_action {
-                Some(NativeAction::Wrap) => &WETH_ADDRESS,
+                Some(NativeAction::Wrap) => &self.wrapped_address,
                 _ => given_token,
             }
         } else {
             given_token
         };
 
-        let checked_token = if *checked_token == *NATIVE_ADDRESS {
+        let checked_token = if *checked_token == *self.native_address {
             match native_action {
-                Some(NativeAction::Unwrap) => &WETH_ADDRESS,
+                Some(NativeAction::Unwrap) => &self.wrapped_address,
                 _ => checked_token,
             }
         } else {
@@ -278,14 +285,14 @@ impl StrategyEncoder for SplitSwapStrategyEncoder {
 
         let mut tokens = Vec::with_capacity(2 + intermediary_tokens.len());
         if wrap {
-            tokens.push(WETH_ADDRESS.clone());
+            tokens.push(self.wrapped_address.clone());
         } else {
             tokens.push(solution.given_token.clone());
         }
         tokens.extend(intermediary_tokens);
 
         if unwrap {
-            tokens.push(WETH_ADDRESS.clone());
+            tokens.push(self.wrapped_address.clone());
         } else {
             tokens.push(solution.checked_token.clone());
         }
@@ -424,19 +431,32 @@ mod tests {
     use std::str::FromStr;
 
     use alloy::hex::encode;
+    use alloy_primitives::hex;
     use num_bigint::BigUint;
     use rstest::rstest;
-    use tycho_core::{dto::ProtocolComponent, Bytes};
+    use tycho_core::{
-
+        dto::{Chain as TychoCoreChain, ProtocolComponent},
-    use super::*;
+        Bytes,
-    use crate::encoding::{
-        evm::constants::{NATIVE_ADDRESS, WETH_ADDRESS},
-        models::Swap,
     };
 
+    use super::*;
+    use crate::encoding::models::Swap;
+
+    fn eth_chain() -> Chain {
+        TychoCoreChain::Ethereum.into()
+    }
+
+    fn eth() -> Bytes {
+        Bytes::from(hex!("0000000000000000000000000000000000000000").to_vec())
+    }
+
+    fn weth() -> Bytes {
+        Bytes::from(hex!("c02aaa39b223fe8d0a0e5c4f27ead9083c756cc2").to_vec())
+    }
+
     fn get_swap_encoder_registry() -> SwapEncoderRegistry {
-        SwapEncoderRegistry::new("src/encoding/config/executor_addresses.json", Chain::Ethereum)
+        let eth_chain = eth_chain();
-            .unwrap()
+        SwapEncoderRegistry::new("src/encoding/config/executor_addresses.json", eth_chain).unwrap()
     }
 
     #[test]
@@ -444,7 +464,7 @@ mod tests {
         let swap_encoder_registry = get_swap_encoder_registry();
         let encoder = ExecutorStrategyEncoder::new(swap_encoder_registry);
 
-        let token_in = Bytes::from("0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2");
+        let token_in = weth();
         let token_out = Bytes::from("0x6b175474e89094c44da98b954eedeac495271d0f");
 
         let swap = Swap {
@@ -550,8 +570,7 @@ mod tests {
         };
         let swap_encoder_registry = get_swap_encoder_registry();
         let encoder =
-            SplitSwapStrategyEncoder::new(private_key, Chain::Ethereum, swap_encoder_registry)
+            SplitSwapStrategyEncoder::new(private_key, eth_chain(), swap_encoder_registry).unwrap();
-                .unwrap();
         let solution = Solution {
             exact_out: false,
             given_token: weth,
@@ -646,17 +665,16 @@ mod tests {
                 protocol_system: "uniswap_v2".to_string(),
                 ..Default::default()
             },
-            token_in: WETH_ADDRESS.clone(),
+            token_in: weth(),
             token_out: dai.clone(),
             split: 0f64,
         };
         let swap_encoder_registry = get_swap_encoder_registry();
         let encoder =
-            SplitSwapStrategyEncoder::new(private_key, Chain::Ethereum, swap_encoder_registry)
+            SplitSwapStrategyEncoder::new(private_key, eth_chain(), swap_encoder_registry).unwrap();
-                .unwrap();
         let solution = Solution {
             exact_out: false,
-            given_token: NATIVE_ADDRESS.clone(),
+            given_token: eth(),
             given_amount: BigUint::from_str("1_000000000000000000").unwrap(),
             checked_token: dai,
             expected_amount: Some(BigUint::from_str("3_000_000000000000000000").unwrap()),
@@ -696,18 +714,17 @@ mod tests {
                 ..Default::default()
             },
             token_in: dai.clone(),
-            token_out: WETH_ADDRESS.clone(),
+            token_out: weth(),
             split: 0f64,
         };
         let swap_encoder_registry = get_swap_encoder_registry();
         let encoder =
-            SplitSwapStrategyEncoder::new(private_key, Chain::Ethereum, swap_encoder_registry)
+            SplitSwapStrategyEncoder::new(private_key, eth_chain(), swap_encoder_registry).unwrap();
-                .unwrap();
         let solution = Solution {
             exact_out: false,
             given_token: dai,
             given_amount: BigUint::from_str("3_000_000000000000000000").unwrap(),
-            checked_token: NATIVE_ADDRESS.clone(),
+            checked_token: eth(),
             expected_amount: Some(BigUint::from_str("1_000000000000000000").unwrap()),
             check_amount: None,
             sender: Bytes::from_str("0xcd09f75E2BF2A4d11F3AB23f1389FcC1621c0cc2").unwrap(),
@@ -742,7 +759,7 @@ mod tests {
         let private_key =
             "0x123456789abcdef123456789abcdef123456789abcdef123456789abcdef1234".to_string();
 
-        let weth = Bytes::from_str("0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2").unwrap();
+        let weth = weth();
         let dai = Bytes::from_str("0x6b175474e89094c44da98b954eedeac495271d0f").unwrap();
         let wbtc = Bytes::from_str("0x2260fac5e5542a773aa44fbcfedf7c193bc2c599").unwrap();
         let usdc = Bytes::from_str("0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48").unwrap();
@@ -791,8 +808,7 @@ mod tests {
         };
         let swap_encoder_registry = get_swap_encoder_registry();
         let encoder =
-            SplitSwapStrategyEncoder::new(private_key, Chain::Ethereum, swap_encoder_registry)
+            SplitSwapStrategyEncoder::new(private_key, eth_chain(), swap_encoder_registry).unwrap();
-                .unwrap();
         let solution = Solution {
             exact_out: false,
             given_token: weth,
@@ -819,7 +835,7 @@ mod tests {
         let private_key =
             "0x123456789abcdef123456789abcdef123456789abcdef123456789abcdef1234".to_string();
         let swap_encoder_registry = get_swap_encoder_registry();
-        SplitSwapStrategyEncoder::new(private_key, Chain::Ethereum, swap_encoder_registry).unwrap()
+        SplitSwapStrategyEncoder::new(private_key, eth_chain(), swap_encoder_registry).unwrap()
     }
 
     #[test]

src/encoding/evm/swap_encoder/swap_encoder_registry.rs

@@ -1,9 +1,7 @@
 use std::{collections::HashMap, fs};
 
-use tycho_core::dto::Chain;
-
 use crate::encoding::{
-    errors::EncodingError, evm::swap_encoder::builder::SwapEncoderBuilder,
+    errors::EncodingError, evm::swap_encoder::builder::SwapEncoderBuilder, models::Chain,
     swap_encoder::SwapEncoder,
 };
 
@@ -15,10 +13,10 @@ pub struct SwapEncoderRegistry {
 impl SwapEncoderRegistry {
     pub fn new(executors_file_path: &str, blockchain: Chain) -> Result<Self, EncodingError> {
         let config_str = fs::read_to_string(executors_file_path)?;
-        let config: HashMap<Chain, HashMap<String, String>> = serde_json::from_str(&config_str)?;
+        let config: HashMap<String, HashMap<String, String>> = serde_json::from_str(&config_str)?;
         let mut encoders = HashMap::new();
         let executors = config
-            .get(&blockchain)
+            .get(&blockchain.name)
             .ok_or(EncodingError::FatalError("No executors found for blockchain".to_string()))?;
         for (protocol, executor_address) in executors {
             let builder = SwapEncoderBuilder::new(protocol, executor_address);

src/encoding/evm/tycho_encoder.rs

@@ -5,8 +5,7 @@ use tycho_core::Bytes;
 
 use crate::encoding::{
     errors::EncodingError,
-    evm::constants::{NATIVE_ADDRESS, WETH_ADDRESS},
+    models::{Chain, NativeAction, Solution, Transaction},
-    models::{NativeAction, Solution, Transaction},
     strategy_encoder::StrategyEncoderRegistry,
     tycho_encoder::TychoEncoder,
 };
@@ -14,13 +13,29 @@ use crate::encoding::{
 pub struct EVMTychoEncoder<S: StrategyEncoderRegistry> {
     strategy_selector: S,
     router_address: Bytes,
+    native_address: Bytes,
+    wrapped_address: Bytes,
 }
 
 impl<S: StrategyEncoderRegistry> EVMTychoEncoder<S> {
-    pub fn new(strategy_selector: S, router_address: String) -> Result<Self, EncodingError> {
+    pub fn new(
+        strategy_selector: S,
+        router_address: String,
+        chain: Chain,
+    ) -> Result<Self, EncodingError> {
         let router_address = Bytes::from_str(&router_address)
             .map_err(|_| EncodingError::FatalError("Invalid router address".to_string()))?;
-        Ok(EVMTychoEncoder { strategy_selector, router_address })
+        if chain.name != *"ethereum" {
+            return Err(EncodingError::InvalidInput(
+                "Currently only Ethereum is supported".to_string(),
+            ));
+        }
+        Ok(EVMTychoEncoder {
+            strategy_selector,
+            router_address,
+            native_address: chain.native_token()?,
+            wrapped_address: chain.wrapped_token()?,
+        })
     }
 }
 
@@ -36,28 +51,30 @@ impl<S: StrategyEncoderRegistry> EVMTychoEncoder<S> {
         }
         if let Some(native_action) = solution.clone().native_action {
             if native_action == NativeAction::Wrap {
-                if solution.given_token != *NATIVE_ADDRESS {
+                if solution.given_token != self.native_address {
                     return Err(EncodingError::FatalError(
-                        "ETH must be the input token in order to wrap".to_string(),
+                        "Native token must be the input token in order to wrap".to_string(),
                     ));
                 }
                 if let Some(first_swap) = solution.swaps.first() {
-                    if first_swap.token_in != *WETH_ADDRESS {
+                    if first_swap.token_in != self.wrapped_address {
                         return Err(EncodingError::FatalError(
-                            "WETH must be the first swap's input in order to wrap".to_string(),
+                            "Wrapped token must be the first swap's input in order to wrap"
+                                .to_string(),
                         ));
                     }
                 }
             } else if native_action == NativeAction::Unwrap {
-                if solution.checked_token != *NATIVE_ADDRESS {
+                if solution.checked_token != self.native_address {
                     return Err(EncodingError::FatalError(
-                        "ETH must be the output token in order to unwrap".to_string(),
+                        "Native token must be the output token in order to unwrap".to_string(),
                     ));
                 }
                 if let Some(last_swap) = solution.swaps.last() {
-                    if last_swap.token_out != *WETH_ADDRESS {
+                    if last_swap.token_out != self.wrapped_address {
                         return Err(EncodingError::FatalError(
-                            "WETH must be the last swap's output in order to unwrap".to_string(),
+                            "Wrapped token must be the last swap's output in order to unwrap"
+                                .to_string(),
                         ));
                     }
                 }
@@ -104,7 +121,7 @@ impl<S: StrategyEncoderRegistry> TychoEncoder<S> for EVMTychoEncoder<S> {
 
 #[cfg(test)]
 mod tests {
-    use tycho_core::dto::{Chain, ProtocolComponent};
+    use tycho_core::dto::{Chain as TychoCoreChain, ProtocolComponent};
 
     use super::*;
     use crate::encoding::{
@@ -115,10 +132,22 @@ mod tests {
         strategy: Box<dyn StrategyEncoder>,
     }
 
+    fn eth_chain() -> Chain {
+        TychoCoreChain::Ethereum.into()
+    }
+
     fn dai() -> Bytes {
         Bytes::from_str("0x6b175474e89094c44da98b954eedeac495271d0f").unwrap()
     }
 
+    fn eth() -> Bytes {
+        Bytes::from_str("0x0000000000000000000000000000000000000000").unwrap()
+    }
+
+    fn weth() -> Bytes {
+        Bytes::from_str("0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2").unwrap()
+    }
+
     impl StrategyEncoderRegistry for MockStrategyRegistry {
         fn new(
             _chain: Chain,
@@ -158,10 +187,11 @@ mod tests {
     }
 
     fn get_mocked_tycho_encoder() -> EVMTychoEncoder<MockStrategyRegistry> {
-        let strategy_selector = MockStrategyRegistry::new(Chain::Ethereum, "", None).unwrap();
+        let strategy_selector = MockStrategyRegistry::new(eth_chain(), "", None).unwrap();
         EVMTychoEncoder::new(
             strategy_selector,
             "0x1234567890abcdef1234567890abcdef12345678".to_string(),
+            eth_chain(),
         )
         .unwrap()
     }
@@ -176,7 +206,7 @@ mod tests {
                 protocol_system: "uniswap_v2".to_string(),
                 ..Default::default()
             },
-            token_in: WETH_ADDRESS.clone(),
+            token_in: weth(),
             token_out: dai(),
             split: 0f64,
         };
@@ -184,7 +214,7 @@ mod tests {
         let solution = Solution {
             exact_out: false,
             given_amount: eth_amount_in.clone(),
-            given_token: NATIVE_ADDRESS.clone(),
+            given_token: eth(),
             router_address: None,
             swaps: vec![swap],
             native_action: Some(NativeAction::Wrap),
@@ -228,14 +258,14 @@ mod tests {
                 protocol_system: "uniswap_v2".to_string(),
                 ..Default::default()
             },
-            token_in: WETH_ADDRESS.clone(),
+            token_in: weth(),
             token_out: dai(),
             split: 0f64,
         };
 
         let solution = Solution {
             exact_out: false,
-            given_token: NATIVE_ADDRESS.clone(),
+            given_token: eth(),
             checked_token: dai(),
             check_amount: None,
             swaps: vec![swap],
@@ -257,14 +287,14 @@ mod tests {
                 protocol_system: "uniswap_v2".to_string(),
                 ..Default::default()
             },
-            token_in: WETH_ADDRESS.clone(),
+            token_in: weth(),
             token_out: dai(),
             split: 0f64,
         };
 
         let solution = Solution {
             exact_out: false,
-            given_token: WETH_ADDRESS.clone(),
+            given_token: weth(),
             swaps: vec![swap],
             native_action: Some(NativeAction::Wrap),
             ..Default::default()
@@ -275,7 +305,9 @@ mod tests {
         assert!(result.is_err());
         assert_eq!(
             result.err().unwrap(),
-            EncodingError::FatalError("ETH must be the input token in order to wrap".to_string())
+            EncodingError::FatalError(
+                "Native token must be the input token in order to wrap".to_string()
+            )
         );
     }
 
@@ -288,14 +320,14 @@ mod tests {
                 protocol_system: "uniswap_v2".to_string(),
                 ..Default::default()
             },
-            token_in: NATIVE_ADDRESS.clone(),
+            token_in: eth(),
             token_out: dai(),
             split: 0f64,
         };
 
         let solution = Solution {
             exact_out: false,
-            given_token: NATIVE_ADDRESS.clone(),
+            given_token: eth(),
             swaps: vec![swap],
             native_action: Some(NativeAction::Wrap),
             ..Default::default()
@@ -307,7 +339,7 @@ mod tests {
         assert_eq!(
             result.err().unwrap(),
             EncodingError::FatalError(
-                "WETH must be the first swap's input in order to wrap".to_string()
+                "Wrapped token must be the first swap's input in order to wrap".to_string()
             )
         );
     }
@@ -317,7 +349,7 @@ mod tests {
         let encoder = get_mocked_tycho_encoder();
         let solution = Solution {
             exact_out: false,
-            given_token: NATIVE_ADDRESS.clone(),
+            given_token: eth(),
             swaps: vec![],
             native_action: Some(NativeAction::Wrap),
             ..Default::default()
@@ -342,13 +374,13 @@ mod tests {
                 ..Default::default()
             },
             token_in: dai(),
-            token_out: WETH_ADDRESS.clone(),
+            token_out: weth(),
             split: 0f64,
         };
 
         let solution = Solution {
             exact_out: false,
-            checked_token: NATIVE_ADDRESS.clone(),
+            checked_token: eth(),
             check_amount: None,
             swaps: vec![swap],
             native_action: Some(NativeAction::Unwrap),
@@ -370,14 +402,14 @@ mod tests {
                 ..Default::default()
             },
             token_in: dai(),
-            token_out: WETH_ADDRESS.clone(),
+            token_out: weth(),
             split: 0f64,
         };
 
         let solution = Solution {
             exact_out: false,
             given_token: dai(),
-            checked_token: WETH_ADDRESS.clone(),
+            checked_token: weth(),
             swaps: vec![swap],
             native_action: Some(NativeAction::Unwrap),
             ..Default::default()
@@ -389,7 +421,7 @@ mod tests {
         assert_eq!(
             result.err().unwrap(),
             EncodingError::FatalError(
-                "ETH must be the output token in order to unwrap".to_string()
+                "Native token must be the output token in order to unwrap".to_string()
             )
         );
     }
@@ -404,13 +436,13 @@ mod tests {
                 ..Default::default()
             },
             token_in: dai(),
-            token_out: NATIVE_ADDRESS.clone(),
+            token_out: eth(),
             split: 0f64,
         };
 
         let solution = Solution {
             exact_out: false,
-            checked_token: NATIVE_ADDRESS.clone(),
+            checked_token: eth(),
             swaps: vec![swap],
             native_action: Some(NativeAction::Unwrap),
             ..Default::default()
@@ -422,7 +454,7 @@ mod tests {
         assert_eq!(
             result.err().unwrap(),
             EncodingError::FatalError(
-                "WETH must be the last swap's output in order to unwrap".to_string()
+                "Wrapped token must be the last swap's output in order to unwrap".to_string()
             )
         );
     }

src/encoding/models.rs

@@ -1,5 +1,11 @@
+use hex;
 use num_bigint::BigUint;
-use tycho_core::{dto::ProtocolComponent, Bytes};
+use tycho_core::{
+    dto::{Chain as TychoCoreChain, ProtocolComponent},
+    Bytes,
+};
+
+use crate::encoding::errors::EncodingError;
 
 #[derive(Clone, Default, Debug)]
 pub struct Solution {
@@ -80,6 +86,63 @@ pub struct EncodingContext {
     pub router_address: Bytes,
 }
 
+#[derive(Clone, PartialEq, Eq, Hash)]
+pub struct ChainId(pub u64);
+
+#[derive(Clone, PartialEq, Eq, Hash)]
+pub struct Chain {
+    pub id: ChainId,
+    pub name: String,
+}
+
+impl ChainId {
+    pub fn id(&self) -> u64 {
+        self.0
+    }
+}
+
+impl From<TychoCoreChain> for ChainId {
+    fn from(chain: TychoCoreChain) -> Self {
+        match chain {
+            TychoCoreChain::Ethereum => ChainId(1),
+            TychoCoreChain::ZkSync => ChainId(324),
+            TychoCoreChain::Arbitrum => ChainId(42161),
+            TychoCoreChain::Starknet => ChainId(0),
+        }
+    }
+}
+
+impl From<TychoCoreChain> for Chain {
+    fn from(chain: TychoCoreChain) -> Self {
+        Chain { id: chain.into(), name: chain.to_string() }
+    }
+}
+
+impl Chain {
+    pub fn native_token(&self) -> Result<Bytes, EncodingError> {
+        match self.id.id() {
+            1 => Ok(Bytes::from(hex::decode("0000000000000000000000000000000000000000").map_err(
+                |_| EncodingError::FatalError("Failed to decode native token".to_string()),
+            )?)),
+            _ => Err(EncodingError::InvalidInput(format!(
+                "Native token not set for chain {:?}. Double check the chain is supported.",
+                self.name
+            ))),
+        }
+    }
+    pub fn wrapped_token(&self) -> Result<Bytes, EncodingError> {
+        match self.id.id() {
+            1 => Ok(Bytes::from(hex::decode("C02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2").map_err(
+                |_| EncodingError::FatalError("Failed to decode wrapped token".to_string()),
+            )?)),
+            _ => Err(EncodingError::InvalidInput(format!(
+                "Wrapped token not set for chain {:?}. Double check the chain is supported.",
+                self.name
+            ))),
+        }
+    }
+}
+
 mod tests {
     use super::*;
 

src/encoding/strategy_encoder.rs

@@ -1,6 +1,10 @@
-use tycho_core::{dto::Chain, Bytes};
+use tycho_core::Bytes;
 
-use crate::encoding::{errors::EncodingError, models::Solution, swap_encoder::SwapEncoder};
+use crate::encoding::{
+    errors::EncodingError,
+    models::{Chain, Solution},
+    swap_encoder::SwapEncoder,
+};
 
 pub trait StrategyEncoder {
     fn encode_strategy(