From 7958915d60a8753167a2fea9253e550d6f7f41d1 Mon Sep 17 00:00:00 2001 From: jordy25519 Date: Thu, 21 Nov 2024 16:42:40 -0500 Subject: [PATCH] Add forced markets to AccounstListBuilder methods --- ...map_builder.rs => account_list_builder.rs} | 113 ++++++++++++------ crates/src/math/leverage.rs | 8 +- crates/src/math/liquidation.rs | 10 +- crates/src/math/mod.rs | 2 +- 4 files changed, 89 insertions(+), 44 deletions(-) rename crates/src/math/{account_map_builder.rs => account_list_builder.rs} (66%) diff --git a/crates/src/math/account_map_builder.rs b/crates/src/math/account_list_builder.rs similarity index 66% rename from crates/src/math/account_map_builder.rs rename to crates/src/math/account_list_builder.rs index 303b9e9..1e66d4d 100644 --- a/crates/src/math/account_map_builder.rs +++ b/crates/src/math/account_list_builder.rs @@ -29,41 +29,64 @@ impl AccountsListBuilder { /// /// * `client` - drift client instance /// * `user` - the account to build against + /// * `force_markets` - additional market accounts that should be included in the account list /// /// It relies on the `client` being subscribed to all the necessary markets and oracles - pub fn try_build(&mut self, client: &DriftClient, user: &User) -> SdkResult { + pub fn try_build( + &mut self, + client: &DriftClient, + user: &User, + force_markets: &[MarketId], + ) -> SdkResult { let mut oracle_markets = HashMap::::with_capacity_and_hasher(16, Default::default()); - let mut spot_markets = Vec::::with_capacity(user.spot_positions.len()); - let mut perp_markets = Vec::::with_capacity(user.perp_positions.len()); + let mut spot_markets = Vec::::new(); + let mut perp_markets = Vec::::new(); let drift_state_account = client.try_get_account::(state_account())?; - for p in user.spot_positions.iter().filter(|p| !p.is_available()) { - let market = client.try_get_spot_market_account(p.market_index)?; + let force_spot_iter = force_markets + .iter() + .filter(|m| m.is_spot()) + .map(|m| m.index()); + let mut spot_market_idxs = ahash::HashSet::from_iter( + user.spot_positions + .iter() + .filter(|p| !p.is_available()) + .map(|p| p.market_index) + .chain(force_spot_iter), + ); + spot_market_idxs.insert(MarketId::QUOTE_SPOT.index()); + + for idx in spot_market_idxs { + let market = client.try_get_spot_market_account(idx)?; oracle_markets.insert(market.oracle, MarketId::spot(market.market_index)); spot_markets.push(market); } - let quote_market = client.try_get_spot_market_account(MarketId::QUOTE_SPOT.index())?; - if oracle_markets - .insert(quote_market.oracle, MarketId::QUOTE_SPOT) - .is_none() - { - spot_markets.push(quote_market); - } + let force_perp_iter = force_markets + .iter() + .filter(|m| m.is_perp()) + .map(|m| m.index()); + let perp_market_idxs = ahash::HashSet::from_iter( + user.perp_positions + .iter() + .filter(|p| !p.is_available()) + .map(|p| p.market_index) + .chain(force_perp_iter), + ); - for p in user.perp_positions.iter().filter(|p| !p.is_available()) { - let market = client.try_get_perp_market_account(p.market_index)?; + for idx in perp_market_idxs { + let market = client.try_get_perp_market_account(idx)?; oracle_markets.insert(market.amm.oracle, MarketId::perp(market.market_index)); perp_markets.push(market); } - for market in spot_markets.iter() { + for market in spot_markets { self.spot_accounts.push( ( market.pubkey, Account { - data: zero_account_to_bytes(*market), + data: zero_account_to_bytes(market), owner: constants::PROGRAM_ID, ..Default::default() }, @@ -72,12 +95,12 @@ impl AccountsListBuilder { ); } - for market in perp_markets.iter() { + for market in perp_markets { self.perp_accounts.push( ( market.pubkey, Account { - data: zero_account_to_bytes(*market), + data: zero_account_to_bytes(market), owner: constants::PROGRAM_ID, ..Default::default() }, @@ -121,34 +144,56 @@ impl AccountsListBuilder { /// /// * `client` - drift client instance /// * `user` - the account to build against + /// * `force_markets` - additional market accounts that should be included in the account list /// /// like `try_build` but will fall back to network queries to fetch market/oracle accounts as required /// if the client is already subscribed to necessary market/oracles then no network requests are made. - pub async fn build(&mut self, client: &DriftClient, user: &User) -> SdkResult { + pub async fn build( + &mut self, + client: &DriftClient, + user: &User, + force_markets: &[MarketId], + ) -> SdkResult { let mut oracle_markets = HashMap::::with_capacity_and_hasher(16, Default::default()); - let mut spot_markets = Vec::::with_capacity(user.spot_positions.len()); - let mut perp_markets = Vec::::with_capacity(user.perp_positions.len()); + let mut spot_markets = Vec::::new(); + let mut perp_markets = Vec::::new(); let drift_state_account = client.try_get_account::(state_account())?; - for p in user.spot_positions.iter().filter(|p| !p.is_available()) { - let market = client.get_spot_market_account(p.market_index).await?; + // TODO: could batch the requests + let force_spot_iter = force_markets + .iter() + .filter(|m| m.is_spot()) + .map(|m| m.index()); + let mut spot_market_idxs = ahash::HashSet::from_iter( + user.spot_positions + .iter() + .filter(|p| !p.is_available()) + .map(|p| p.market_index) + .chain(force_spot_iter), + ); + spot_market_idxs.insert(MarketId::QUOTE_SPOT.index()); + + for market_idx in spot_market_idxs.iter() { + let market = client.get_spot_market_account(*market_idx).await?; oracle_markets.insert(market.oracle, MarketId::spot(market.market_index)); spot_markets.push(market); } - let quote_market = client - .get_spot_market_account(MarketId::QUOTE_SPOT.index()) - .await?; - if oracle_markets - .insert(quote_market.oracle, MarketId::QUOTE_SPOT) - .is_none() - { - spot_markets.push(quote_market); - } + let force_perp_iter = force_markets + .iter() + .filter(|m| m.is_perp()) + .map(|m| m.index()); + let perp_market_idxs = ahash::HashSet::from_iter( + user.perp_positions + .iter() + .filter(|p| !p.is_available()) + .map(|p| p.market_index) + .chain(force_perp_iter), + ); - for p in user.perp_positions.iter().filter(|p| !p.is_available()) { - let market = client.get_perp_market_account(p.market_index).await?; + for market_idx in perp_market_idxs.iter() { + let market = client.get_perp_market_account(*market_idx).await?; oracle_markets.insert(market.amm.oracle, MarketId::perp(market.market_index)); perp_markets.push(market); } diff --git a/crates/src/math/leverage.rs b/crates/src/math/leverage.rs index ca0cbaa..a3708c9 100644 --- a/crates/src/math/leverage.rs +++ b/crates/src/math/leverage.rs @@ -1,7 +1,7 @@ use solana_sdk::pubkey::Pubkey; use super::{ - account_map_builder::AccountsListBuilder, + account_list_builder::AccountsListBuilder, constants::{AMM_RESERVE_PRECISION, BASE_PRECISION, MARGIN_PRECISION, PRICE_PRECISION}, }; use crate::{ @@ -17,7 +17,7 @@ use crate::{ pub fn get_leverage(client: &DriftClient, user: &User) -> SdkResult { let mut builder = AccountsListBuilder::default(); - let mut accounts = builder.try_build(client, user)?; + let mut accounts = builder.try_build(client, user, &[])?; let margin_calculation = calculate_margin_requirement_and_total_collateral_and_liability_info( user, &mut accounts, @@ -47,7 +47,7 @@ pub fn get_leverage(client: &DriftClient, user: &User) -> SdkResult { pub fn get_spot_asset_value(client: &DriftClient, user: &User) -> SdkResult { let mut builder = AccountsListBuilder::default(); - let mut accounts = builder.try_build(client, user)?; + let mut accounts = builder.try_build(client, user, &[])?; let margin_calculation = calculate_margin_requirement_and_total_collateral_and_liability_info( user, @@ -111,7 +111,7 @@ pub trait UserMargin { impl UserMargin for DriftClient { fn calculate_margin_info(&self, user: &User) -> SdkResult { let mut builder = AccountsListBuilder::default(); - let mut accounts = builder.try_build(self, user)?; + let mut accounts = builder.try_build(self, user, &[])?; calculate_margin_requirement_and_total_collateral_and_liability_info( user, &mut accounts, diff --git a/crates/src/math/liquidation.rs b/crates/src/math/liquidation.rs index 4f0cf60..a4bf3d5 100644 --- a/crates/src/math/liquidation.rs +++ b/crates/src/math/liquidation.rs @@ -9,7 +9,7 @@ use crate::{ MarginContextMode, }, math::{ - account_map_builder::AccountsListBuilder, + account_list_builder::AccountsListBuilder, constants::{ AMM_RESERVE_PRECISION_I128, BASE_PRECISION_I128, MARGIN_PRECISION, QUOTE_PRECISION_I128, QUOTE_PRECISION_I64, SPOT_WEIGHT_PRECISION, @@ -51,7 +51,7 @@ pub async fn calculate_liquidation_price_and_unrealized_pnl( // build a list of all user positions for margin calculations let mut builder = AccountsListBuilder::default(); - let mut accounts_list = builder.build(client, user).await?; + let mut accounts_list = builder.build(client, user, &[]).await?; let oracle = accounts_list .oracles @@ -124,7 +124,7 @@ pub async fn calculate_liquidation_price( market_index: u16, ) -> SdkResult { let mut accounts_builder = AccountsListBuilder::default(); - let mut account_maps = accounts_builder.build(client, user).await?; + let mut account_maps = accounts_builder.build(client, user, &[]).await?; let perp_market = client .program_data() .perp_market_config_by_index(market_index) @@ -294,7 +294,7 @@ pub fn calculate_margin_requirements( ) -> SdkResult { calculate_margin_requirements_inner( user, - &mut AccountsListBuilder::default().try_build(client, user)?, + &mut AccountsListBuilder::default().try_build(client, user, &[])?, ) } @@ -337,7 +337,7 @@ pub fn calculate_collateral( let mut accounts_builder = AccountsListBuilder::default(); calculate_collateral_inner( user, - &mut accounts_builder.try_build(client, user)?, + &mut accounts_builder.try_build(client, user, &[])?, margin_requirement_type, ) } diff --git a/crates/src/math/mod.rs b/crates/src/math/mod.rs index e535049..a7a3b30 100644 --- a/crates/src/math/mod.rs +++ b/crates/src/math/mod.rs @@ -3,7 +3,7 @@ use crate::drift_idl::{ types::{MarginCalculationMode, MarginRequirementType, MarketIdentifier}, }; -pub mod account_map_builder; +pub mod account_list_builder; pub mod auction; pub mod constants; pub mod leverage;