Merge pull request #1587 from drift-labs/nour/add-mapping-datum

Nour/add mapping datum

Author: moosecat2

programs/drift/src/error.rs

@@ -635,8 +635,8 @@ pub enum ErrorCode {
     InvalidSignedMsgUserOrdersResize,
     #[msg("Invalid Constituent")]
     InvalidConstituent,
-    #[msg("Misatch amm mapping and constituent target weights")]
-    MismatchAmmConstituentMappingAndConstituentTargetWeights,
+    #[msg("Invalid Amm Constituent Mapping argument")]
+    InvalidAmmConstituentMappingArgument,
 }
 
 #[macro_export]

programs/drift/src/instructions/admin.rs

@@ -4600,6 +4600,62 @@ pub fn handle_initialize_constituent<'info>(
     Ok(())
 }
 
+#[derive(AnchorSerialize, AnchorDeserialize, Clone, Default)]
+pub struct InitializeAmmConstituentMappingDatum {
+    pub constituent_index: u16,
+    pub perp_market_index: u16,
+}
+
+pub fn handle_add_amm_constituent_data<'info>(
+    ctx: Context<AddAmmConstituentMappingData>,
+    init_amm_constituent_mapping_data: Vec<InitializeAmmConstituentMappingDatum>,
+) -> Result<()> {
+    let amm_mapping = &mut ctx.accounts.amm_constituent_mapping;
+    let constituent_target_weights = &ctx.accounts.constituent_target_weights;
+    let state = &ctx.accounts.state;
+    let mut current_len = amm_mapping.data.len();
+
+    for init_datum in init_amm_constituent_mapping_data {
+        let perp_market_index = init_datum.perp_market_index;
+
+        validate!(
+            perp_market_index < state.number_of_markets,
+            ErrorCode::InvalidAmmConstituentMappingArgument,
+            "perp_market_index too large compared to number of markets"
+        )?;
+
+        validate!(
+            (init_datum.constituent_index as usize) < constituent_target_weights.data.len(),
+            ErrorCode::InvalidAmmConstituentMappingArgument,
+            "constituent_index too large compared to number of constituents in target weights"
+        )?;
+
+        let constituent_index = init_datum.constituent_index;
+        let mut datum = AmmConstituentDatum::default();
+        datum.perp_market_index = perp_market_index;
+        datum.constituent_index = constituent_index;
+
+        // Check if the datum already exists
+        let exists = amm_mapping.data.iter().any(|d| {
+            d.perp_market_index == perp_market_index && d.constituent_index == constituent_index
+        });
+
+        validate!(
+            !exists,
+            ErrorCode::InvalidAmmConstituentMappingArgument,
+            "AmmConstituentDatum already exists for perp_market_index {} and constituent_index {}",
+            perp_market_index,
+            constituent_index
+        )?;
+
+        // Add the new datum to the mapping
+        current_len += 1;
+        amm_mapping.data.resize(current_len, datum);
+    }
+
+    Ok(())
+}
+
 #[derive(Accounts)]
 pub struct Initialize<'info> {
     #[account(mut)]
@@ -5409,8 +5465,43 @@ pub struct InitializeConstituent<'info> {
         payer = admin,
     )]
     pub constituent: AccountLoader<'info, Constituent>,
+    pub rent: Sysvar<'info, Rent>,
+    pub system_program: Program<'info, System>,
+}
 
+#[derive(Accounts)]
+#[instruction(
+    lp_pool_name: [u8; 32],
+    market_index_constituent_index_pairs:  Vec<(u16, u16)>,
+)]
+pub struct AddAmmConstituentMappingData<'info> {
     #[account(mut)]
-    pub payer: Signer<'info>,
+    pub admin: Signer<'info>,
+
+    #[account(
+        seeds = [b"lp_pool", lp_pool_name.as_ref()],
+        bump,
+    )]
+    pub lp_pool: AccountLoader<'info, LPPool>,
+
+    #[account(
+        mut,
+        seeds = [AMM_MAP_PDA_SEED.as_ref(), lp_pool.key().as_ref()],
+        bump,
+        realloc = AmmConstituentMapping::space(amm_constituent_mapping.data.len() + market_index_constituent_index_pairs.len()),
+        realloc::payer = admin,
+        realloc::zero = false,
+    )]
+    pub amm_constituent_mapping: Box<Account<'info, AmmConstituentMapping>>,
+    #[account(
+        mut,
+        seeds = [CONSTITUENT_TARGET_WEIGHT_PDA_SEED.as_ref(), lp_pool.key().as_ref()],
+        bump,
+        realloc = ConstituentTargetWeights::space(constituent_target_weights.data.len() + 1 as usize),
+        realloc::payer = admin,
+        realloc::zero = false,
+    )]
+    pub constituent_target_weights: Box<Account<'info, ConstituentTargetWeights>>,
+    pub state: Box<Account<'info, State>>,
     pub system_program: Program<'info, System>,
 }

programs/drift/src/lib.rs

@@ -1706,6 +1706,33 @@ pub mod drift {
     ) -> Result<()> {
         handle_update_protected_maker_mode_config(ctx, max_users, reduce_only, current_users)
     }
+
+    pub fn initialize_constituent(
+        ctx: Context<InitializeConstituent>,
+        lp_pool_name: [u8; 32],
+        spot_market_index: u16,
+        decimals: u8,
+        max_weight_deviation: i64,
+        swap_fee_min: i64,
+        swap_fee_max: i64,
+    ) -> Result<()> {
+        handle_initialize_constituent(
+            ctx,
+            spot_market_index,
+            decimals,
+            max_weight_deviation,
+            swap_fee_min,
+            swap_fee_max,
+        )
+    }
+
+    pub fn add_amm_constituent_mapping_data(
+        ctx: Context<AddAmmConstituentMappingData>,
+        lp_pool_name: [u8; 32],
+        init_amm_constituent_mapping_data: Vec<InitializeAmmConstituentMappingDatum>,
+    ) -> Result<()> {
+        handle_add_amm_constituent_data(ctx, init_amm_constituent_mapping_data)
+    }
 }
 
 #[cfg(not(feature = "no-entrypoint"))]

programs/drift/src/state/lp_pool.rs

@@ -208,7 +208,7 @@ pub struct BLPosition {
     pub market_index: u16,
     /// Whether the position is deposit or borrow
     pub balance_type: SpotBalanceType,
-    pub padding: [u8; 4],
+    pub padding: [u8; 5],
 }
 
 impl SpotBalance for BLPosition {
@@ -269,7 +269,7 @@ pub struct Constituent {
 }
 
 impl Size for Constituent {
-    const SIZE: usize = 108;
+    const SIZE: usize = 112;
 }
 
 impl Constituent {
@@ -433,7 +433,7 @@ impl ConstituentTargetWeights {
 
     pub fn validate(&self) -> DriftResult<()> {
         validate!(
-            self.data.len() >= 0 && self.data.len() <= 128,
+            self.data.len() <= 128,
             ErrorCode::DefaultError,
             "Number of constituents len must be between 1 and 128"
         )?;

sdk/src/adminClient.ts

@@ -15,6 +15,7 @@ import {
 	ContractTier,
 	AssetTier,
 	SpotFulfillmentConfigStatus,
+	InitAmmConstituentMappingDatum,
 } from './types';
 import { DEFAULT_MARKET_NAME, encodeName } from './userName';
 import { BN } from '@coral-xyz/anchor';
@@ -40,6 +41,7 @@ import {
 	getLpPoolPublicKey,
 	getAmmConstituentMappingPublicKey,
 	getConstituentTargetWeightsPublicKey,
+	getConstituentPublicKey,
 } from './addresses/pda';
 import { squareRootBN } from './math/utils';
 import {
@@ -4272,4 +4274,111 @@ export class AdminClient extends DriftClient {
 			createAtaIx,
 		];
 	}
+
+	public async initializeConstituent(
+		lpPoolName: number[],
+		spotMarketIndex: number,
+		decimals: number,
+		maxWeightDeviation: BN,
+		swapFeeMin: BN,
+		swapFeeMax: BN
+	): Promise<TransactionSignature> {
+		const ixs = await this.getInitializeConstituentIx(
+			lpPoolName,
+			spotMarketIndex,
+			decimals,
+			maxWeightDeviation,
+			swapFeeMin,
+			swapFeeMax
+		);
+		const tx = await this.buildTransaction(ixs);
+		const { txSig } = await this.sendTransaction(tx, []);
+		return txSig;
+	}
+
+	public async getInitializeConstituentIx(
+		lpPoolName: number[],
+		spotMarketIndex: number,
+		decimals: number,
+		maxWeightDeviation: BN,
+		swapFeeMin: BN,
+		swapFeeMax: BN
+	): Promise<TransactionInstruction[]> {
+		const lpPool = getLpPoolPublicKey(this.program.programId, lpPoolName);
+		const constituentTargetWeights = getConstituentTargetWeightsPublicKey(
+			this.program.programId,
+			lpPool
+		);
+		const constituent = getConstituentPublicKey(
+			this.program.programId,
+			lpPool,
+			spotMarketIndex
+		);
+		return [
+			this.program.instruction.initializeConstituent(
+				lpPoolName,
+				spotMarketIndex,
+				decimals,
+				maxWeightDeviation,
+				swapFeeMin,
+				swapFeeMax,
+				{
+					accounts: {
+						admin: this.wallet.publicKey,
+						lpPool,
+						constituentTargetWeights,
+						constituent,
+						rent: SYSVAR_RENT_PUBKEY,
+						systemProgram: SystemProgram.programId,
+					},
+					signers: [],
+				}
+			),
+		];
+	}
+
+	public async addInitAmmConstituentMappingData(
+		lpPoolName: number[],
+		marketIndexConstituentIndexPairs: InitAmmConstituentMappingDatum[]
+	): Promise<TransactionSignature> {
+		const ixs = await this.getAddInitAmmConstituentMappingDataIx(
+			lpPoolName,
+			marketIndexConstituentIndexPairs
+		);
+		const tx = await this.buildTransaction(ixs);
+		const { txSig } = await this.sendTransaction(tx, []);
+		return txSig;
+	}
+
+	public async getAddInitAmmConstituentMappingDataIx(
+		lpPoolName: number[],
+		marketIndexConstituentIndexPairs: InitAmmConstituentMappingDatum[]
+	): Promise<TransactionInstruction[]> {
+		const lpPool = getLpPoolPublicKey(this.program.programId, lpPoolName);
+		const ammConstituentMapping = getAmmConstituentMappingPublicKey(
+			this.program.programId,
+			lpPool
+		);
+		const constituentTargetWeights = getConstituentTargetWeightsPublicKey(
+			this.program.programId,
+			lpPool
+		);
+		return [
+			this.program.instruction.addAmmConstituentMappingData(
+				lpPoolName,
+				marketIndexConstituentIndexPairs,
+				{
+					accounts: {
+						admin: this.wallet.publicKey,
+						lpPool,
+						ammConstituentMapping,
+						constituentTargetWeights,
+						rent: SYSVAR_RENT_PUBKEY,
+						systemProgram: SystemProgram.programId,
+						state: await this.getStatePublicKey(),
+					},
+				}
+			),
+		];
+	}
 }