update aum considers amm cache (#1701)

Author: NourAlharithi

programs/drift/src/instructions/lp_pool.rs

@@ -28,12 +28,12 @@ use crate::{
         events::{LPMintRedeemRecord, LPSwapRecord},
         lp_pool::{
             calculate_target_weight, AmmConstituentDatum, AmmConstituentMappingFixed, Constituent,
-            ConstituentCorrelations, ConstituentCorrelationsFixed, ConstituentTargetBaseFixed,
-            LPPool, TargetsDatum, WeightValidationFlags, CONSTITUENT_CORRELATIONS_PDA_SEED,
+            ConstituentCorrelationsFixed, ConstituentTargetBaseFixed, LPPool, TargetsDatum,
+            WeightValidationFlags, CONSTITUENT_CORRELATIONS_PDA_SEED,
         },
         oracle::OraclePriceData,
         oracle_map::OracleMap,
-        perp_market::{AmmCacheFixed, CacheInfo, MarketStatus, AMM_POSITIONS_CACHE},
+        perp_market::{AmmCacheFixed, CacheInfo, AMM_POSITIONS_CACHE},
         perp_market_map::MarketSet,
         spot_market::{SpotBalanceType, SpotMarket},
         spot_market_map::get_writable_spot_market_set_from_many,
@@ -261,6 +261,23 @@ pub fn handle_update_lp_pool_aum<'c: 'info, 'info>(
         "Constituent target weights PDA does not match expected PDA"
     )?;
 
+    let amm_cache_key: &Pubkey = &ctx.accounts.amm_cache.key();
+    let amm_cache: AccountZeroCopy<'_, CacheInfo, AmmCacheFixed> =
+        ctx.accounts.amm_cache.load_zc()?;
+    let expected_amm_pda = &Pubkey::create_program_address(
+        &[
+            AMM_POSITIONS_CACHE.as_ref(),
+            amm_cache.fixed.bump.to_le_bytes().as_ref(),
+        ],
+        &crate::ID,
+    )
+    .map_err(|_| ErrorCode::InvalidPDA)?;
+    validate!(
+        amm_cache_key.eq(expected_amm_pda),
+        ErrorCode::InvalidPDA,
+        "Amm cache PDA does not match expected PDA"
+    )?;
+
     let mut aum: u128 = 0;
     let mut crypto_delta = 0_i128;
     let mut oldest_slot = u64::MAX;
@@ -353,6 +370,14 @@ pub fn handle_update_lp_pool_aum<'c: 'info, 'info>(
         aum = aum.safe_add(constituent_aum.cast()?)?;
     }
 
+    for cache_datum in amm_cache.iter() {
+        if cache_datum.quote_owed_from_lp > 0 {
+            aum = aum.safe_sub(cache_datum.quote_owed_from_lp.abs().cast::<u128>()?)?;
+        } else {
+            aum = aum.safe_add(cache_datum.quote_owed_from_lp.abs().cast::<u128>()?)?;
+        }
+    }
+
     lp_pool.oldest_oracle_slot = oldest_slot;
     lp_pool.last_aum = aum;
     lp_pool.last_aum_slot = slot;
@@ -1293,6 +1318,8 @@ pub struct UpdateLPPoolAum<'info> {
     /// CHECK: checked in ConstituentTargetBaseZeroCopy checks
     #[account(mut)]
     pub constituent_target_base: AccountInfo<'info>,
+    /// CHECK: checked in AmmCacheZeroCopy checks
+    pub amm_cache: AccountInfo<'info>,
 }
 
 /// `in`/`out` is in the program's POV for this swap. So `user_in_token_account` is the user owned token account

programs/drift/src/state/lp_pool.rs

@@ -256,6 +256,13 @@ impl LPPool {
         let lp_amount = if self.last_aum == 0 {
             token_amount_usd.safe_div(token_precision_denominator)?
         } else {
+            msg!("token_amount_usd: {}", token_amount_usd);
+            msg!("dlp_total_supply: {}", dlp_total_supply);
+            msg!("last_aum: {}", self.last_aum);
+            msg!(
+                "token_precision_denominator: {}",
+                token_precision_denominator
+            );
             token_amount_usd
                 .safe_mul(dlp_total_supply as u128)?
                 .safe_div(self.last_aum.safe_mul(token_precision_denominator)?)?

programs/drift/src/state/zero_copy.rs

@@ -89,7 +89,9 @@ where
     validate!(
         acct.owner == &program_id,
         ErrorCode::DefaultError,
-        "invalid owner",
+        "invalid owner {}, program_id: {}",
+        acct.owner,
+        program_id,
     )?;
 
     let data = acct.try_borrow_data().safe_unwrap()?;

sdk/src/driftClient.ts

@@ -9815,6 +9815,7 @@ export class DriftClient {
 					this.program.programId,
 					lpPool.pubkey
 				),
+				ammCache: getAmmCachePublicKey(this.program.programId),
 			},
 			remainingAccounts,
 		});

sdk/src/idl/drift.json

@@ -7649,6 +7649,11 @@
           "name": "constituentTargetBase",
           "isMut": true,
           "isSigner": false
+        },
+        {
+          "name": "ammCache",
+          "isMut": false,
+          "isSigner": false
         }
       ],
       "args": []

sdk/src/index.ts

@@ -12,7 +12,6 @@ export * from './accounts/fetch';
 export * from './accounts/webSocketDriftClientAccountSubscriber';
 export * from './accounts/webSocketInsuranceFundStakeAccountSubscriber';
 export * from './accounts/webSocketHighLeverageModeConfigAccountSubscriber';
-export * from './accounts/webSocketConstituentAccountSubscriber';
 export * from './accounts/bulkAccountLoader';
 export * from './accounts/bulkUserSubscription';
 export * from './accounts/bulkUserStatsSubscription';
@@ -23,9 +22,11 @@ export * from './accounts/pollingUserAccountSubscriber';
 export * from './accounts/pollingUserStatsAccountSubscriber';
 export * from './accounts/pollingInsuranceFundStakeAccountSubscriber';
 export * from './accounts/pollingHighLeverageModeConfigAccountSubscriber';
-export * from './accounts/pollingConstituentAccountSubscriber';
 export * from './accounts/basicUserAccountSubscriber';
 export * from './accounts/oneShotUserAccountSubscriber';
+export * from './constituentMap/constituentMap';
+export * from './constituentMap/pollingConstituentAccountSubscriber';
+export * from './constituentMap/webSocketConstituentAccountSubscriber';
 export * from './accounts/types';
 export * from './addresses/pda';
 export * from './adminClient';

tests/lpPool.ts

@@ -52,6 +52,7 @@ import {
 	mockOracleNoProgram,
 	mockUSDCMint,
 	mockUserUSDCAccountWithAuthority,
+	overWriteMintAccount,
 	overWritePerpMarket,
 	overWriteSpotMarket,
 } from './testHelpers';
@@ -924,14 +925,69 @@ describe('LP Pool', () => {
 			getConstituentPublicKey(program.programId, lpPoolKey, 0)
 		)) as ConstituentAccount;
 
+		let lpPool = (await adminClient.program.account.lpPool.fetch(
+			lpPoolKey
+		)) as LPPoolAccount;
+
 		assert(ammCache.cache[0].quoteOwedFromLp.eq(owedAmount.divn(2)));
 		assert(constituent.tokenBalance.eq(ZERO));
+		assert(lpPool.lastAum.eq(ZERO));
+
+		// Deposit here to DLP to make sure aum calc work with perp market debt
+		await overWriteMintAccount(
+			bankrunContextWrapper,
+			lpPool.mint,
+			BigInt(lpPool.lastAum.toNumber())
+		);
+		await adminClient.lpPoolAddLiquidity({
+			lpPool,
+			inAmount: new BN(1000).mul(QUOTE_PRECISION),
+			minMintAmount: new BN(1),
+			inMarketIndex: 0,
+		});
+
+		await adminClient.updateLpPoolAum(lpPool, [0, 1, 2]);
+
+		lpPool = (await adminClient.program.account.lpPool.fetch(
+			lpPoolKey
+		)) as LPPoolAccount;
+
+		let aum = new BN(0);
+		for (let i = 0; i <= 2; i++) {
+			const constituent = (await adminClient.program.account.constituent.fetch(
+				getConstituentPublicKey(program.programId, lpPoolKey, i)
+			)) as ConstituentAccount;
+			aum = aum.add(
+				constituent.tokenBalance
+					.mul(constituent.lastOraclePrice)
+					.div(QUOTE_PRECISION)
+			);
+		}
+
+		// Overwrite the amm cache with amount owed
+		ammCache = (await adminClient.program.account.ammCache.fetch(
+			getAmmCachePublicKey(program.programId)
+		)) as AmmCache;
+		for (let i = 0; i <= ammCache.cache.length - 1; i++) {
+			aum = aum.sub(ammCache.cache[i].quoteOwedFromLp);
+		}
+		assert(lpPool.lastAum.eq(aum));
 	});
 
 	it('perp market will transfer with the constituent vault if it should send more than its owed', async () => {
+		let lpPool = (await adminClient.program.account.lpPool.fetch(
+			lpPoolKey
+		)) as LPPoolAccount;
+		const aumBefore = lpPool.lastAum;
+		let constituent = (await adminClient.program.account.constituent.fetch(
+			getConstituentPublicKey(program.programId, lpPoolKey, 0)
+		)) as ConstituentAccount;
+
 		let ammCache = (await adminClient.program.account.ammCache.fetch(
 			getAmmCachePublicKey(program.programId)
 		)) as AmmCache;
+
+		const balanceBefore = constituent.tokenBalance;
 		const owedAmount = ammCache.cache[0].quoteOwedFromLp;
 
 		// Give the perp market half of its owed amount
@@ -954,11 +1010,16 @@ describe('LP Pool', () => {
 		ammCache = (await adminClient.program.account.ammCache.fetch(
 			getAmmCachePublicKey(program.programId)
 		)) as AmmCache;
-		const constituent = (await adminClient.program.account.constituent.fetch(
+		constituent = (await adminClient.program.account.constituent.fetch(
 			getConstituentPublicKey(program.programId, lpPoolKey, 0)
 		)) as ConstituentAccount;
 
+		lpPool = (await adminClient.program.account.lpPool.fetch(
+			lpPoolKey
+		)) as LPPoolAccount;
+
 		assert(ammCache.cache[0].quoteOwedFromLp.eq(ZERO));
-		assert(constituent.tokenBalance.eq(owedAmount));
+		assert(constituent.tokenBalance.eq(balanceBefore.add(owedAmount)));
+		assert(lpPool.lastAum.eq(aumBefore.add(owedAmount)));
 	});
 });