feat: early return in usv4 swap encoder for second swap, add utils

foundry/test/executors/UniswapV4Executor.t.sol

@@ -197,7 +197,7 @@ contract UniswapV4ExecutorTest is Test, Constants {
     }
 
     function testMultipleSwapIntegration() public {
-        // USDE -> USDT -> WBTC ->
+        // USDE -> USDT -> WBTC
         // Generated by the Tycho swap encoder - test_encode_uniswap_v4_combined
         bytes memory protocolData =
             hex"4c9edd5852cd905f086c759e8383e09bff1e68b32260fac5e5542a773aa44fbcfedf7c193bc2c5990000000000000000000000000000000000000000000000000000000000000001015615deb798bb3e4dfa0139dfa1b3d433cc23b72f91dd7346dac17f958d2ee523a2206206994597c13d831ec70000640000012260fac5e5542a773aa44fbcfedf7c193bc2c599000bb800003c";

src/encoding/evm/swap_encoder/swap_encoders.rs

@@ -7,7 +7,9 @@ use crate::encoding::{
     errors::EncodingError,
     evm::{
         approvals::protocol_approvals_manager::ProtocolApprovalsManager,
-        utils::{bytes_to_address, encode_function_selector, pad_to_fixed_size},
+        utils::{
+            bytes_to_address, encode_function_selector, get_static_attribute, pad_to_fixed_size,
+        },
     },
     models::{EncodingContext, Swap},
     swap_encoder::SwapEncoder,
@@ -105,16 +107,7 @@ impl SwapEncoder for UniswapV3SwapEncoder {
         let zero_to_one = Self::get_zero_to_one(token_in_address, token_out_address);
         let component_id = Address::from_str(&swap.component.id)
             .map_err(|_| EncodingError::FatalError("Invalid USV3 component id".to_string()))?;
-        let pool_fee_bytes = swap
+        let pool_fee_bytes = get_static_attribute(&swap, "fee")?;
-            .component
-            .static_attributes
-            .get("fee")
-            .ok_or_else(|| {
-                EncodingError::FatalError(
-                    "Pool fee not found in Uniswap v3 static attributes".to_string(),
-                )
-            })?
-            .to_vec();
 
         let pool_fee_u24 = pad_to_fixed_size::<3>(&pool_fee_bytes)
             .map_err(|_| EncodingError::FatalError("Failed to extract fee bytes".to_string()))?;
@@ -183,10 +176,27 @@ impl SwapEncoder for UniswapV4SwapEncoder {
         swap: Swap,
         encoding_context: EncodingContext,
     ) -> Result<Vec<u8>, EncodingError> {
-        let mut first_swap = false;
+        let fee = get_static_attribute(&swap, "fee")?;
-        if encoding_context.group_token_in == Some(swap.token_in.clone()) {
+
-            first_swap = true;
+        let pool_fee_u24 = pad_to_fixed_size::<3>(&fee)
+            .map_err(|_| EncodingError::FatalError("Failed to pad fee bytes".to_string()))?;
+
+        let tick_spacing = get_static_attribute(&swap, "tickSpacing")?;
+
+        let pool_tick_spacing_u24 = pad_to_fixed_size::<3>(&tick_spacing).map_err(|_| {
+            EncodingError::FatalError("Failed to pad tick spacing bytes".to_string())
+        })?;
+
+        // Early check if this is not the first swap
+        if encoding_context.group_token_in != Some(swap.token_in.clone()) {
+            return Ok(Self::encode_pool_params(
+                bytes_to_address(&swap.token_out)?,
+                pool_fee_u24,
+                pool_tick_spacing_u24,
+            ));
         }
+
+        // This is the first swap, compute all necessary values
         let token_in_address = bytes_to_address(&swap.token_in)?;
         let token_out_address = bytes_to_address(&swap.token_out)?;
         let group_token_in = if let Some(group_token_in) = encoding_context.group_token_in {
@@ -208,46 +218,9 @@ impl SwapEncoder for UniswapV4SwapEncoder {
         let zero_to_one = Self::get_zero_to_one(token_in_address, token_out_address);
         let callback_executor = bytes_to_address(&encoding_context.router_address)?;
 
-        let fee = swap
-            .component
-            .static_attributes
-            .get("fee")
-            .ok_or_else(|| {
-                EncodingError::FatalError(
-                    "Pool fee not found in Uniswap v4 static attributes".to_string(),
-                )
-            })?
-            .to_vec();
-
-        let pool_fee_u24 = pad_to_fixed_size::<3>(&fee)
-            .map_err(|_| EncodingError::FatalError("Failed to extract fee bytes".to_string()))?;
-
-        let tick_spacing = swap
-            .component
-            .static_attributes
-            .get("tickSpacing")
-            .ok_or_else(|| {
-                EncodingError::FatalError(
-                    "Pool tick spacing not found in Uniswap v4 static attributes".to_string(),
-                )
-            })?
-            .to_vec();
-
-        let pool_tick_spacing_u24 = pad_to_fixed_size::<3>(&tick_spacing).map_err(|_| {
-            EncodingError::FatalError("Failed to extract tick spacing bytes".to_string())
-        })?;
-
         let pool_params =
             Self::encode_pool_params(token_out_address, pool_fee_u24, pool_tick_spacing_u24);
 
-        if !first_swap {
-            return Ok(Self::encode_pool_params(
-                token_out_address,
-                pool_fee_u24,
-                pool_tick_spacing_u24,
-            ));
-        }
-
         let args = (
             group_token_in,
             group_token_out,
@@ -527,9 +500,9 @@ mod tests {
         assert_eq!(
             hex_swap,
             String::from(concat!(
-                // token in
+                // group token in
                 "4c9edd5852cd905f086c759e8383e09bff1e68b3",
-                // token out
+                // group token out
                 "dac17f958d2ee523a2206206994597c13d831ec7",
                 // amount out min (0 as u128)
                 "0000000000000000000000000000000000000000000000000000000000000001",
@@ -540,11 +513,11 @@ mod tests {
                 // callback selector for "unlockCallback(bytes)"
                 "91dd7346",
                 // pool params:
-                // - intermediary token (20 bytes)
+                // - intermediary token
                 "dac17f958d2ee523a2206206994597c13d831ec7",
-                // - fee (3 bytes)
+                // - fee
                 "000064",
-                // - tick spacing (3 bytes)
+                // - tick spacing
                 "000001"
             ))
         );
@@ -556,6 +529,7 @@ mod tests {
         let tick_spacing = BigInt::from(60);
         let encoded_pool_fee = Bytes::from(fee.to_signed_bytes_be());
         let encoded_tick_spacing = Bytes::from(tick_spacing.to_signed_bytes_be());
+        let group_token_in = Bytes::from("0x4c9EDD5852cd905f086C759E8383e09bff1E68B3"); // USDE
         let token_in = Bytes::from("0xdAC17F958D2ee523a2206206994597C13D831ec7"); // USDT
         let token_out = Bytes::from("0x2260FAC5E5542a773Aa44fBCfeDf7C193bc2C599"); // WBTC
 
@@ -580,9 +554,9 @@ mod tests {
             receiver: Bytes::from("0x0000000000000000000000000000000000000001"),
             exact_out: false,
             router_address: Bytes::zero(20),
-            // Different from token_in and token_out
+            group_token_in: Some(group_token_in),
-            group_token_in: Some(Bytes::zero(20)),
+            // Token out is the same as the group token out
-            group_token_out: Some(Bytes::zero(20)),
+            group_token_out: Some(token_out),
             amount_out_min: Some(BigUint::from(1u128)),
         };
 

src/encoding/evm/utils.rs

@@ -4,7 +4,10 @@ use alloy_primitives::{aliases::U24, keccak256, Address, FixedBytes, Keccak256,
 use num_bigint::BigUint;
 use tycho_core::Bytes;
 
-use crate::encoding::{errors::EncodingError, models::Solution};
+use crate::encoding::{
+    errors::EncodingError,
+    models::{Solution, Swap},
+};
 
 /// Safely converts a `Bytes` object to an `Address` object.
 ///
@@ -92,7 +95,7 @@ pub fn get_token_position(tokens: Vec<Bytes>, token: Bytes) -> Result<U8, Encodi
     Ok(position)
 }
 
-/// Pads a byte slice to a fixed size array with leading zeros.
+/// Pads a byte slice to a fixed size array of N bytes.
 pub fn pad_to_fixed_size<const N: usize>(input: &[u8]) -> Result<[u8; N], EncodingError> {
     let mut padded = [0u8; N];
     let start = N - input.len();
@@ -105,3 +108,15 @@ pub fn encode_function_selector(selector: &str) -> FixedBytes<4> {
     let hash = keccak256(selector.as_bytes());
     FixedBytes::<4>::from([hash[0], hash[1], hash[2], hash[3]])
 }
+
+/// Extracts a static attribute from a swap.
+pub fn get_static_attribute(swap: &Swap, attribute_name: &str) -> Result<Vec<u8>, EncodingError> {
+    Ok(swap
+        .component
+        .static_attributes
+        .get(attribute_name)
+        .ok_or_else(|| {
+            EncodingError::FatalError(format!("Attribute {} not found", attribute_name))
+        })?
+        .to_vec())
+}