refactor to remove need for dropping in wrapper (#170)

CKS-Systems · Oct 10, 2024 · 9b236a5 · 9b236a5
1 parent 181e63a
commit 9b236a5
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 40 deletions.
diff --git a/programs/wrapper/src/processors/deposit.rs b/programs/wrapper/src/processors/deposit.rs
@@ -1,7 +1,7 @@
-use std::{cell::Ref, mem::size_of};
+use std::cell::Ref;
 
 use borsh::{BorshDeserialize, BorshSerialize};
-use hypertree::{get_helper, DataIndex, RBNode};
+use hypertree::DataIndex;
 use manifest::{
     program::{deposit_instruction, invoke},
     state::MarketFixed,
@@ -15,10 +15,8 @@ use solana_program::{
     pubkey::Pubkey,
 };
 
-use crate::{market_info::MarketInfo, wrapper_state::ManifestWrapperStateFixed};
-
 use super::shared::{
-    check_signer, get_market_info_index_for_market, sync, WrapperStateAccountInfo,
+    check_signer, get_trader_index_hint_for_market, sync, WrapperStateAccountInfo,
 };
 
 #[derive(BorshDeserialize, BorshSerialize)]
@@ -53,34 +51,28 @@ pub(crate) fn process_deposit(
     let mint_account_info: MintAccountInfo =
         MintAccountInfo::new(next_account_info(account_iter)?)?;
 
-    let market_fixed: Ref<MarketFixed> = market.get_fixed()?;
-    let base_mint: Pubkey = *market_fixed.get_base_mint();
-    let quote_mint: Pubkey = *market_fixed.get_quote_mint();
-    let mint: &Pubkey = if &trader_token_account.try_borrow_data()?[0..32] == base_mint.as_ref() {
-        &base_mint
-    } else {
-        &quote_mint
+    let mint: Pubkey = {
+        let market_fixed: Ref<MarketFixed> = market.get_fixed()?;
+        let base_mint: &Pubkey = market_fixed.get_base_mint();
+        let quote_mint: &Pubkey = market_fixed.get_quote_mint();
+        if &trader_token_account.try_borrow_data()?[0..32] == base_mint.as_ref() {
+            *base_mint
+        } else {
+            *quote_mint
+        }
     };
-    drop(market_fixed);
 
     let WrapperDepositParams { amount_atoms } = WrapperDepositParams::try_from_slice(data)?;
 
-    let market_info_index: DataIndex =
-        get_market_info_index_for_market(&wrapper_state, market.info.key);
-    let wrapper_data: Ref<&mut [u8]> = wrapper_state.info.try_borrow_data()?;
-    let (_fixed_data, wrapper_dynamic_data) =
-        wrapper_data.split_at(size_of::<ManifestWrapperStateFixed>());
-    let market_info: MarketInfo =
-        *get_helper::<RBNode<MarketInfo>>(wrapper_dynamic_data, market_info_index).get_value();
-    let trader_index_hint: Option<DataIndex> = Some(market_info.trader_index);
-    drop(wrapper_data);
+    let trader_index_hint: Option<DataIndex> =
+        get_trader_index_hint_for_market(&wrapper_state, &market.info.key)?;
 
     // Call the deposit CPI
     invoke(
         &deposit_instruction(
             market.key,
             owner.key,
-            mint,
+            &mint,
             amount_atoms,
             trader_token_account.key,
             *token_program.key,

diff --git a/programs/wrapper/src/processors/shared.rs b/programs/wrapper/src/processors/shared.rs
@@ -266,6 +266,21 @@ pub(crate) fn get_market_info_index_for_market(
     market_info_index
 }
 
+pub(crate) fn get_trader_index_hint_for_market(
+    wrapper_state: &WrapperStateAccountInfo,
+    market_key: &Pubkey,
+) -> Result<Option<DataIndex>, ProgramError> {
+    let market_info_index: DataIndex = get_market_info_index_for_market(wrapper_state, market_key);
+
+    let wrapper_data: Ref<&mut [u8]> = wrapper_state.info.try_borrow_data()?;
+    let (_fixed_data, wrapper_dynamic_data) =
+        wrapper_data.split_at(size_of::<ManifestWrapperStateFixed>());
+    let market_info: MarketInfo =
+        *get_helper::<RBNode<MarketInfo>>(wrapper_dynamic_data, market_info_index).get_value();
+    let trader_index_hint: Option<DataIndex> = Some(market_info.trader_index);
+    Ok(trader_index_hint)
+}
+
 /// Validation for wrapper account
 #[derive(Clone)]
 pub struct WrapperStateAccountInfo<'a, 'info> {

diff --git a/programs/wrapper/src/processors/withdraw.rs b/programs/wrapper/src/processors/withdraw.rs
@@ -1,7 +1,7 @@
-use std::{cell::Ref, mem::size_of};
+use std::cell::Ref;
 
 use borsh::{BorshDeserialize, BorshSerialize};
-use hypertree::{get_helper, DataIndex, RBNode};
+use hypertree::DataIndex;
 use manifest::{
     program::{invoke, withdraw_instruction},
     state::MarketFixed,
@@ -15,10 +15,8 @@ use solana_program::{
     pubkey::Pubkey,
 };
 
-use crate::{market_info::MarketInfo, wrapper_state::ManifestWrapperStateFixed};
-
 use super::shared::{
-    check_signer, get_market_info_index_for_market, sync, WrapperStateAccountInfo,
+    check_signer, get_trader_index_hint_for_market, sync, WrapperStateAccountInfo,
 };
 
 #[derive(BorshDeserialize, BorshSerialize)]
@@ -52,8 +50,8 @@ pub(crate) fn process_withdraw(
     let mint_account_info: MintAccountInfo =
         MintAccountInfo::new(next_account_info(account_iter)?)?;
 
-    let market_fixed: Ref<MarketFixed> = market.get_fixed()?;
     let mint: Pubkey = {
+        let market_fixed: Ref<MarketFixed> = market.get_fixed()?;
         let base_mint: &Pubkey = market_fixed.get_base_mint();
         let quote_mint: &Pubkey = market_fixed.get_quote_mint();
         if &trader_token_account.try_borrow_data()?[0..32] == base_mint.as_ref() {
@@ -62,21 +60,12 @@ pub(crate) fn process_withdraw(
             *quote_mint
         }
     };
-    drop(market_fixed);
 
     // Params are a direct pass through.
     let WrapperWithdrawParams { amount_atoms } = WrapperWithdrawParams::try_from_slice(data)?;
 
-    // TODO: Make a helper for get_trader_index_hint_for_market
-    let market_info_index: DataIndex =
-        get_market_info_index_for_market(&wrapper_state, market.info.key);
-    let wrapper_data: Ref<&mut [u8]> = wrapper_state.info.try_borrow_data()?;
-    let (_fixed_data, wrapper_dynamic_data) =
-        wrapper_data.split_at(size_of::<ManifestWrapperStateFixed>());
-    let market_info: MarketInfo =
-        *get_helper::<RBNode<MarketInfo>>(wrapper_dynamic_data, market_info_index).get_value();
-    let trader_index_hint: Option<DataIndex> = Some(market_info.trader_index);
-    drop(wrapper_data);
+    let trader_index_hint: Option<DataIndex> =
+        get_trader_index_hint_for_market(&wrapper_state, &market.info.key)?;
 
     // Call the withdraw CPI
     invoke(