prettify (#1702)

Author: NourAlharithi

programs/drift/src/instructions/admin.rs

@@ -5100,6 +5100,27 @@ pub fn handle_add_amm_constituent_data<'info>(
     Ok(())
 }
 
+pub fn handle_update_constituent_correlation_data<'info>(
+    ctx: Context<UpdateConstituentCorrelation>,
+    index1: u16,
+    index2: u16,
+    corr: i64,
+) -> Result<()> {
+    let constituent_correlations = &mut ctx.accounts.constituent_correlations;
+    constituent_correlations.set_correlation(index1, index2, corr)?;
+
+    msg!(
+        "Updated correlation between constituent {} and {} to {}",
+        index1,
+        index2,
+        corr
+    );
+
+    constituent_correlations.validate()?;
+
+    Ok(())
+}
+
 pub fn handle_begin_lp_swap<'c: 'info, 'info>(
     ctx: Context<'_, '_, 'c, 'info, LPTakerSwap<'info>>,
     in_market_index: u16,
@@ -6459,6 +6480,24 @@ pub struct RemoveAmmConstituentMappingData<'info> {
     pub state: Box<Account<'info, State>>,
 }
 
+#[derive(Accounts)]
+pub struct UpdateConstituentCorrelation<'info> {
+    #[account(
+        mut,
+        constraint = admin.key() == admin_hot_wallet::id() || admin.key() == state.admin
+    )]
+    pub admin: Signer<'info>,
+    pub lp_pool: AccountLoader<'info, LPPool>,
+
+    #[account(
+        mut,
+        seeds = [CONSTITUENT_CORRELATIONS_PDA_SEED.as_ref(), lp_pool.key().as_ref()],
+        bump = constituent_correlations.bump,
+    )]
+    pub constituent_correlations: Box<Account<'info, ConstituentCorrelations>>,
+    pub state: Box<Account<'info, State>>,
+}
+
 #[derive(Accounts)]
 #[instruction(
     in_market_index: u16,

programs/drift/src/lib.rs

@@ -1802,6 +1802,15 @@ pub mod drift {
         handle_remove_amm_constituent_mapping_data(ctx, perp_market_index, constituent_index)
     }
 
+    pub fn update_constituent_correlation_data(
+        ctx: Context<UpdateConstituentCorrelation>,
+        index1: u16,
+        index2: u16,
+        correlation: i64,
+    ) -> Result<()> {
+        handle_update_constituent_correlation_data(ctx, index1, index2, correlation)
+    }
+
     pub fn update_lp_constituent_target_base<'c: 'info, 'info>(
         ctx: Context<'_, '_, 'c, 'info, UpdateConstituentTargetBase<'info>>,
     ) -> Result<()> {

sdk/src/adminClient.ts

@@ -4769,6 +4769,50 @@ export class AdminClient extends DriftClient {
 		];
 	}
 
+	public async updateConstituentCorrelationData(
+		lpPoolName: number[],
+		index1: number,
+		index2: number,
+		correlation: BN
+	): Promise<TransactionSignature> {
+		const ixs = await this.getUpdateConstituentCorrelationDataIx(
+			lpPoolName,
+			index1,
+			index2,
+			correlation
+		);
+		const tx = await this.buildTransaction(ixs);
+		const { txSig } = await this.sendTransaction(tx, []);
+		return txSig;
+	}
+
+	public async getUpdateConstituentCorrelationDataIx(
+		lpPoolName: number[],
+		index1: number,
+		index2: number,
+		correlation: BN
+	): Promise<TransactionInstruction[]> {
+		const lpPool = getLpPoolPublicKey(this.program.programId, lpPoolName);
+		return [
+			this.program.instruction.updateConstituentCorrelationData(
+				index1,
+				index2,
+				correlation,
+				{
+					accounts: {
+						admin: this.wallet.publicKey,
+						lpPool,
+						constituentCorrelations: getConstituentCorrelationsPublicKey(
+							this.program.programId,
+							lpPool
+						),
+						state: await this.getStatePublicKey(),
+					},
+				}
+			),
+		];
+	}
+
 	/**
 	 * Get the drift begin_swap and end_swap instructions
 	 *

sdk/src/idl/drift.json

@@ -7591,6 +7591,45 @@
         }
       ]
     },
+    {
+      "name": "updateConstituentCorrelationData",
+      "accounts": [
+        {
+          "name": "admin",
+          "isMut": true,
+          "isSigner": true
+        },
+        {
+          "name": "lpPool",
+          "isMut": false,
+          "isSigner": false
+        },
+        {
+          "name": "constituentCorrelations",
+          "isMut": true,
+          "isSigner": false
+        },
+        {
+          "name": "state",
+          "isMut": false,
+          "isSigner": false
+        }
+      ],
+      "args": [
+        {
+          "name": "index1",
+          "type": "u16"
+        },
+        {
+          "name": "index2",
+          "type": "u16"
+        },
+        {
+          "name": "correlation",
+          "type": "i64"
+        }
+      ]
+    },
     {
       "name": "updateLpConstituentTargetBase",
       "accounts": [

tests/lpPool.ts

@@ -477,7 +477,7 @@ describe('LP Pool', () => {
 		assert(ammCache.cache[0].oraclePrice.eq(new BN(200000000)));
 	});
 
-	it('can update constituent properties', async () => {
+	it('can update constituent properties and correlations', async () => {
 		const constituentPublicKey = getConstituentPublicKey(
 			program.programId,
 			lpPoolKey,
@@ -505,6 +505,20 @@ describe('LP Pool', () => {
 			)) as ConstituentTargetBase;
 		expect(targets).to.not.be.null;
 		assert(targets.targets[constituent.constituentIndex].costToTradeBps == 10);
+
+		await adminClient.updateConstituentCorrelationData(
+			encodeName(lpPoolName),
+			0,
+			1,
+			PERCENTAGE_PRECISION.muln(87).divn(100)
+		);
+
+		await adminClient.updateConstituentCorrelationData(
+			encodeName(lpPoolName),
+			0,
+			1,
+			PERCENTAGE_PRECISION
+		);
 	});
 
 	it('fails adding datum with bad params', async () => {