From 6e6787582169aa99d9530f35dee69ce7218172b9 Mon Sep 17 00:00:00 2001
From: Diana Carvalho <diana@propellerheads.xyz>
Date: Tue, 14 Jan 2025 10:45:34 +0000
Subject: [PATCH] feat: Add StrategySelector

This way the user could be able to extend this easily
---
 src/encoding/mod.rs               |  1 +
 src/encoding/permit2.rs           |  2 +-
 src/encoding/router_encoder.rs    | 31 +++++++++++--------------------
 src/encoding/strategy_encoder.rs  |  2 --
 src/encoding/strategy_selector.rs | 23 +++++++++++++++++++++++
 5 files changed, 36 insertions(+), 23 deletions(-)
 create mode 100644 src/encoding/strategy_selector.rs

diff --git a/src/encoding/mod.rs b/src/encoding/mod.rs
index 5cbe4b8..8631ba1 100644
--- a/src/encoding/mod.rs
+++ b/src/encoding/mod.rs
@@ -3,5 +3,6 @@ mod models;
 mod permit2;
 mod router_encoder;
 mod strategy_encoder;
+mod strategy_selector;
 mod swap_encoder;
 mod utils;
diff --git a/src/encoding/permit2.rs b/src/encoding/permit2.rs
index 3a68bbf..1a527be 100644
--- a/src/encoding/permit2.rs
+++ b/src/encoding/permit2.rs
@@ -1,4 +1,4 @@
-use alloy_primitives::{Address, U256};
+use alloy_primitives::U256;
 use num_bigint::BigUint;
 use std::str::FromStr;
 use tycho_core::Bytes;
diff --git a/src/encoding/router_encoder.rs b/src/encoding/router_encoder.rs
index ec9912d..b13d083 100644
--- a/src/encoding/router_encoder.rs
+++ b/src/encoding/router_encoder.rs
@@ -1,23 +1,24 @@
-use crate::encoding::models::{Order, Solution, PROPELLER_ROUTER_ADDRESS};
+use crate::encoding::models::{Solution, PROPELLER_ROUTER_ADDRESS};
 use crate::encoding::permit2::{Permit2, PermitRequest};
-use crate::encoding::strategy_encoder::{
-    SequentialExactInStrategyEncoder, SingleSwapStrategyEncoder, SlipSwapStrategyEncoder,
-    StrategyEncoder,
-};
+use crate::encoding::strategy_encoder::StrategyEncoder;
+use crate::encoding::strategy_selector::StrategySelector;
 use crate::encoding::utils::{encode_input, ple_encode};
 use alloy_sol_types::SolValue;
 use anyhow::Error;
-use std::str::FromStr;
 
-struct RouterEncoder {}
-
-impl RouterEncoder {
+struct RouterEncoder<S: StrategySelector> {
+    strategy_selector: S,
+}
+impl<S: StrategySelector> RouterEncoder<S> {
+    pub fn new(strategy_selector: S) -> Self {
+        RouterEncoder { strategy_selector }
+    }
     pub fn encode_router_calldata(&self, solution: Solution) -> Result<Vec<u8>, Error> {
         let permit_calldata = self.handle_approvals(&solution)?; // TODO: where should we append this?
         let mut calldata_list: Vec<Vec<u8>> = Vec::new();
         let encode_for_batch_execute = solution.orders.len() > 1;
         for order in solution.orders {
-            let strategy = self.get_strategy(&order);
+            let strategy = self.strategy_selector.select_strategy(&order);
             let contract_interaction = strategy.encode_strategy(order, encode_for_batch_execute)?;
             calldata_list.push(contract_interaction);
         }
@@ -44,14 +45,4 @@ impl RouterEncoder {
         }
         Ok(Permit2::new().encode_permit(permits))
     }
-
-    fn get_strategy(&self, order: &Order) -> &dyn StrategyEncoder {
-        if order.swaps.len() == 1 {
-            &SingleSwapStrategyEncoder {}
-        } else if order.swaps.iter().all(|s| s.split == 0.0) {
-            &SequentialExactInStrategyEncoder {}
-        } else {
-            &SlipSwapStrategyEncoder {}
-        }
-    }
 }
diff --git a/src/encoding/strategy_encoder.rs b/src/encoding/strategy_encoder.rs
index 63593ad..c3328d2 100644
--- a/src/encoding/strategy_encoder.rs
+++ b/src/encoding/strategy_encoder.rs
@@ -1,8 +1,6 @@
-use alloy_primitives::Address;
 use alloy_sol_types::SolValue;
 use anyhow::Error;
 use num_bigint::BigUint;
-use std::cmp::min;
 
 use crate::encoding::models::{
     ActionType, EncodingContext, NativeAction, Order, PROPELLER_ROUTER_ADDRESS,
diff --git a/src/encoding/strategy_selector.rs b/src/encoding/strategy_selector.rs
new file mode 100644
index 0000000..6faa47a
--- /dev/null
+++ b/src/encoding/strategy_selector.rs
@@ -0,0 +1,23 @@
+use crate::encoding::models::Order;
+use crate::encoding::strategy_encoder::{
+    SequentialExactInStrategyEncoder, SingleSwapStrategyEncoder, SlipSwapStrategyEncoder,
+    StrategyEncoder,
+};
+
+pub trait StrategySelector {
+    fn select_strategy(&self, order: &Order) -> Box<dyn StrategyEncoder>;
+}
+
+pub struct DefaultStrategySelector;
+
+impl StrategySelector for DefaultStrategySelector {
+    fn select_strategy(&self, order: &Order) -> Box<dyn StrategyEncoder> {
+        if order.swaps.len() == 1 {
+            Box::new(SingleSwapStrategyEncoder {})
+        } else if order.swaps.iter().all(|s| s.split == 0.0) {
+            Box::new(SequentialExactInStrategyEncoder {})
+        } else {
+            Box::new(SlipSwapStrategyEncoder {})
+        }
+    }
+}