Nour/cu profiling (#1870)

moosecat2 · web-flow · commit a2ff5631a158 · 2025-09-10T15:48:41.000-07:00
* add CU profiling test

* reduce CUs for target base
diff --git a/programs/drift/src/instructions/lp_pool.rs b/programs/drift/src/instructions/lp_pool.rs
@@ -51,7 +51,10 @@ use super::optional_accounts::{get_whitelist_token, load_maps, AccountMaps};
 use crate::controller::spot_balance::update_spot_market_cumulative_interest;
 use crate::controller::token::{receive, send_from_program_vault_with_signature_seeds};
 use crate::instructions::constraints::*;
-use crate::state::lp_pool::{CONSTITUENT_PDA_SEED, LP_POOL_TOKEN_VAULT_PDA_SEED};
+use crate::state::lp_pool::{
+    AmmInventoryAndPrices, ConstituentIndexAndDecimalAndPrice, CONSTITUENT_PDA_SEED,
+    LP_POOL_TOKEN_VAULT_PDA_SEED,
+};
 
 pub fn handle_update_constituent_target_base<'c: 'info, 'info>(
     ctx: Context<'_, '_, 'c, 'info, UpdateConstituentTargetBase<'info>>,
@@ -68,7 +71,6 @@ pub fn handle_update_constituent_target_base<'c: 'info, 'info>(
     amm_cache.check_oracle_staleness(slot, MAX_AMM_CACHE_ORACLE_STALENESS_FOR_TARGET_CALC)?;
     amm_cache.check_perp_market_staleness(slot, MAX_AMM_CACHE_STALENESS_FOR_TARGET_CALC)?;
 
-    let state = &ctx.accounts.state;
     let mut constituent_target_base: AccountZeroCopyMut<
         '_,
         TargetsDatum,
@@ -81,13 +83,6 @@ pub fn handle_update_constituent_target_base<'c: 'info, 'info>(
         "Constituent target base lp pool pubkey does not match lp pool pubkey",
     )?;
 
-    let num_constituents = constituent_target_base.len();
-    for datum in constituent_target_base.iter() {
-        msg!("weight datum: {:?}", datum);
-    }
-
-    let slot = Clock::get()?.slot;
-
     let amm_constituent_mapping: AccountZeroCopy<
         '_,
         AmmConstituentDatum,
@@ -103,7 +98,7 @@ pub fn handle_update_constituent_target_base<'c: 'info, 'info>(
     let constituent_map =
         ConstituentMap::load(&ConstituentSet::new(), &lp_pool_key, remaining_accounts)?;
 
-    let mut amm_inventories: Vec<(u16, i64, i64)> =
+    let mut amm_inventories: Vec<AmmInventoryAndPrices> =
         Vec::with_capacity(amm_constituent_mapping.len() as usize);
     for (_, datum) in amm_constituent_mapping.iter().enumerate() {
         let cache_info = amm_cache.get(datum.perp_market_index as u32);
@@ -117,39 +112,29 @@ pub fn handle_update_constituent_target_base<'c: 'info, 'info>(
             continue;
         }
 
-        amm_inventories.push((
-            datum.perp_market_index,
-            cache_info.position,
-            cache_info.oracle_price,
-        ));
+        amm_inventories.push(AmmInventoryAndPrices {
+            perp_market_index: datum.perp_market_index,
+            inventory: cache_info.position,
+            price: cache_info.oracle_price,
+        });
     }
 
     if amm_inventories.is_empty() {
         msg!("No valid inventories found for constituent target weights update");
         return Ok(());
     }
 
-    let mut constituent_indexes_and_decimals_and_prices: Vec<(u16, u8, i64)> =
+    let mut constituent_indexes_and_decimals_and_prices: Vec<ConstituentIndexAndDecimalAndPrice> =
         Vec::with_capacity(constituent_map.0.len());
     for (index, loader) in &constituent_map.0 {
         let constituent_ref = loader.load()?;
-        constituent_indexes_and_decimals_and_prices.push((
-            *index,
-            constituent_ref.decimals,
-            constituent_ref.last_oracle_price,
-        ));
+        constituent_indexes_and_decimals_and_prices.push(ConstituentIndexAndDecimalAndPrice {
+            constituent_index: *index,
+            decimals: constituent_ref.decimals,
+            price: constituent_ref.last_oracle_price,
+        });
     }
 
-    let exists_invalid_constituent_index = constituent_indexes_and_decimals_and_prices
-        .iter()
-        .any(|(index, _, _)| *index as u32 >= num_constituents);
-
-    validate!(
-        !exists_invalid_constituent_index,
-        ErrorCode::InvalidUpdateConstituentTargetBaseArgument,
-        "Constituent index larger than numr of constituent target weights"
-    )?;
-
     constituent_target_base.update_target_base(
         &amm_constituent_mapping,
         amm_inventories.as_slice(),
diff --git a/programs/drift/src/state/lp_pool.rs b/programs/drift/src/state/lp_pool.rs
@@ -1154,46 +1154,84 @@ pub fn calculate_target_weight(
 }
 
 /// Update target base based on amm_inventory and mapping
+pub struct AmmInventoryAndPrices {
+    pub perp_market_index: u16,
+    pub inventory: i64,
+    pub price: i64,
+}
+
+pub struct ConstituentIndexAndDecimalAndPrice {
+    pub constituent_index: u16,
+    pub decimals: u8,
+    pub price: i64,
+}
+
 impl<'a> AccountZeroCopyMut<'a, TargetsDatum, ConstituentTargetBaseFixed> {
     pub fn update_target_base(
         &mut self,
         mapping: &AccountZeroCopy<'a, AmmConstituentDatum, AmmConstituentMappingFixed>,
         // (perp market index, inventory, price)
-        amm_inventory_and_prices: &[(u16, i64, i64)],
-        constituents_indexes_and_decimals_and_prices: &[(u16, u8, i64)],
+        amm_inventory_and_prices: &[AmmInventoryAndPrices],
+        constituents_indexes_and_decimals_and_prices: &[ConstituentIndexAndDecimalAndPrice],
         slot: u64,
     ) -> DriftResult<Vec<i128>> {
+        // Precompute notional by perp market index
+        let mut notional_by_perp: Vec<(u16, i128)> =
+            Vec::with_capacity(amm_inventory_and_prices.len());
+        for &AmmInventoryAndPrices {
+            perp_market_index,
+            inventory,
+            price,
+        } in amm_inventory_and_prices.iter()
+        {
+            let notional = (inventory as i128)
+                .safe_mul(price as i128)?
+                .safe_div(BASE_PRECISION_I128)?;
+            notional_by_perp.push((perp_market_index, notional));
+        }
+        notional_by_perp.sort_by_key(|&(idx, _)| idx);
+
+        #[inline]
+        fn find_notional(sorted: &[(u16, i128)], idx: u16) -> Option<i128> {
+            match sorted.binary_search_by_key(&idx, |&(p, _)| p) {
+                Ok(i) => Some(sorted[i].1),
+                Err(_) => None,
+            }
+        }
+
         let mut results = Vec::with_capacity(constituents_indexes_and_decimals_and_prices.len());
-        for (i, (constituent_index, decimals, price)) in
-            constituents_indexes_and_decimals_and_prices
-                .iter()
-                .enumerate()
+        for (
+            i,
+            &ConstituentIndexAndDecimalAndPrice {
+                constituent_index,
+                decimals,
+                price,
+            },
+        ) in constituents_indexes_and_decimals_and_prices
+            .iter()
+            .enumerate()
         {
             let mut target_notional = 0i128;
 
-            for (perp_market_index, inventory, price) in amm_inventory_and_prices.iter() {
-                if let Some(idx) = mapping.iter().position(|d| {
-                    &d.perp_market_index == perp_market_index
-                        && d.constituent_index == *constituent_index
-                }) {
-                    let weight = mapping.get(idx as u32).weight; // PERCENTAGE_PRECISION
-
-                    let notional: i128 = (*inventory as i128)
-                        .safe_mul(*price as i128)?
-                        .safe_div(BASE_PRECISION_I128)?;
+            for d in mapping.iter() {
+                if d.constituent_index != constituent_index {
+                    continue;
+                }
 
+                if let Some(perp_notional) = find_notional(&notional_by_perp, d.perp_market_index) {
+                    let w = d.weight as i128;
                     target_notional = target_notional.saturating_add(
-                        notional
-                            .saturating_mul(weight as i128)
+                        perp_notional
+                            .saturating_mul(w)
                             .saturating_div(PERCENTAGE_PRECISION_I128),
                     );
                 }
             }
 
             let cell = self.get_mut(i as u32);
             let target_base = target_notional
-                .safe_mul(10_i128.pow(*decimals as u32))?
-                .safe_div(*price as i128)?
+                .safe_mul(10_i128.pow(decimals as u32))?
+                .safe_div(price as i128)?
                 * -1; // Want to target opposite sign of total scaled notional inventory
 
             msg!(
diff --git a/programs/drift/src/state/lp_pool/tests.rs b/programs/drift/src/state/lp_pool/tests.rs
@@ -68,17 +68,49 @@ mod tests {
             }
         };
 
-        let amm_inventory_and_price: Vec<(u16, i64, i64)> = vec![
-            (0, 4 * BASE_PRECISION_I64, 100_000 * PRICE_PRECISION_I64), // $400k BTC
-            (1, 2000 * BASE_PRECISION_I64, 200 * PRICE_PRECISION_I64),  // $400k SOL
-            (2, 200 * BASE_PRECISION_I64, 1500 * PRICE_PRECISION_I64),  // $300k ETH
-            (3, 16500 * BASE_PRECISION_I64, PRICE_PRECISION_I64),       // $16.5k FARTCOIN
+        let amm_inventory_and_price: Vec<AmmInventoryAndPrices> = vec![
+            AmmInventoryAndPrices {
+                perp_market_index: 0,
+                inventory: 4 * BASE_PRECISION_I64,
+                price: 100_000 * PRICE_PRECISION_I64,
+            }, // $400k BTC
+            AmmInventoryAndPrices {
+                perp_market_index: 1,
+                inventory: 2000 * BASE_PRECISION_I64,
+                price: 200 * PRICE_PRECISION_I64,
+            }, // $400k SOL
+            AmmInventoryAndPrices {
+                perp_market_index: 2,
+                inventory: 200 * BASE_PRECISION_I64,
+                price: 1500 * PRICE_PRECISION_I64,
+            }, // $300k ETH
+            AmmInventoryAndPrices {
+                perp_market_index: 3,
+                inventory: 16500 * BASE_PRECISION_I64,
+                price: PRICE_PRECISION_I64,
+            }, // $16.5k FARTCOIN
         ];
         let constituents_indexes_and_decimals_and_prices = vec![
-            (0, 6, 100_000 * PRICE_PRECISION_I64),
-            (1, 6, 200 * PRICE_PRECISION_I64),
-            (2, 6, 1500 * PRICE_PRECISION_I64),
-            (3, 6, PRICE_PRECISION_I64), // USDC
+            ConstituentIndexAndDecimalAndPrice {
+                constituent_index: 0,
+                decimals: 6,
+                price: 100_000 * PRICE_PRECISION_I64,
+            },
+            ConstituentIndexAndDecimalAndPrice {
+                constituent_index: 1,
+                decimals: 6,
+                price: 200 * PRICE_PRECISION_I64,
+            },
+            ConstituentIndexAndDecimalAndPrice {
+                constituent_index: 2,
+                decimals: 6,
+                price: 1500 * PRICE_PRECISION_I64,
+            },
+            ConstituentIndexAndDecimalAndPrice {
+                constituent_index: 3,
+                decimals: 6,
+                price: PRICE_PRECISION_I64,
+            }, // USDC
         ];
         let aum = 2_000_000 * QUOTE_PRECISION; // $2M AUM
 
@@ -112,7 +144,7 @@ mod tests {
                 calculate_target_weight(
                     base.cast::<i64>().unwrap(),
                     &SpotMarket::default_quote_market(),
-                    amm_inventory_and_price.get(index).unwrap().2,
+                    amm_inventory_and_price.get(index).unwrap().price,
                     aum,
                 )
                 .unwrap()
@@ -155,8 +187,17 @@ mod tests {
             }
         };
 
-        let amm_inventory_and_prices: Vec<(u16, i64, i64)> = vec![(0, 1_000_000, 1_000_000)];
-        let constituents_indexes_and_decimals_and_prices = vec![(1, 6, 1_000_000)];
+        let amm_inventory_and_prices: Vec<AmmInventoryAndPrices> = vec![AmmInventoryAndPrices {
+            perp_market_index: 0,
+            inventory: 1_000_000,
+            price: 1_000_000,
+        }];
+        let constituents_indexes_and_decimals_and_prices =
+            vec![ConstituentIndexAndDecimalAndPrice {
+                constituent_index: 1,
+                decimals: 6,
+                price: 1_000_000,
+            }];
         let aum = 1_000_000;
         let now_ts = 1000;
 
@@ -215,8 +256,17 @@ mod tests {
         };
 
         let price = PRICE_PRECISION_I64;
-        let amm_inventory_and_prices: Vec<(u16, i64, i64)> = vec![(0, BASE_PRECISION_I64, price)];
-        let constituents_indexes_and_decimals_and_prices = vec![(1, 6, price)];
+        let amm_inventory_and_prices: Vec<AmmInventoryAndPrices> = vec![AmmInventoryAndPrices {
+            perp_market_index: 0,
+            inventory: BASE_PRECISION_I64,
+            price,
+        }];
+        let constituents_indexes_and_decimals_and_prices =
+            vec![ConstituentIndexAndDecimalAndPrice {
+                constituent_index: 1,
+                decimals: 6,
+                price,
+            }];
         let aum = 1_000_000;
         let now_ts = 1234;
 
@@ -293,9 +343,23 @@ mod tests {
             }
         };
 
-        let amm_inventory_and_prices: Vec<(u16, i64, i64)> = vec![(0, 1_000_000_000, 1_000_000)];
-        let constituents_indexes_and_decimals_and_prices =
-            vec![(1, 6, 1_000_000), (2, 6, 1_000_000)];
+        let amm_inventory_and_prices: Vec<AmmInventoryAndPrices> = vec![AmmInventoryAndPrices {
+            perp_market_index: 0,
+            inventory: 1_000_000_000,
+            price: 1_000_000,
+        }];
+        let constituents_indexes_and_decimals_and_prices = vec![
+            ConstituentIndexAndDecimalAndPrice {
+                constituent_index: 1,
+                decimals: 6,
+                price: 1_000_000,
+            },
+            ConstituentIndexAndDecimalAndPrice {
+                constituent_index: 2,
+                decimals: 6,
+                price: 1_000_000,
+            },
+        ];
 
         let aum = 1_000_000;
         let now_ts = 999;
@@ -330,7 +394,7 @@ mod tests {
                     constituents_indexes_and_decimals_and_prices
                         .get(i as usize)
                         .unwrap()
-                        .2,
+                        .price,
                     aum,
                 )
                 .unwrap(),
@@ -368,8 +432,17 @@ mod tests {
             }
         };
 
-        let amm_inventory_and_prices: Vec<(u16, i64, i64)> = vec![(0, 1_000_000, 142_000_000)];
-        let constituents_indexes_and_decimals_and_prices = vec![(1, 6, 142_000_000)];
+        let amm_inventory_and_prices: Vec<AmmInventoryAndPrices> = vec![AmmInventoryAndPrices {
+            perp_market_index: 0,
+            inventory: 1_000_000,
+            price: 142_000_000,
+        }];
+        let constituents_indexes_and_decimals_and_prices =
+            vec![ConstituentIndexAndDecimalAndPrice {
+                constituent_index: 1,
+                decimals: 6,
+                price: 142_000_000,
+            }];
 
         let prices = vec![142_000_000];
         let aum = 0;
diff --git a/sdk/src/driftClient.ts b/sdk/src/driftClient.ts
@@ -10352,9 +10352,7 @@ export class DriftClient {
 	public async getUpdateConstituentOracleInfoIx(
 		constituent: ConstituentAccount
 	): Promise<TransactionInstruction> {
-		const spotMarket = await this.getSpotMarketAccount(
-			constituent.spotMarketIndex
-		);
+		const spotMarket = this.getSpotMarketAccount(constituent.spotMarketIndex);
 		return this.program.instruction.updateConstituentOracleInfo({
 			accounts: {
 				keeper: this.wallet.publicKey,
diff --git a/tests/lpPool.ts b/tests/lpPool.ts
@@ -12,7 +12,6 @@ import {
 	Transaction,
 } from '@solana/web3.js';
 import {
-	createAssociatedTokenAccountIdempotentInstruction,
 	createAssociatedTokenAccountInstruction,
 	createInitializeMint2Instruction,
 	createMintToInstruction,
diff --git a/tests/lpPoolCUs.ts b/tests/lpPoolCUs.ts
