Dlp refactor amm cache

programs/drift/src/error.rs

@@ -692,6 +692,8 @@ pub enum ErrorCode {
     Unauthorized,
     #[msg("Invalid Lp Pool Id for Operation")]
     InvalidLpPoolId,
+    #[msg("MarketIndexNotFoundAmmCache")]
+    MarketIndexNotFoundAmmCache,
 }
 
 #[macro_export]

programs/drift/src/instructions/admin.rs

@@ -1107,17 +1107,6 @@ pub fn handle_initialize_perp_market(
 
     safe_increment!(state.number_of_markets, 1);
 
-    let amm_cache = &mut ctx.accounts.amm_cache;
-    let current_len = amm_cache.cache.len();
-    amm_cache
-        .cache
-        .resize_with(current_len + 1, CacheInfo::default);
-    let current_market_info = amm_cache.cache.get_mut(current_len).unwrap();
-    current_market_info.slot = clock_slot;
-    current_market_info.oracle = perp_market.amm.oracle;
-    current_market_info.oracle_source = u8::from(perp_market.amm.oracle_source);
-    amm_cache.validate(state)?;
-
     controller::amm::update_concentration_coef(perp_market, concentration_coef_scale)?;
     crate::dlog!(oracle_price);
 
@@ -1140,28 +1129,32 @@ pub fn handle_initialize_amm_cache(ctx: Context<InitializeAmmCache>) -> Result<(
     Ok(())
 }
 
-pub fn handle_resize_amm_cache(ctx: Context<ResizeAmmCache>) -> Result<()> {
+pub fn handle_add_market_to_amm_cache(ctx: Context<AddMarketToAmmCache>) -> Result<()> {
     let amm_cache = &mut ctx.accounts.amm_cache;
-    let state = &ctx.accounts.state;
+    let perp_market = ctx.accounts.perp_market.load()?;
+
+    for cache_info in amm_cache.cache.iter() {
+        validate!(
+            cache_info.market_index != perp_market.market_index,
+            ErrorCode::DefaultError,
+            "Market index {} already in amm cache",
+            perp_market.market_index
+        )?;
+    }
+
     let current_size = amm_cache.cache.len();
-    let new_size = (state.number_of_markets as usize).min(current_size + 20_usize);
+    let new_size = current_size.saturating_add(1);
 
     msg!(
         "resizing amm cache from {} entries to {}",
         current_size,
         new_size
     );
 
-    let growth = new_size.saturating_sub(current_size);
-    validate!(
-        growth <= 20,
-        ErrorCode::DefaultError,
-        "cannot grow amm_cache by more than 20 entries in a single resize (requested +{})",
-        growth
-    )?;
-
-    amm_cache.cache.resize_with(new_size, CacheInfo::default);
-    amm_cache.validate(state)?;
+    amm_cache.cache.resize_with(new_size, || CacheInfo {
+        market_index: perp_market.market_index,
+        ..CacheInfo::default()
+    });
 
     Ok(())
 }
@@ -5511,15 +5504,6 @@ pub struct InitializePerpMarket<'info> {
         payer = admin
     )]
     pub perp_market: AccountLoader<'info, PerpMarket>,
-    #[account(
-        mut,
-        seeds = [AMM_POSITIONS_CACHE.as_ref()],
-        bump = amm_cache.bump,
-        realloc = AmmCache::space(amm_cache.cache.len() + 1_usize),
-        realloc::payer = admin,
-        realloc::zero = false,
-    )]
-    pub amm_cache: Box<Account<'info, AmmCache>>,
     /// CHECK: checked in `initialize_perp_market`
     pub oracle: AccountInfo<'info>,
     pub rent: Sysvar<'info, Rent>,
@@ -5547,7 +5531,7 @@ pub struct InitializeAmmCache<'info> {
 }
 
 #[derive(Accounts)]
-pub struct ResizeAmmCache<'info> {
+pub struct AddMarketToAmmCache<'info> {
     #[account(
         mut,
         constraint = admin.key() == admin_hot_wallet::id() || admin.key() == state.admin
@@ -5558,11 +5542,12 @@ pub struct ResizeAmmCache<'info> {
         mut,
         seeds = [AMM_POSITIONS_CACHE.as_ref()],
         bump,
-        realloc = AmmCache::space(amm_cache.cache.len() + (state.number_of_markets as usize - amm_cache.cache.len()).min(20_usize)),
+        realloc = AmmCache::space(amm_cache.cache.len() + 1),
         realloc::payer = admin,
         realloc::zero = false,
     )]
     pub amm_cache: Box<Account<'info, AmmCache>>,
+    pub perp_market: AccountLoader<'info, PerpMarket>,
     pub rent: Sysvar<'info, Rent>,
     pub system_program: Program<'info, System>,
 }

programs/drift/src/instructions/keeper.rs

@@ -3378,7 +3378,7 @@ pub fn handle_settle_perp_to_lp_pool<'c: 'info, 'info>(
             continue;
         }
 
-        let cached_info = amm_cache.get_mut(perp_market.market_index as u32);
+        let cached_info = amm_cache.get_for_market_index_mut(perp_market.market_index)?;
 
         // Early validation checks
         if slot.saturating_sub(cached_info.oracle_slot) > SETTLE_AMM_ORACLE_MAX_DELAY {
@@ -3594,7 +3594,7 @@ pub fn handle_update_amm_cache<'c: 'info, 'info>(
         if perp_market.lp_status == 0 {
             continue;
         }
-        let cached_info = amm_cache.get_mut(perp_market.market_index as u32);
+        let cached_info = amm_cache.get_for_market_index_mut(perp_market.market_index)?;
 
         validate!(
             perp_market.oracle_id() == cached_info.oracle_id()?,

programs/drift/src/instructions/lp_admin.rs

@@ -922,20 +922,6 @@ pub fn handle_update_initial_amm_cache_info<'c: 'info, 'info>(
     Ok(())
 }
 
-pub fn handle_reset_amm_cache(ctx: Context<ResetAmmCache>) -> Result<()> {
-    let state = &ctx.accounts.state;
-    let amm_cache = &mut ctx.accounts.amm_cache;
-
-    amm_cache.cache.clear();
-    amm_cache
-        .cache
-        .resize_with(state.number_of_markets as usize, CacheInfo::default);
-    amm_cache.validate(state)?;
-
-    msg!("AMM cache reset. markets: {}", state.number_of_markets);
-    Ok(())
-}
-
 #[derive(Debug, Clone, Copy, AnchorSerialize, AnchorDeserialize, PartialEq, Eq)]
 pub struct OverrideAmmCacheParams {
     pub quote_owed_from_lp_pool: Option<i64>,

programs/drift/src/instructions/lp_pool.rs

@@ -43,6 +43,7 @@ use crate::{
     },
     validate,
 };
+use std::collections::BTreeMap;
 use std::iter::Peekable;
 use std::slice::Iter;
 
@@ -96,31 +97,33 @@ 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<AmmInventoryAndPricesAndSlots> =
-        Vec::with_capacity(amm_cache.len() as usize);
+    let mut amm_inventories: BTreeMap<u16, AmmInventoryAndPricesAndSlots> = BTreeMap::new();
     for (_, cache_info) in amm_cache.iter().enumerate() {
         if cache_info.lp_status_for_perp_market == 0 {
             continue;
         }
 
-        amm_inventories.push(AmmInventoryAndPricesAndSlots {
-            inventory: {
-                let scaled_position = cache_info
-                    .position
-                    .safe_mul(cache_info.amm_position_scalar as i64)?
-                    .safe_div(100)?;
-
-                scaled_position.clamp(
-                    -cache_info.amm_inventory_limit,
-                    cache_info.amm_inventory_limit,
-                )
+        amm_inventories.insert(
+            cache_info.market_index,
+            AmmInventoryAndPricesAndSlots {
+                inventory: {
+                    let scaled_position = cache_info
+                        .position
+                        .safe_mul(cache_info.amm_position_scalar as i64)?
+                        .safe_div(100)?;
+
+                    scaled_position.clamp(
+                        -cache_info.amm_inventory_limit,
+                        cache_info.amm_inventory_limit,
+                    )
+                },
+                price: cache_info.oracle_price,
+                last_oracle_slot: cache_info.oracle_slot,
+                last_position_slot: cache_info.slot,
             },
-            price: cache_info.oracle_price,
-            last_oracle_slot: cache_info.oracle_slot,
-            last_position_slot: cache_info.slot,
-        });
+        );
     }
-    msg!("amm inventories: {:?}", amm_inventories);
+    msg!("amm inventories:{:?}", amm_inventories);
 
     if amm_inventories.is_empty() {
         msg!("No valid inventories found for constituent target weights update");
@@ -140,7 +143,7 @@ pub fn handle_update_constituent_target_base<'c: 'info, 'info>(
 
     constituent_target_base.update_target_base(
         &amm_constituent_mapping,
-        amm_inventories.as_slice(),
+        &amm_inventories,
         constituent_indexes_and_decimals_and_prices.as_mut_slice(),
         slot,
     )?;
@@ -1010,7 +1013,7 @@ pub fn handle_lp_pool_remove_liquidity<'c: 'info, 'info>(
         if cache_info.last_fee_pool_token_amount != 0 && cache_info.last_settle_slot != slot {
             msg!(
                 "Market {} has not been settled in current slot. Last slot: {}",
-                i,
+                cache_info.market_index,
                 cache_info.last_settle_slot
             );
             return Err(ErrorCode::AMMCacheStale.into());

programs/drift/src/lib.rs

@@ -1020,10 +1020,10 @@ pub mod drift {
         handle_initialize_amm_cache(ctx)
     }
 
-    pub fn resize_amm_cache<'c: 'info, 'info>(
-        ctx: Context<'_, '_, 'c, 'info, ResizeAmmCache<'info>>,
+    pub fn add_market_to_amm_cache<'c: 'info, 'info>(
+        ctx: Context<'_, '_, 'c, 'info, AddMarketToAmmCache<'info>>,
     ) -> Result<()> {
-        handle_resize_amm_cache(ctx)
+        handle_add_market_to_amm_cache(ctx)
     }
 
     pub fn update_initial_amm_cache_info<'c: 'info, 'info>(
@@ -2086,12 +2086,6 @@ pub mod drift {
         handle_override_amm_cache_info(ctx, market_index, override_params)
     }
 
-    pub fn reset_amm_cache<'c: 'info, 'info>(
-        ctx: Context<'_, '_, 'c, 'info, ResetAmmCache<'info>>,
-    ) -> Result<()> {
-        handle_reset_amm_cache(ctx)
-    }
-
     pub fn lp_pool_swap<'c: 'info, 'info>(
         ctx: Context<'_, '_, 'c, 'info, LPPoolSwap<'info>>,
         in_market_index: u16,

programs/drift/src/state/amm_cache.rs

@@ -10,11 +10,9 @@ use crate::state::oracle::MMOraclePriceData;
 use crate::state::oracle_map::OracleIdentifier;
 use crate::state::perp_market::PerpMarket;
 use crate::state::spot_market::{SpotBalance, SpotMarket};
-use crate::state::state::State;
 use crate::state::traits::Size;
 use crate::state::zero_copy::HasLen;
 use crate::state::zero_copy::{AccountZeroCopy, AccountZeroCopyMut};
-use crate::validate;
 use crate::OracleSource;
 use crate::{impl_zero_copy_loader, OracleGuardRails};
 
@@ -53,11 +51,12 @@ pub struct CacheInfo {
     pub amm_inventory_limit: i64,
     pub oracle_price: i64,
     pub oracle_slot: u64,
+    pub market_index: u16,
     pub oracle_source: u8,
     pub oracle_validity: u8,
     pub lp_status_for_perp_market: u8,
     pub amm_position_scalar: u8,
-    pub _padding: [u8; 36],
+    pub _padding: [u8; 34],
 }
 
 impl Size for CacheInfo {
@@ -86,7 +85,8 @@ impl Default for CacheInfo {
             quote_owed_from_lp_pool: 0i64,
             lp_status_for_perp_market: 0u8,
             amm_position_scalar: 0u8,
-            _padding: [0u8; 36],
+            market_index: 0u16,
+            _padding: [0u8; 34],
         }
     }
 }
@@ -183,15 +183,6 @@ impl AmmCache {
         8 + 8 + 4 + num_markets * CacheInfo::SIZE
     }
 
-    pub fn validate(&self, state: &State) -> DriftResult<()> {
-        validate!(
-            self.cache.len() <= state.number_of_markets as usize,
-            ErrorCode::DefaultError,
-            "Number of amm positions is no larger than number of markets"
-        )?;
-        Ok(())
-    }
-
     pub fn update_perp_market_fields(&mut self, perp_market: &PerpMarket) -> DriftResult<()> {
         let cache_info = self.cache.get_mut(perp_market.market_index as usize);
         if let Some(cache_info) = cache_info {
@@ -238,6 +229,15 @@ impl AmmCache {
 impl_zero_copy_loader!(AmmCache, crate::id, AmmCacheFixed, CacheInfo);
 
 impl<'a> AccountZeroCopy<'a, CacheInfo, AmmCacheFixed> {
+    pub fn get_for_market_index(&self, market_index: u16) -> DriftResult<&CacheInfo> {
+        for cache_info in self.iter() {
+            if cache_info.market_index == market_index {
+                return Ok(cache_info);
+            }
+        }
+        Err(ErrorCode::MarketIndexNotFoundAmmCache.into())
+    }
+
     pub fn check_settle_staleness(&self, slot: u64, threshold_slot_diff: u64) -> DriftResult<()> {
         for (i, cache_info) in self.iter().enumerate() {
             if cache_info.slot == 0 {
@@ -284,6 +284,27 @@ impl<'a> AccountZeroCopy<'a, CacheInfo, AmmCacheFixed> {
 }
 
 impl<'a> AccountZeroCopyMut<'a, CacheInfo, AmmCacheFixed> {
+    pub fn get_for_market_index_mut(&mut self, market_index: u16) -> DriftResult<&mut CacheInfo> {
+        let pos = {
+            let mut found: Option<u32> = None;
+            for i in 0..self.len() {
+                let cache_info = self.get(i);
+                if cache_info.market_index == market_index {
+                    found = Some(i);
+                    break;
+                }
+            }
+            found
+        };
+
+        if let Some(i) = pos {
+            Ok(self.get_mut(i))
+        } else {
+            msg!("Market index not found in amm cache: {}", market_index);
+            Err(ErrorCode::MarketIndexNotFoundAmmCache.into())
+        }
+    }
+
     pub fn update_amount_owed_from_lp_pool(
         &mut self,
         perp_market: &PerpMarket,