chore: Refactor TransferOptimization to be a struct instead of a trait

--- don't change below this line ---
ENG-4446 Took 39 minutes
This commit is contained in:
Diana Carvalho
2025-04-22 12:14:21 +01:00
parent ca349a8e11
commit 87de8bd637
2 changed files with 69 additions and 75 deletions

View File

@@ -32,17 +32,15 @@ use crate::encoding::{
/// * `permit2`: Permit2, responsible for managing permit2 operations and providing necessary
/// signatures and permit2 objects for calling the router
/// * `selector`: String, the selector for the swap function in the router contract
/// * `native_address`: Address of the chain's native token
/// * `wrapped_address`: Address of the chain's wrapped token
/// * `router_address`: Address of the router to be used to execute swaps
/// * `transfer_optimization`: TransferOptimization, responsible for optimizing the token transfers
#[derive(Clone)]
pub struct SingleSwapStrategyEncoder {
swap_encoder_registry: SwapEncoderRegistry,
permit2: Option<Permit2>,
selector: String,
native_address: Bytes,
wrapped_address: Bytes,
router_address: Bytes,
transfer_optimization: TransferOptimization,
}
impl SingleSwapStrategyEncoder {
@@ -60,13 +58,17 @@ impl SingleSwapStrategyEncoder {
"singleSwap(uint256,address,address,uint256,bool,bool,address,bytes)".to_string(),
)
};
let permit2_is_active = permit2.is_some();
Ok(Self {
permit2,
selector,
swap_encoder_registry,
native_address: chain.native_token()?,
wrapped_address: chain.wrapped_token()?,
router_address,
transfer_optimization: TransferOptimization::new(
chain.native_token()?,
chain.wrapped_token()?,
permit2_is_active,
),
})
}
@@ -80,8 +82,6 @@ impl SingleSwapStrategyEncoder {
}
}
impl TransferOptimization for SingleSwapStrategyEncoder {}
impl StrategyEncoder for SingleSwapStrategyEncoder {
fn encode_strategy(&self, solution: Solution) -> Result<(Vec<u8>, Bytes), EncodingError> {
let grouped_swaps = group_swaps(solution.clone().swaps);
@@ -127,15 +127,9 @@ 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_type(
swap.clone(),
solution.given_token.clone(),
self.native_address.clone(),
self.wrapped_address.clone(),
self.permit2.clone().is_some(),
wrap,
false,
);
let transfer_type = self
.transfer_optimization
.get_transfer_type(swap.clone(), solution.given_token.clone(), wrap, false);
let encoding_context = EncodingContext {
receiver: swap_receiver.clone(),
@@ -215,6 +209,7 @@ impl StrategyEncoder for SingleSwapStrategyEncoder {
/// * `router_address`: Address of the router to be used to execute swaps
/// * `sequential_swap_validator`: SequentialSwapValidator, responsible for checking validity of
/// sequential swap solutions
/// * `transfer_optimization`: TransferOptimization, responsible for optimizing the token transfers
#[derive(Clone)]
pub struct SequentialSwapStrategyEncoder {
swap_encoder_registry: SwapEncoderRegistry,
@@ -224,10 +219,9 @@ pub struct SequentialSwapStrategyEncoder {
native_address: Bytes,
wrapped_address: Bytes,
sequential_swap_validator: SequentialSwapValidator,
transfer_optimization: TransferOptimization,
}
impl TransferOptimization for SequentialSwapStrategyEncoder {}
impl SequentialSwapStrategyEncoder {
pub fn new(
chain: Chain,
@@ -244,6 +238,7 @@ impl SequentialSwapStrategyEncoder {
.to_string(),
)
};
let permit2_is_active = permit2.is_some();
Ok(Self {
permit2,
selector,
@@ -252,6 +247,11 @@ impl SequentialSwapStrategyEncoder {
native_address: chain.native_token()?,
wrapped_address: chain.wrapped_token()?,
sequential_swap_validator: SequentialSwapValidator,
transfer_optimization: TransferOptimization::new(
chain.native_token()?,
chain.wrapped_token()?,
permit2_is_active,
),
})
}
@@ -329,15 +329,14 @@ 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_type(
swap.clone(),
solution.given_token.clone(),
self.native_address.clone(),
self.wrapped_address.clone(),
self.permit2.clone().is_some(),
wrap,
in_between_swap_optimization,
);
let transfer_type = self
.transfer_optimization
.get_transfer_type(
swap.clone(),
solution.given_token.clone(),
wrap,
in_between_swap_optimization,
);
let encoding_context = EncodingContext {
receiver: swap_receiver.clone(),
@@ -422,6 +421,7 @@ impl StrategyEncoder for SequentialSwapStrategyEncoder {
/// * `split_swap_validator`: SplitSwapValidator, responsible for checking validity of split swap
/// solutions
/// * `router_address`: Address of the router to be used to execute swaps
/// * `transfer_optimization`: TransferOptimization, responsible for optimizing the token transfers
#[derive(Clone)]
pub struct SplitSwapStrategyEncoder {
swap_encoder_registry: SwapEncoderRegistry,
@@ -431,6 +431,7 @@ pub struct SplitSwapStrategyEncoder {
wrapped_address: Bytes,
split_swap_validator: SplitSwapValidator,
router_address: Bytes,
transfer_optimization: TransferOptimization,
}
impl SplitSwapStrategyEncoder {
@@ -449,7 +450,7 @@ impl SplitSwapStrategyEncoder {
.to_string(),
)
};
let permit2_is_active = permit2.is_some();
Ok(Self {
permit2,
selector,
@@ -458,6 +459,11 @@ impl SplitSwapStrategyEncoder {
wrapped_address: chain.wrapped_token()?,
split_swap_validator: SplitSwapValidator,
router_address,
transfer_optimization: TransferOptimization::new(
chain.native_token()?,
chain.wrapped_token()?,
permit2_is_active,
),
})
}
@@ -481,8 +487,6 @@ impl SplitSwapStrategyEncoder {
}
}
impl TransferOptimization for SplitSwapStrategyEncoder {}
impl StrategyEncoder for SplitSwapStrategyEncoder {
fn encode_strategy(&self, solution: Solution) -> Result<(Vec<u8>, Bytes), EncodingError> {
self.split_swap_validator
@@ -566,15 +570,9 @@ 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_type(
swap.clone(),
solution.given_token.clone(),
self.native_address.clone(),
self.wrapped_address.clone(),
self.permit2.clone().is_some(),
wrap,
false,
);
let transfer_type = self
.transfer_optimization
.get_transfer_type(swap.clone(), solution.given_token.clone(), wrap, false);
let encoding_context = EncodingContext {
receiver: swap_receiver.clone(),

View File

@@ -5,17 +5,23 @@ use crate::encoding::{
models::{Swap, TransferType},
};
/// A trait that defines how the tokens will be transferred into the given pool given the solution.
pub trait TransferOptimization {
/// A struct that defines how the tokens will be transferred into the given pool given the solution.
#[derive(Clone)]
pub struct TransferOptimization {
native_token: Bytes,
wrapped_token: Bytes,
permit2: bool,
}
impl TransferOptimization {
pub fn new(native_token: Bytes, wrapped_token: Bytes, permit2: bool) -> Self {
TransferOptimization { native_token, wrapped_token, permit2 }
}
/// Returns the transfer method that should be used for the given swap and solution.
#[allow(clippy::too_many_arguments)]
fn get_transfer_type(
pub fn get_transfer_type(
&self,
swap: Swap,
given_token: Bytes,
native_token: Bytes,
wrapped_token: Bytes,
permit2: bool,
wrap: bool,
in_between_swap_optimization: bool,
) -> TransferType {
@@ -24,22 +30,22 @@ pub trait TransferOptimization {
let is_first_swap = swap.token_in == given_token;
if swap.token_in == native_token {
if swap.token_in == self.native_token {
// Funds are already in router. All protocols currently take care of native transfers.
TransferType::None
} else if (swap.token_in == wrapped_token) && wrap {
} else if (swap.token_in == self.wrapped_token) && wrap {
// Wrapping already happened in the router so we can just use a normal transfer.
TransferType::TransferToProtocol
} else if is_first_swap {
if in_transfer_required {
if permit2 {
if self.permit2 {
// Transfer from swapper to pool using permit2.
TransferType::TransferPermit2ToProtocol
} else {
// Transfer from swapper to pool.
TransferType::TransferFromToProtocol
}
} else if permit2 {
} else if self.permit2 {
// Transfer from swapper to router using permit2.
TransferType::TransferPermit2ToRouter
} else {
@@ -63,9 +69,6 @@ mod tests {
use super::*;
struct MockStrategy {}
impl TransferOptimization for MockStrategy {}
fn weth() -> Bytes {
Bytes::from(hex!("c02aaa39b223fe8d0a0e5c4f27ead9083c756cc2").to_vec())
}
@@ -94,9 +97,8 @@ mod tests {
token_out: dai(),
split: 0f64,
};
let strategy = MockStrategy {};
let transfer_method =
strategy.get_transfer_type(swap.clone(), weth(), eth(), weth(), true, false, false);
let optimization = TransferOptimization::new(eth(), weth(), true);
let transfer_method = optimization.get_transfer_type(swap.clone(), weth(), false, false);
assert_eq!(transfer_method, TransferType::TransferPermit2ToProtocol);
}
@@ -112,9 +114,8 @@ mod tests {
token_out: dai(),
split: 0f64,
};
let strategy = MockStrategy {};
let transfer_method =
strategy.get_transfer_type(swap.clone(), weth(), eth(), weth(), false, false, false);
let optimization = TransferOptimization::new(eth(), weth(), false);
let transfer_method = optimization.get_transfer_type(swap.clone(), weth(), false, false);
assert_eq!(transfer_method, TransferType::TransferFromToProtocol);
}
@@ -131,9 +132,8 @@ mod tests {
token_out: dai(),
split: 0f64,
};
let strategy = MockStrategy {};
let transfer_method =
strategy.get_transfer_type(swap.clone(), eth(), eth(), weth(), false, false, false);
let optimization = TransferOptimization::new(eth(), weth(), false);
let transfer_method = optimization.get_transfer_type(swap.clone(), eth(), false, false);
assert_eq!(transfer_method, TransferType::None);
}
@@ -150,9 +150,8 @@ mod tests {
token_out: dai(),
split: 0f64,
};
let strategy = MockStrategy {};
let transfer_method =
strategy.get_transfer_type(swap.clone(), eth(), eth(), weth(), false, true, false);
let optimization = TransferOptimization::new(eth(), weth(), false);
let transfer_method = optimization.get_transfer_type(swap.clone(), eth(), true, false);
assert_eq!(transfer_method, TransferType::TransferToProtocol);
}
@@ -169,9 +168,8 @@ mod tests {
token_out: dai(),
split: 0f64,
};
let strategy = MockStrategy {};
let transfer_method =
strategy.get_transfer_type(swap.clone(), weth(), eth(), weth(), false, false, false);
let optimization = TransferOptimization::new(eth(), weth(), false);
let transfer_method = optimization.get_transfer_type(swap.clone(), weth(), false, false);
assert_eq!(transfer_method, TransferType::TransferToProtocol);
}
@@ -188,9 +186,8 @@ mod tests {
token_out: dai(),
split: 0f64,
};
let strategy = MockStrategy {};
let transfer_method =
strategy.get_transfer_type(swap.clone(), weth(), eth(), weth(), false, false, false);
let optimization = TransferOptimization::new(eth(), weth(), false);
let transfer_method = optimization.get_transfer_type(swap.clone(), weth(), false, false);
assert_eq!(transfer_method, TransferType::None);
}
@@ -207,9 +204,8 @@ mod tests {
token_out: dai(),
split: 0f64,
};
let strategy = MockStrategy {};
let transfer_method =
strategy.get_transfer_type(swap.clone(), weth(), eth(), weth(), false, false, true);
let optimization = TransferOptimization::new(eth(), weth(), false);
let transfer_method = optimization.get_transfer_type(swap.clone(), weth(), false, true);
assert_eq!(transfer_method, TransferType::None);
}
}