From f873bb5187d8f0e24894c15d7ed189b27c5edcf7 Mon Sep 17 00:00:00 2001
From: TAMARA LIPOWSKI <data.lipowski@extaccount.com>
Date: Tue, 4 Feb 2025 10:01:38 -0500
Subject: [PATCH] refactor: check given/checked tokens before swap tokens

It's a more intuitive order, since these are higher-level attributes.

Also remove the validation from the trait - it doesn't have to be part of the trait.
---
 src/encoding/evm/tycho_encoder.rs | 85 ++++++++++++++++---------------
 src/encoding/tycho_encoder.rs     |  2 -
 2 files changed, 44 insertions(+), 43 deletions(-)

diff --git a/src/encoding/evm/tycho_encoder.rs b/src/encoding/evm/tycho_encoder.rs
index 44c3f9f..76fbd1d 100644
--- a/src/encoding/evm/tycho_encoder.rs
+++ b/src/encoding/evm/tycho_encoder.rs
@@ -32,6 +32,50 @@ impl<S: StrategySelector> EVMTychoEncoder<S> {
         Ok(EVMTychoEncoder { strategy_selector, signer, chain, router_address })
     }
 }
+
+impl<S: StrategySelector> EVMTychoEncoder<S> {
+    fn validate_solution(&self, solution: &Solution) -> Result<(), EncodingError> {
+        if solution.exact_out {
+            return Err(EncodingError::FatalError(
+                "Currently only exact input solutions are supported".to_string(),
+            ));
+        }
+        if solution.swaps.is_empty() {
+            return Err(EncodingError::FatalError("No swaps found in solution".to_string()));
+        }
+        if let Some(native_action) = solution.clone().native_action {
+            if native_action == NativeAction::Wrap {
+                if solution.given_token != *NATIVE_ADDRESS {
+                    return Err(EncodingError::FatalError(
+                        "ETH must be the input token in order to wrap".to_string(),
+                    ));
+                }
+                if let Some(first_swap) = solution.swaps.first() {
+                    if first_swap.token_in != *WETH_ADDRESS {
+                        return Err(EncodingError::FatalError(
+                            "WETH must be the first swap's input in order to wrap".to_string(),
+                        ));
+                    }
+                }
+            } else if native_action == NativeAction::Unwrap {
+                if solution.checked_token != *NATIVE_ADDRESS {
+                    return Err(EncodingError::FatalError(
+                        "ETH must be the output token in order to unwrap".to_string(),
+                    ));
+                }
+                if let Some(last_swap) = solution.swaps.last() {
+                    if last_swap.token_out != *WETH_ADDRESS {
+                        return Err(EncodingError::FatalError(
+                            "WETH must be the last swap's output in order to unwrap".to_string(),
+                        ));
+                    }
+                }
+            }
+        }
+        Ok(())
+    }
+}
+
 impl<S: StrategySelector> TychoEncoder<S> for EVMTychoEncoder<S> {
     fn encode_router_calldata(
         &self,
@@ -67,47 +111,6 @@ impl<S: StrategySelector> TychoEncoder<S> for EVMTychoEncoder<S> {
         }
         Ok(transactions)
     }
-
-    fn validate_solution(&self, solution: &Solution) -> Result<(), EncodingError> {
-        if solution.exact_out {
-            return Err(EncodingError::FatalError(
-                "Currently only exact input solutions are supported".to_string(),
-            ));
-        }
-        if solution.swaps.is_empty() {
-            return Err(EncodingError::FatalError("No swaps found in solution".to_string()));
-        }
-        if let Some(native_action) = solution.clone().native_action {
-            if native_action == NativeAction::Wrap {
-                if let Some(first_swap) = solution.swaps.first() {
-                    if first_swap.token_in != *WETH_ADDRESS {
-                        return Err(EncodingError::FatalError(
-                            "WETH must be the first swap's input in order to wrap".to_string(),
-                        ));
-                    }
-                }
-                if solution.given_token != *NATIVE_ADDRESS {
-                    return Err(EncodingError::FatalError(
-                        "ETH must be the input token in order to wrap".to_string(),
-                    ));
-                }
-            } else if native_action == NativeAction::Unwrap {
-                if let Some(last_swap) = solution.swaps.last() {
-                    if last_swap.token_out != *WETH_ADDRESS {
-                        return Err(EncodingError::FatalError(
-                            "WETH must be the last swap's output in order to unwrap".to_string(),
-                        ));
-                    }
-                }
-                if solution.checked_token != *NATIVE_ADDRESS {
-                    return Err(EncodingError::FatalError(
-                        "ETH must be the output token in order to unwrap".to_string(),
-                    ));
-                }
-            }
-        }
-        Ok(())
-    }
 }
 
 #[cfg(test)]
diff --git a/src/encoding/tycho_encoder.rs b/src/encoding/tycho_encoder.rs
index 0c92820..bfec250 100644
--- a/src/encoding/tycho_encoder.rs
+++ b/src/encoding/tycho_encoder.rs
@@ -10,6 +10,4 @@ pub trait TychoEncoder<S: StrategySelector> {
         &self,
         solutions: Vec<Solution>,
     ) -> Result<Vec<Transaction>, EncodingError>;
-
-    fn validate_solution(&self, solution: &Solution) -> Result<(), EncodingError>;
 }