refactor(solana-contracts): refactor equalize related fns into common

swim-io · Oct 6, 2022 · 49e04a7 · 49e04a7
1 parent b88f868
commit 49e04a7
Show file tree

Hide file tree

Showing 9 changed files with 102 additions and 225 deletions.
diff --git a/packages/solana-contracts/programs/two-pool/src/common.rs b/packages/solana-contracts/programs/two-pool/src/common.rs
@@ -1,4 +1,8 @@
-use {arrayvec::ArrayVec, std::fmt::Debug};
+use {
+    crate::{decimal::U128, TOKEN_COUNT},
+    arrayvec::ArrayVec,
+    std::fmt::Debug,
+};
 
 //final unwraps are safe because we know that there is enough capacity
 
@@ -11,3 +15,55 @@ pub fn create_result_array<T: Debug, E: Debug, const SIZE: usize>(
 ) -> Result<[T; SIZE], E> {
     Ok((0..SIZE).into_iter().map(closure).collect::<Result<ArrayVec<_, SIZE>, _>>()?.into_inner().unwrap())
 }
+
+pub fn to_equalized(value: u64, equalizer: u8) -> U128 {
+    if equalizer > 0 {
+        U128::from(value) * U128::ten_to_the(equalizer)
+    } else {
+        U128::from(value)
+    }
+}
+
+pub fn from_equalized(value: U128, equalizer: u8) -> u64 {
+    if equalizer > 0 {
+        ((value + U128::ten_to_the(equalizer - 1) * 5u64) / U128::ten_to_the(equalizer)).as_u64()
+    } else {
+        value.as_u64()
+    }
+}
+
+pub fn array_equalize(amounts: [u64; TOKEN_COUNT], equalizers: [u8; TOKEN_COUNT]) -> [U128; TOKEN_COUNT] {
+    amounts
+        .iter()
+        .zip(equalizers.iter())
+        .map(|(&amount, &equalizer)| to_equalized(amount, equalizer))
+        .collect::<Vec<_>>()
+        .as_slice()
+        .try_into()
+        .unwrap()
+}
+
+/// `result_from_equalized` takes in a user's amount, the user's equalizer, the governance mint amount,
+/// the lp decimal equalizer, and the latest depth, and returns the user's amount, the governance mint
+/// amount, and the latest depth
+///
+/// Arguments:
+///
+/// * `user_amount`: The amount of tokens the user is staking
+/// * `user_equalizer`: The equalizer of the user's token.
+/// * `governance_mint_amount`: The amount of governance tokens that will be minted to the user.
+/// * `lp_decimal_equalizer`: The equalizer for the LP token. should always be pool_state.lp_decimal_equalizer
+/// * `latest_depth`: The amount of liquidity in the pool.
+pub fn result_from_equalized(
+    user_amount: U128,
+    user_equalizer: u8,
+    governance_mint_amount: U128,
+    lp_decimal_equalizer: u8,
+    latest_depth: U128,
+) -> (u64, u64, u128) {
+    (
+        from_equalized(user_amount, user_equalizer),
+        from_equalized(governance_mint_amount, lp_decimal_equalizer),
+        latest_depth.as_u128(),
+    )
+}
diff --git a/packages/solana-contracts/programs/two-pool/src/instructions/add.rs b/packages/solana-contracts/programs/two-pool/src/instructions/add.rs
@@ -1,8 +1,6 @@
 use {
-    crate::{decimal::U128, error::*, gen_pool_signer_seeds, invariant::Invariant, TwoPool, TOKEN_COUNT},
-    anchor_lang::{
-        prelude::*,
-    },
+    crate::{common, decimal::U128, error::*, gen_pool_signer_seeds, invariant::Invariant, TwoPool, TOKEN_COUNT},
+    anchor_lang::prelude::*,
     anchor_spl::{
         associated_token::get_associated_token_address,
         token::{self, Mint, Token, TokenAccount},
@@ -125,33 +123,25 @@ pub fn handle_add(
     let lp_total_supply = ctx.accounts.lp_mint.supply;
     //initial add to pool must add all tokens
     if lp_total_supply == 0 {
-        for i in 0..TOKEN_COUNT {
-            require_gt!(input_amounts[i], 0u64, PoolError::InitialAddRequiresAllTokens);
-        }
+        require!(input_amounts.iter().all(|&x| x > 0), PoolError::InitialAddRequiresAllTokens);
     }
 
     let pool = &ctx.accounts.pool;
     let user_token_accounts = [&ctx.accounts.user_token_account_0, &ctx.accounts.user_token_account_1];
     let pool_token_accounts = [&ctx.accounts.pool_token_account_0, &ctx.accounts.pool_token_account_1];
-    // let pool_balances = pool_token_accounts
-    //   .iter()
-    //   .map(|account| account.amount)
-    //   .collect::<Vec<_>>()
-    //   .as_slice()
-    //   .try_into().unwrap();
     let pool_balances = [ctx.accounts.pool_token_account_0.amount, ctx.accounts.pool_token_account_1.amount];
     let current_ts = Clock::get()?.unix_timestamp;
     require_gt!(current_ts, 0i64, PoolError::InvalidTimestamp);
     let (user_amount, governance_mint_amount, latest_depth) = Invariant::<TOKEN_COUNT>::add(
-        &array_equalize(input_amounts, pool.token_decimal_equalizers),
-        &array_equalize(pool_balances, pool.token_decimal_equalizers),
+        &common::array_equalize(input_amounts, pool.token_decimal_equalizers),
+        &common::array_equalize(pool_balances, pool.token_decimal_equalizers),
         pool.amp_factor.get(current_ts),
         pool.lp_fee.get(),
         pool.governance_fee.get(),
-        to_equalized(lp_total_supply, pool.lp_decimal_equalizer),
+        common::to_equalized(lp_total_supply, pool.lp_decimal_equalizer),
         pool.previous_depth.into(),
     )?;
-    let (mint_amount, governance_mint_amount, latest_depth) = result_from_equalized(
+    let (mint_amount, governance_mint_amount, latest_depth) = common::result_from_equalized(
         user_amount,
         pool.lp_decimal_equalizer,
         governance_mint_amount,
@@ -160,20 +150,22 @@ pub fn handle_add(
     );
     require_gte!(mint_amount, minimum_mint_amount, PoolError::OutsideSpecifiedLimits);
     let mut token_accounts = zip(user_token_accounts.into_iter(), pool_token_accounts.into_iter());
-    for i in 0..TOKEN_COUNT {
+    for input_amount in input_amounts.iter().take(TOKEN_COUNT) {
         let (user_token_account, pool_token_account) = token_accounts.next().unwrap();
-        if input_amounts[i] > 0 {
+        if *input_amount > 0u64 {
+            let pool_token_account = pool_token_account.to_account_info();
+            let user_token_account = user_token_account.to_account_info();
+            let user_transfer_authority = ctx.accounts.user_transfer_authority.to_account_info();
             token::transfer(
                 CpiContext::new(
                     ctx.accounts.token_program.to_account_info(),
                     token::Transfer {
-                        // source
-                        from: user_token_account.to_account_info(),
-                        to: pool_token_account.to_account_info(),
-                        authority: ctx.accounts.user_transfer_authority.to_account_info(),
+                        from: user_token_account.clone(),
+                        to: pool_token_account.clone(),
+                        authority: user_transfer_authority.clone(),
                     },
                 ),
-                input_amounts[i],
+                *input_amount,
             )?;
         }
     }
@@ -211,187 +203,4 @@ pub fn handle_add(
     let pool_state = &mut ctx.accounts.pool;
     pool_state.previous_depth = latest_depth;
     Ok(mint_amount)
-    // msg!("add_and_wormhole_transfer");
-    // handle_approve(&ctx, &pool_add_params)?;
-    //
-    // let mint_amount = handle_add_liquidity(&ctx, pool_add_params)?;
-    // msg!("finished add_liquidity. mint_amount: {}", mint_amount);
-    //
-    // revoke(&ctx)
-    //
-    // // Ok(mint_amount)
-}
-
-pub fn to_equalized(value: u64, equalizer: u8) -> U128 {
-    if equalizer > 0 {
-        U128::from(value) * U128::ten_to_the(equalizer)
-    } else {
-        U128::from(value)
-    }
 }
-
-pub fn from_equalized(value: U128, equalizer: u8) -> u64 {
-    if equalizer > 0 {
-        ((value + U128::ten_to_the(equalizer - 1) * 5u64) / U128::ten_to_the(equalizer)).as_u64()
-    } else {
-        value.as_u64()
-    }
-}
-
-pub fn array_equalize(amounts: [u64; TOKEN_COUNT], equalizers: [u8; TOKEN_COUNT]) -> [U128; TOKEN_COUNT] {
-    amounts
-        .iter()
-        .zip(equalizers.iter())
-        .map(|(&amount, &equalizer)| to_equalized(amount, equalizer))
-        .collect::<Vec<_>>()
-        .as_slice()
-        .try_into()
-        .unwrap()
-}
-
-/// `result_from_equalized` takes in a user's amount, the user's equalizer, the governance mint amount,
-/// the lp decimal equalizer, and the latest depth, and returns the user's amount, the governance mint
-/// amount, and the latest depth
-///
-/// Arguments:
-///
-/// * `user_amount`: The amount of tokens the user is staking
-/// * `user_equalizer`: The equalizer of the user's token.
-/// * `governance_mint_amount`: The amount of governance tokens that will be minted to the user.
-/// * `lp_decimal_equalizer`: The equalizer for the LP token. should always be pool_state.lp_decimal_equalizer
-/// * `latest_depth`: The amount of liquidity in the pool.
-pub fn result_from_equalized(
-    user_amount: U128,
-    user_equalizer: u8,
-    governance_mint_amount: U128,
-    lp_decimal_equalizer: u8,
-    latest_depth: U128,
-) -> (u64, u64, u128) {
-    (
-        from_equalized(user_amount, user_equalizer),
-        from_equalized(governance_mint_amount, lp_decimal_equalizer),
-        latest_depth.as_u128(),
-    )
-}
-//
-// pub fn handle_approve(
-// 	ctx: &Context<Add>,
-// 	pool_add_params: &AddParams,
-// ) -> Result<()> {
-// 	msg!("[handle_approve]: approve");
-// 	token::approve(
-// 		CpiContext::new(
-// 			ctx.accounts.token_program.to_account_info(),
-// 			token::Approve {
-// 				// source
-// 				to: ctx.accounts.user_token_account_0.to_account_info(),
-// 				delegate: ctx.accounts.pool_auth.to_account_info(),
-// 				authority: ctx.accounts.payer.to_account_info(),
-// 			},
-// 		),
-// 		pool_add_params.input_amounts[0],
-// 	)?;
-//
-// 	token::approve(
-// 		CpiContext::new(
-// 			ctx.accounts.token_program.to_account_info(),
-// 			token::Approve {
-// 				// source
-// 				to: ctx.accounts.user_token_account_1.to_account_info(),
-// 				delegate: ctx.accounts.pool_auth.to_account_info(),
-// 				authority: ctx.accounts.payer.to_account_info(),
-// 			},
-// 		),
-// 		pool_add_params.input_amounts[1],
-// 	)?;
-// 	msg!("[handle_approve]: finished approves");
-// 	Ok(())
-// }
-//
-// pub fn handle_add_liquidity(
-// 	ctx: &Context<Add>,
-// 	pool_add_params: AddParams
-// ) -> Result<u64> {
-// 	let pool_program = &ctx.accounts.pool_program;
-// 	let pool = &ctx.accounts.pool_state;
-// 	let pool_auth = &ctx.accounts.pool_auth;
-// 	let pool_token_account_0 = &ctx.accounts.pool_token_account_0;
-// 	let pool_token_account_1 = &ctx.accounts.pool_token_account_1;
-// 	let lp_mint = &ctx.accounts.lp_mint;
-// 	let governance_fee = &ctx.accounts.governance_fee;
-// 	// let user_transfer_auth = &ctx.accounts.user_transfer_authority;
-// 	let user_token_account_0 = &ctx.accounts.user_token_account_0;
-// 	let user_token_account_1 = &ctx.accounts.user_token_account_1;
-// 	let user_lp_token_account = &ctx.accounts.user_lp_token_account;
-// 	// let token_program = &self.token_program;
-//
-// 	let add_defi_ix = two_pool::instruction::DeFiInstruction::Add{
-// 		input_amounts: pool_add_params.input_amounts,
-// 		minimum_mint_amount: pool_add_params.minimum_mint_amount
-// 	};
-// 	let add_ix = two_pool::instruction::create_defi_ix(
-// 		add_defi_ix,
-// 		&pool_program.key(),
-// 		&pool.key(),
-// 		&pool_auth.key(),
-// 		&[
-// 			pool_token_account_0.key(),
-// 			pool_token_account_1.key()
-// 		],
-// 		&lp_mint.key(),
-// 		&governance_fee.key(),
-// 		&pool_auth.key(),
-// 		// &user_transfer_auth.key(),
-// 		&[
-// 			user_token_account_0.key(),
-// 			user_token_account_1.key()
-// 		],
-// 		Some(&user_lp_token_account.key()),
-// 	)?;
-// 	invoke(
-// 		&add_ix,
-// 		&ctx.accounts.to_account_infos(),
-// 	)?;
-// 	let (program_id, data) = get_return_data().unwrap();
-// 	//.unwrap_or_else(PropellerError::InvalidCpiReturnValue);
-// 	require_keys_eq!(program_id, ctx.accounts.pool_program.key(), PropellerError::InvalidCpiReturnProgramId);
-// 	Ok(u64::try_from_slice(&data).map_err(|_| PropellerError::InvalidCpiReturnValue)?)
-// }
-//
-// pub fn revoke(ctx: &Context<Add>) -> Result<()> {
-// 	msg!("[revoke]: revoke");
-// 	token::revoke(
-// 		CpiContext::new(
-// 			ctx.accounts.token_program.to_account_info(),
-// 			token::Revoke {
-// 				// source
-// 				source: ctx.accounts.user_token_account_0.to_account_info(),
-// 				authority: ctx.accounts.payer.to_account_info(),
-// 			},
-// 		)
-// 	)?;
-// 	token::revoke(
-// 		CpiContext::new(
-// 			ctx.accounts.token_program.to_account_info(),
-// 			token::Revoke {
-// 				// source
-// 				source: ctx.accounts.user_token_account_1.to_account_info(),
-// 				authority: ctx.accounts.payer.to_account_info(),
-// 			},
-// 		)
-// 	)?;
-// 	msg!("Revoked delegate authority for user_token_accounts");
-// 	token::revoke(
-// 		CpiContext::new(
-// 			ctx.accounts.token_program.to_account_info(),
-// 			token::Revoke {
-// 				// source
-// 				source: ctx.accounts.user_lp_token_account.to_account_info(),
-// 				authority: ctx.accounts.payer.to_account_info(),
-// 			},
-// 		)
-// 	)?;
-// 	msg!("Revoked delegate authority for user_lp_token_account");
-// 	msg!("[revoke]: finished revoke");
-// 	Ok(())
-// }
diff --git a/packages/solana-contracts/programs/two-pool/src/instructions/marginal_prices.rs b/packages/solana-contracts/programs/two-pool/src/instructions/marginal_prices.rs
@@ -1,12 +1,12 @@
 use {
     crate::{
-        array_equalize, common::create_array, error::*, invariant::Invariant, to_equalized,
+        common::{array_equalize, create_array, to_equalized},
+        error::*,
+        invariant::Invariant,
         BorshDecimal, TwoPool, TOKEN_COUNT,
     },
     anchor_lang::prelude::*,
-    anchor_spl::{
-        token::{Mint, TokenAccount},
-    },
+    anchor_spl::token::{Mint, TokenAccount},
 };
 
 #[derive(Accounts)]

diff --git a/packages/solana-contracts/programs/two-pool/src/instructions/remove_exact_burn.rs b/packages/solana-contracts/programs/two-pool/src/instructions/remove_exact_burn.rs
@@ -1,6 +1,9 @@
 use {
     crate::{
-        array_equalize, error::*, gen_pool_signer_seeds, invariant::Invariant, result_from_equalized, to_equalized,
+        common::{array_equalize, result_from_equalized, to_equalized},
+        error::*,
+        gen_pool_signer_seeds,
+        invariant::Invariant,
         TwoPool, TOKEN_COUNT,
     },
     anchor_lang::prelude::*,

diff --git a/packages/solana-contracts/programs/two-pool/src/instructions/remove_exact_output.rs b/packages/solana-contracts/programs/two-pool/src/instructions/remove_exact_output.rs
@@ -1,7 +1,10 @@
 use {
     crate::{
-        array_equalize, error::*, gen_pool_signer_seeds, get_current_ts, invariant::Invariant, result_from_equalized,
-        to_equalized, TwoPool, TOKEN_COUNT,
+        common::{array_equalize, result_from_equalized, to_equalized},
+        error::*,
+        gen_pool_signer_seeds, get_current_ts,
+        invariant::Invariant,
+        TwoPool, TOKEN_COUNT,
     },
     anchor_lang::prelude::*,
     anchor_spl::{

diff --git a/packages/solana-contracts/programs/two-pool/src/instructions/swap_exact_input.rs b/packages/solana-contracts/programs/two-pool/src/instructions/swap_exact_input.rs
@@ -1,6 +1,9 @@
 use {
     crate::{
-        array_equalize, error::*, gen_pool_signer_seeds, invariant::Invariant, result_from_equalized, to_equalized,
+        common::{array_equalize, result_from_equalized, to_equalized},
+        error::*,
+        gen_pool_signer_seeds,
+        invariant::Invariant,
         TwoPool, TOKEN_COUNT,
     },
     anchor_lang::prelude::*,