From fa462ee9f3cce3d339f0bc1645b7ce8bd6d42cf0 Mon Sep 17 00:00:00 2001
From: Diana Carvalho <diana@propellerheads.xyz>
Date: Mon, 13 Jan 2025 17:11:21 +0000
Subject: [PATCH] feat: Handle native actions

---
 src/encoding/models.rs           |  7 +++++++
 src/encoding/router_encoder.rs   |  2 --
 src/encoding/strategy_encoder.rs | 35 +++++++++++++++++---------------
 3 files changed, 26 insertions(+), 18 deletions(-)

diff --git a/src/encoding/models.rs b/src/encoding/models.rs
index 7076d49..5086ca4 100644
--- a/src/encoding/models.rs
+++ b/src/encoding/models.rs
@@ -8,6 +8,12 @@ pub struct Solution {
     pub router_address: Option<Address>,
 }
 
+#[derive(Clone)]
+pub enum NativeAction {
+    Wrap,
+    Unwrap,
+}
+
 pub struct Order {
     /// True if the order is an exact output order.
     pub exact_out: bool,
@@ -30,6 +36,7 @@ pub struct Order {
 
     pub slippage: f64,
     pub min_checked_amount: Option<BigUint>,
+    pub native_action: Option<NativeAction>,
 }
 
 #[derive(Clone)]
diff --git a/src/encoding/router_encoder.rs b/src/encoding/router_encoder.rs
index 958c140..011d100 100644
--- a/src/encoding/router_encoder.rs
+++ b/src/encoding/router_encoder.rs
@@ -28,8 +28,6 @@ impl RouterEncoder {
         let encode_for_batch_execute = solution.orders.len() > 1;
         for order in solution.orders {
             let strategy = self.get_strategy(&order);
-            // TODO: handle native action??
-
             let contract_interaction = strategy.encode_strategy(
                 order,
                 if solution.router_address.is_some() {
diff --git a/src/encoding/strategy_encoder.rs b/src/encoding/strategy_encoder.rs
index dafd1c4..a33f23b 100644
--- a/src/encoding/strategy_encoder.rs
+++ b/src/encoding/strategy_encoder.rs
@@ -4,7 +4,7 @@ use anyhow::Error;
 use num_bigint::BigUint;
 use std::cmp::min;
 
-use crate::encoding::models::{ActionType, EncodingContext, Order};
+use crate::encoding::models::{ActionType, EncodingContext, NativeAction, Order};
 use crate::encoding::swap_encoder::{get_swap_encoder, get_swap_executor_address};
 use crate::encoding::utils::{biguint_to_u256, bytes_to_address, encode_input, ple_encode};
 
@@ -94,24 +94,27 @@ impl StrategyEncoder for SequentialExactInStrategyEncoder {
             )
         };
         let encoded_swaps = ple_encode(swaps);
+
+        let (mut unwrap, mut wrap) = (false, false);
+        if order.native_action.is_some() {
+            match order.native_action.unwrap() {
+                NativeAction::Wrap => wrap = true,
+                NativeAction::Unwrap => unwrap = true,
+            }
+        }
+        let method_calldata = (
+            wrap,
+            unwrap,
+            biguint_to_u256(&order.given_amount),
+            biguint_to_u256(&min_checked_amount),
+            encoded_swaps,
+        )
+            .abi_encode();
         if encode_for_batch_execute {
-            let args = (
-                action_type as u16,
-                biguint_to_u256(&order.given_amount),
-                biguint_to_u256(&min_checked_amount),
-                encoded_swaps,
-            );
+            let args = (action_type as u16, method_calldata);
             Ok(args.abi_encode())
         } else {
-            Ok(encode_input(
-                selector,
-                (
-                    biguint_to_u256(&order.given_amount),
-                    biguint_to_u256(&min_checked_amount),
-                    encoded_swaps,
-                )
-                    .abi_encode(),
-            ))
+            Ok(encode_input(selector, method_calldata))
         }
     }
 }