fix: consider wrapping scenario when getting transfer type

src/encoding/evm/strategy_encoder/strategy_encoders.rs

@@ -40,6 +40,7 @@ pub struct SingleSwapStrategyEncoder {
     permit2: Option<Permit2>,
     selector: String,
     native_address: Bytes,
+    wrapped_address: Bytes,
     router_address: Bytes,
 }
 
@@ -63,6 +64,7 @@ impl SingleSwapStrategyEncoder {
             selector,
             swap_encoder_registry,
             native_address: chain.native_token()?,
+            wrapped_address: chain.wrapped_token()?,
             router_address,
         })
     }
@@ -121,11 +123,13 @@ impl StrategyEncoder for SingleSwapStrategyEncoder {
 
         let mut grouped_protocol_data: Vec<u8> = vec![];
         for swap in grouped_swap.swaps.iter() {
-            let transfer_type = self.get_transfer_method(
+            let transfer_type = self.get_transfer_type(
                 swap.clone(),
                 solution.given_token.clone(),
                 self.native_address.clone(),
+                self.wrapped_address.clone(),
                 self.permit2.clone().is_some(),
+                wrap,
             );
 
             let encoding_context = EncodingContext {
@@ -294,11 +298,13 @@ impl StrategyEncoder for SequentialSwapStrategyEncoder {
 
             let mut grouped_protocol_data: Vec<u8> = vec![];
             for swap in grouped_swap.swaps.iter() {
-                let transfer_type = self.get_transfer_method(
+                let transfer_type = self.get_transfer_type(
                     swap.clone(),
                     solution.given_token.clone(),
                     self.native_address.clone(),
+                    self.wrapped_address.clone(),
                     self.permit2.clone().is_some(),
+                    wrap,
                 );
 
                 let encoding_context = EncodingContext {
@@ -521,11 +527,13 @@ impl StrategyEncoder for SplitSwapStrategyEncoder {
 
             let mut grouped_protocol_data: Vec<u8> = vec![];
             for swap in grouped_swap.swaps.iter() {
-                let transfer_type = self.get_transfer_method(
+                let transfer_type = self.get_transfer_type(
                     swap.clone(),
                     solution.given_token.clone(),
                     self.native_address.clone(),
+                    self.wrapped_address.clone(),
                     self.permit2.clone().is_some(),
+                    wrap,
                 );
 
                 let encoding_context = EncodingContext {

src/encoding/evm/strategy_encoder/transfer_optimizations.rs

@@ -8,22 +8,29 @@ use crate::encoding::{
 /// A trait that defines how the tokens will be transferred into the given pool given the solution.
 pub trait TransferOptimization {
     /// Returns the transfer method that should be used for the given swap and solution.
-    fn get_transfer_method(
+    fn get_transfer_type(
         &self,
         swap: Swap,
         given_token: Bytes,
         native_token: Bytes,
+        wrapped_token: Bytes,
         permit2: bool,
+        wrap: bool,
     ) -> TransferType {
         let send_funds_to_pool: bool =
             IN_TRANSFER_OPTIMIZABLE_PROTOCOLS.contains(&swap.component.protocol_system.as_str());
         let funds_expected_in_router: bool =
             PROTOCOLS_EXPECTING_FUNDS_IN_ROUTER.contains(&swap.component.protocol_system.as_str());
 
-        if (swap.token_in == given_token) && send_funds_to_pool {
+        // In the case of wrapping, check if the swap's token in is the wrapped token to
+        // determine if it's the first swap. Otherwise, compare to the given token.
+        let is_first_swap =
+            (swap.token_in == given_token) || ((swap.token_in == wrapped_token) && wrap);
+
+        if is_first_swap && send_funds_to_pool {
             if swap.token_in == native_token {
-                // Funds are already in router. Transfer from router to pool.
+                // Funds are already in router. Protocol takes care of native transfer.
-                TransferType::Transfer
+                TransferType::None
             } else if permit2 {
                 // Transfer from swapper to pool using permit2.
                 TransferType::Permit2Transfer
@@ -31,7 +38,7 @@ pub trait TransferOptimization {
                 // Transfer from swapper to pool.
                 TransferType::TransferFrom
             }
-        } else if (swap.token_in == given_token) && funds_expected_in_router {
+        } else if is_first_swap && funds_expected_in_router {
             if swap.token_in == native_token {
                 // Funds already in router. Do nothing.
                 TransferType::None
@@ -70,8 +77,13 @@ mod tests {
         Bytes::from(hex!("6b175474e89094c44da98b954eedeac495271d0f").to_vec())
     }
 
+    fn usdc() -> Bytes {
+        Bytes::from(hex!("0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48").to_vec())
+    }
+
     #[test]
     fn test_first_swap_transfer_from_permit2() {
+        // The swap token is the same as the given token, which is not the native token
         let swap = Swap {
             component: ProtocolComponent {
                 protocol_system: "uniswap_v2".to_string(),
@@ -82,12 +94,14 @@ mod tests {
             split: 0f64,
         };
         let strategy = MockStrategy {};
-        let transfer_method = strategy.get_transfer_method(swap.clone(), weth(), eth(), true);
+        let transfer_method =
+            strategy.get_transfer_type(swap.clone(), weth(), eth(), weth(), true, false);
         assert_eq!(transfer_method, TransferType::Permit2Transfer);
     }
 
     #[test]
     fn test_first_swap_transfer_from() {
+        // The swap token is the same as the given token, which is not the native token
         let swap = Swap {
             component: ProtocolComponent {
                 protocol_system: "uniswap_v2".to_string(),
@@ -98,12 +112,34 @@ mod tests {
             split: 0f64,
         };
         let strategy = MockStrategy {};
-        let transfer_method = strategy.get_transfer_method(swap.clone(), weth(), eth(), false);
+        let transfer_method =
+            strategy.get_transfer_type(swap.clone(), weth(), eth(), weth(), false, false);
         assert_eq!(transfer_method, TransferType::TransferFrom);
     }
 
     #[test]
-    fn test_first_swap_transfer() {
+    fn test_first_swap_native() {
+        // The swap token is the same as the given token, and it's the native token.
+        // No transfer action is needed.
+        let swap = Swap {
+            component: ProtocolComponent {
+                protocol_system: "uniswap_v2".to_string(),
+                ..Default::default()
+            },
+            token_in: eth(),
+            token_out: dai(),
+            split: 0f64,
+        };
+        let strategy = MockStrategy {};
+        let transfer_method =
+            strategy.get_transfer_type(swap.clone(), eth(), eth(), weth(), false, false);
+        assert_eq!(transfer_method, TransferType::None);
+    }
+
+    #[test]
+    fn test_first_swap_wrapped() {
+        // The swap token is NOT the same as the given token, but we are wrapping.
+        // Since the swap's token in is the wrapped token - this is the first swap.
         let swap = Swap {
             component: ProtocolComponent {
                 protocol_system: "uniswap_v2".to_string(),
@@ -114,7 +150,27 @@ mod tests {
             split: 0f64,
         };
         let strategy = MockStrategy {};
-        let transfer_method = strategy.get_transfer_method(swap.clone(), dai(), eth(), false);
+        let transfer_method =
+            strategy.get_transfer_type(swap.clone(), eth(), eth(), weth(), false, true);
+        assert_eq!(transfer_method, TransferType::TransferFrom);
+    }
+
+    #[test]
+    fn test_not_first_swap() {
+        // The swap token is NOT the same as the given token, and we are NOT wrapping.
+        // Thus, this is not the first swap.
+        let swap = Swap {
+            component: ProtocolComponent {
+                protocol_system: "uniswap_v2".to_string(),
+                ..Default::default()
+            },
+            token_in: usdc(),
+            token_out: dai(),
+            split: 0f64,
+        };
+        let strategy = MockStrategy {};
+        let transfer_method =
+            strategy.get_transfer_type(swap.clone(), weth(), eth(), weth(), false, false);
         assert_eq!(transfer_method, TransferType::Transfer);
     }
 }