Add wrapping logic to the scripts

mira-amm · Sep 5, 2024 · 17d388a · 17d388a
1 parent 18db31c
commit 17d388a
Show file tree

Hide file tree

Showing 6 changed files with 58 additions and 7 deletions.
diff --git a/libraries/utils/src/blockchain_utils.sw b/libraries/utils/src/blockchain_utils.sw
@@ -1,7 +1,7 @@
 library;
 
 use std::block::height;
-use interfaces::data_structures::PoolId;
+use interfaces::{asset_wrapper::AssetWrapper, data_structures::PoolId};
 use std::hash::*;
 
 /// Validates that the provided deadline hasn't passed yet
@@ -18,3 +18,41 @@ pub fn get_lp_asset(contract_id: ContractId, pool_id: PoolId) -> (b256, AssetId)
 pub fn is_stable(pool_id: PoolId) -> bool {
     pool_id.2
 }
+
+/// Checks if the provided asset is a wrapped asset and wraps the provided amount of it in that case
+pub fn wrap_if_needed(
+    asset_wrapper_contract: ContractId,
+    asset_id: AssetId,
+    amount: u64,
+) {
+    let asset_wrapper = abi(AssetWrapper, asset_wrapper_contract.into());
+    match asset_wrapper.get_underlying_asset(asset_id) {
+        Some(underlying_asset) => {
+            asset_wrapper
+                .wrap {
+                    asset_id: underlying_asset.into(),
+                    coins: amount,
+                }();
+        },
+        None => {}
+    };
+}
+
+/// Checks if the provided asset is a wrapped asset and unwraps the provided amount of it in that case
+pub fn unwrap_if_needed(
+    asset_wrapper_contract: ContractId,
+    asset_id: AssetId,
+    amount: u64,
+) {
+    let asset_wrapper = abi(AssetWrapper, asset_wrapper_contract.into());
+    match asset_wrapper.get_underlying_asset(asset_id) {
+        Some(_) => {
+            asset_wrapper
+                .unwrap {
+                    asset_id: asset_id.into(),
+                    coins: amount,
+                }();
+        },
+        None => {}
+    };
+}
diff --git a/scripts/add_liquidity_script/src/main.sw b/scripts/add_liquidity_script/src/main.sw
@@ -2,11 +2,12 @@ script;
 
 use interfaces::{data_structures::{Asset, PoolId}, mira_amm::MiraAMM};
 use math::pool_math::get_deposit_amounts;
-use utils::blockchain_utils::check_deadline;
+use utils::blockchain_utils::{check_deadline, wrap_if_needed};
 use std::asset::transfer;
 
 configurable {
     AMM_CONTRACT_ID: ContractId = ContractId::from(0x0000000000000000000000000000000000000000000000000000000000000000),
+    ASSET_WRAPPER_CONTRACT_ID: ContractId = ContractId::from(0x0000000000000000000000000000000000000000000000000000000000000000),
 }
 
 fn main(
@@ -36,6 +37,8 @@ fn main(
             .reserve_1,
     );
 
+    wrap_if_needed(ASSET_WRAPPER_CONTRACT_ID, pool_id.0, amount_0);
+    wrap_if_needed(ASSET_WRAPPER_CONTRACT_ID, pool_id.1, amount_1);
     transfer(Identity::ContractId(AMM_CONTRACT_ID), pool_id.0, amount_0);
     transfer(Identity::ContractId(AMM_CONTRACT_ID), pool_id.1, amount_1);
 

diff --git a/scripts/create_pool_and_add_liquidity_script/src/main.sw b/scripts/create_pool_and_add_liquidity_script/src/main.sw
@@ -2,10 +2,11 @@ script;
 
 use interfaces::{data_structures::{Asset, PoolId}, mira_amm::MiraAMM};
 use std::asset::transfer;
-use utils::blockchain_utils::check_deadline;
+use utils::blockchain_utils::{check_deadline, wrap_if_needed};
 
 configurable {
     AMM_CONTRACT_ID: ContractId = ContractId::from(0x0000000000000000000000000000000000000000000000000000000000000000),
+    ASSET_WRAPPER_CONTRACT_ID: ContractId = ContractId::from(0x0000000000000000000000000000000000000000000000000000000000000000),
 }
 
 fn main(
@@ -30,6 +31,8 @@ fn main(
         is_stable,
     );
 
+    wrap_if_needed(ASSET_WRAPPER_CONTRACT_ID, pool_id.0, amount_0_desired);
+    wrap_if_needed(ASSET_WRAPPER_CONTRACT_ID, pool_id.1, amount_1_desired);
     transfer(
         Identity::ContractId(AMM_CONTRACT_ID),
         pool_id.0,

diff --git a/scripts/remove_liquidity_script/src/main.sw b/scripts/remove_liquidity_script/src/main.sw
@@ -2,11 +2,11 @@ script;
 
 use interfaces::{data_structures::{Asset, PoolId}, mira_amm::MiraAMM};
 use math::pool_math::get_deposit_amounts;
-use utils::blockchain_utils::{check_deadline, get_lp_asset};
-use std::asset::transfer;
+use utils::blockchain_utils::{check_deadline, get_lp_asset, unwrap_if_needed};
 
 configurable {
     AMM_CONTRACT_ID: ContractId = ContractId::from(0x0000000000000000000000000000000000000000000000000000000000000000),
+    ASSET_WRAPPER_CONTRACT_ID: ContractId = ContractId::from(0x0000000000000000000000000000000000000000000000000000000000000000),
 }
 
 fn main(
@@ -28,5 +28,8 @@ fn main(
 
     require(amount_0 >= amount_0_min, "Insufficient amount");
     require(amount_1 >= amount_1_min, "Insufficient amount");
+
+    unwrap_if_needed(ASSET_WRAPPER_CONTRACT_ID, pool_id.0, amount_0);
+    unwrap_if_needed(ASSET_WRAPPER_CONTRACT_ID, pool_id.1, amount_1);
     (amount_0, amount_1)
 }
diff --git a/scripts/swap_exact_input_script/src/main.sw b/scripts/swap_exact_input_script/src/main.sw
@@ -2,11 +2,12 @@ script;
 
 use interfaces::{data_structures::PoolId, mira_amm::MiraAMM};
 use math::pool_math::get_amounts_out;
-use utils::blockchain_utils::check_deadline;
+use utils::blockchain_utils::{check_deadline, wrap_if_needed};
 use std::{asset::transfer, bytes::Bytes};
 
 configurable {
     AMM_CONTRACT_ID: ContractId = ContractId::from(0x0000000000000000000000000000000000000000000000000000000000000000),
+    ASSET_WRAPPER_CONTRACT_ID: ContractId = ContractId::from(0x0000000000000000000000000000000000000000000000000000000000000000),
 }
 
 fn main(
@@ -26,6 +27,7 @@ fn main(
         "Insufficient output amount",
     );
 
+    wrap_if_needed(ASSET_WRAPPER_CONTRACT_ID, asset_in, amount_in);
     transfer(Identity::ContractId(AMM_CONTRACT_ID), asset_in, amount_in);
     let amm = abi(MiraAMM, AMM_CONTRACT_ID.into());
 

diff --git a/scripts/swap_exact_output_script/src/main.sw b/scripts/swap_exact_output_script/src/main.sw
@@ -2,11 +2,12 @@ script;
 
 use interfaces::{data_structures::PoolId, mira_amm::MiraAMM};
 use math::pool_math::get_amounts_in;
-use utils::blockchain_utils::check_deadline;
+use utils::blockchain_utils::{check_deadline, wrap_if_needed};
 use std::{asset::transfer, bytes::Bytes};
 
 configurable {
     AMM_CONTRACT_ID: ContractId = ContractId::from(0x0000000000000000000000000000000000000000000000000000000000000000),
+    ASSET_WRAPPER_CONTRACT_ID: ContractId = ContractId::from(0x0000000000000000000000000000000000000000000000000000000000000000),
 }
 
 fn main(
@@ -23,6 +24,7 @@ fn main(
     let (first_amount_in, first_asset) = amounts_in.get(amounts_in.len() - 1).unwrap();
     require(first_amount_in <= amount_in_max, "Exceeding input amount");
 
+    wrap_if_needed(ASSET_WRAPPER_CONTRACT_ID, first_asset, first_amount_in);
     transfer(
         Identity::ContractId(AMM_CONTRACT_ID),
         first_asset,