diff --git a/src/encoding/evm/strategy_encoder/strategy_encoders.rs b/src/encoding/evm/strategy_encoder/strategy_encoders.rs index 043b9d5..2536812 100644 --- a/src/encoding/evm/strategy_encoder/strategy_encoders.rs +++ b/src/encoding/evm/strategy_encoder/strategy_encoders.rs @@ -40,6 +40,7 @@ pub struct SingleSwapStrategyEncoder { permit2: Option, 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 = 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 = 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 = 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 { diff --git a/src/encoding/evm/strategy_encoder/transfer_optimizations.rs b/src/encoding/evm/strategy_encoder/transfer_optimizations.rs index 8d63612..ff43ed1 100644 --- a/src/encoding/evm/strategy_encoder/transfer_optimizations.rs +++ b/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. - TransferType::Transfer + // Funds are already in router. Protocol takes care of native 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); } }