program: add pool ids to spot market (#1250)

* init pool work

* add some invariants

* add more invariants

* start adding test

* add ts to update user pool id

* add tests for margin

* fix error message

* disable pools on mainnet

* fix linters

* CHANGELOG

* prettify 2

CHANGELOG.md

@@ -9,6 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Features
 
+- program: add spot market pool ids ([#1250](https://github.com/drift-labs/protocol-v2/pull/1250))
+
 ### Fixes
 
 ### Breaking

package.json

@@ -46,7 +46,8 @@
 		"prettify": "prettier --check './sdk/src/**/*.ts' './tests/**.ts' './cli/**.ts'",
 		"prettify:fix": "prettier --write './sdk/src/**/*.ts' './tests/**.ts' './cli/**.ts'",
 		"lint": "eslint . --ext ts --quiet",
-		"lint:fix": "eslint . --ext ts --fix"
+		"lint:fix": "eslint . --ext ts --fix",
+		"update-idl": "cp target/idl/drift.json sdk/src/idl/drift.json"
 	},
 	"engines": {
 		"node": ">=12"

programs/drift/src/controller/liquidation.rs

@@ -107,6 +107,13 @@ pub fn liquidate_perp(
         "liquidator bankrupt",
     )?;
 
+    validate!(
+        liquidator.pool_id == 0,
+        ErrorCode::InvalidPoolId,
+        "liquidator pool id ({}) != 0",
+        liquidator.pool_id
+    )?;
+
     let market = perp_market_map.get_ref(&market_index)?;
 
     validate!(
@@ -690,6 +697,13 @@ pub fn liquidate_perp_with_fill(
         "user bankrupt",
     )?;
 
+    validate!(
+        liquidator.pool_id == 0,
+        ErrorCode::InvalidPoolId,
+        "liquidator pool id ({}) != 0",
+        liquidator.pool_id
+    )?;
+
     validate!(
         !liquidator.is_bankrupt(),
         ErrorCode::UserBankrupt,
@@ -1155,6 +1169,14 @@ pub fn liquidate_spot(
         asset_market_index
     )?;
 
+    validate!(
+        liquidator.pool_id == asset_spot_market.pool_id,
+        ErrorCode::InvalidPoolId,
+        "liquidator pool id ({}) != asset spot market pool id ({})",
+        liquidator.pool_id,
+        asset_spot_market.pool_id
+    )?;
+
     drop(asset_spot_market);
 
     let liability_spot_market = spot_market_map.get_ref(&liability_market_index)?;
@@ -1166,6 +1188,14 @@ pub fn liquidate_spot(
         liability_market_index
     )?;
 
+    validate!(
+        liquidator.pool_id == liability_spot_market.pool_id,
+        ErrorCode::InvalidPoolId,
+        "liquidator pool id ({}) != liablity spot market pool id ({})",
+        liquidator.pool_id,
+        liability_spot_market.pool_id
+    )?;
+
     drop(liability_spot_market);
 
     // validate user and liquidator have spot balances
@@ -1687,6 +1717,13 @@ pub fn liquidate_borrow_for_perp_pnl(
         "liquidator bankrupt",
     )?;
 
+    validate!(
+        liquidator.pool_id == 0,
+        ErrorCode::InvalidPoolId,
+        "liquidator pool id ({}) != 0",
+        liquidator.pool_id
+    )?;
+
     let perp_market = perp_market_map.get_ref(&perp_market_index)?;
 
     validate!(
@@ -2156,6 +2193,13 @@ pub fn liquidate_perp_pnl_for_deposit(
         "liquidator bankrupt",
     )?;
 
+    validate!(
+        liquidator.pool_id == 0,
+        ErrorCode::InvalidPoolId,
+        "liquidator pool id ({}) != 0",
+        liquidator.pool_id
+    )?;
+
     let asset_spot_market = spot_market_map.get_ref(&asset_market_index)?;
 
     validate!(

programs/drift/src/controller/orders.rs

@@ -168,6 +168,13 @@ pub fn place_perp_order(
         "Market is being initialized"
     )?;
 
+    validate!(
+        user.pool_id == 0,
+        ErrorCode::InvalidPoolId,
+        "user pool id ({}) != 0",
+        user.pool_id
+    )?;
+
     validate!(
         !market.is_in_settlement(now),
         ErrorCode::MarketPlaceOrderPaused,
@@ -1233,6 +1240,14 @@ pub fn fill_perp_order(
     let is_filler_maker = makers_and_referrer.0.contains_key(&filler_key);
     let (mut filler, mut filler_stats) = if !is_filler_maker && !is_filler_taker {
         let filler = load_mut!(filler)?;
+
+        validate!(
+            filler.pool_id == 0,
+            ErrorCode::InvalidPoolId,
+            "filler pool id ({}) != 0",
+            filler.pool_id
+        )?;
+
         if filler.authority != user.authority {
             (Some(filler), Some(load_mut!(filler_stats)?))
         } else {
@@ -3521,6 +3536,13 @@ pub fn place_spot_order(
     let force_reduce_only = spot_market.is_reduce_only();
     let step_size = spot_market.order_step_size;
 
+    validate!(
+        user.pool_id == 0,
+        ErrorCode::InvalidPoolId,
+        "user pool id ({}) != 0",
+        user.pool_id
+    )?;
+
     validate!(
         !matches!(spot_market.status, MarketStatus::Initialized),
         ErrorCode::MarketBeingInitialized,
@@ -3802,6 +3824,14 @@ pub fn fill_spot_order(
     let is_filler_maker = makers_and_referrer.0.contains_key(&filler_key);
     let (mut filler, mut filler_stats) = if !is_filler_maker && !is_filler_taker {
         let filler = load_mut!(filler)?;
+
+        validate!(
+            filler.pool_id == 0,
+            ErrorCode::InvalidPoolId,
+            "filler pool id ({}) != 0",
+            filler.pool_id
+        )?;
+
         if filler.authority != user.authority {
             (Some(filler), Some(load_mut!(filler_stats)?))
         } else {

programs/drift/src/error.rs

@@ -609,6 +609,8 @@ pub enum ErrorCode {
     SwiftOrderDoesNotExist,
     #[msg("Swift order id cannot be 0s")]
     InvalidSwiftOrderId,
+    #[msg("Invalid pool id")]
+    InvalidPoolId,
 }
 
 #[macro_export]

programs/drift/src/instructions/admin.rs

@@ -157,12 +157,6 @@ pub fn handle_initialize_spot_market(
             ErrorCode::InvalidSpotMarketInitialization,
             "For OracleSource::QuoteAsset, oracle must be default public key"
         )?;
-
-        validate!(
-            spot_market_index == QUOTE_SPOT_MARKET_INDEX,
-            ErrorCode::InvalidSpotMarketInitialization,
-            "For OracleSource::QuoteAsset, spot_market_index must be QUOTE_SPOT_MARKET_INDEX"
-        )?;
     } else {
         OracleMap::validate_oracle_account_info(&ctx.accounts.oracle)?;
     }
@@ -314,7 +308,8 @@ pub fn handle_initialize_spot_market(
         fuel_boost_maker: 0,
         fuel_boost_insurance: 0,
         token_program,
-        padding: [0; 41],
+        pool_id: 0,
+        padding: [0; 40],
         insurance_fund: InsuranceFund {
             vault: *ctx.accounts.insurance_fund_vault.to_account_info().key,
             unstaking_period: THIRTEEN_DAY,
@@ -327,6 +322,36 @@ pub fn handle_initialize_spot_market(
     Ok(())
 }
 
+#[access_control(
+    spot_market_valid(&ctx.accounts.spot_market)
+)]
+pub fn handle_update_spot_market_pool_id(
+    ctx: Context<AdminUpdateSpotMarket>,
+    pool_id: u8,
+) -> Result<()> {
+    #[cfg(all(feature = "mainnet-beta", not(feature = "anchor-test")))]
+    {
+        panic!("pools disabled on mainnet-beta");
+    }
+
+    let mut spot_market = load_mut!(ctx.accounts.spot_market)?;
+    msg!(
+        "updating spot market {} pool id to {}",
+        spot_market.market_index,
+        pool_id
+    );
+
+    validate!(
+        spot_market.status == MarketStatus::Initialized,
+        ErrorCode::DefaultError,
+        "Market must be just initialized to update pool"
+    )?;
+
+    spot_market.pool_id = pool_id;
+
+    Ok(())
+}
+
 pub fn handle_initialize_serum_fulfillment_config(
     ctx: Context<InitializeSerumFulfillmentConfig>,
     market_index: u16,
@@ -858,7 +883,7 @@ pub fn handle_initialize_perp_market(
         fuel_boost_position: 0,
         fuel_boost_taker: 0,
         fuel_boost_maker: 0,
-        padding1: 0,
+        pool_id: 0,
         high_leverage_margin_ratio_initial: 0,
         high_leverage_margin_ratio_maintenance: 0,
         padding: [0; 38],

programs/drift/src/instructions/user.rs

@@ -28,6 +28,7 @@ use crate::instructions::optional_accounts::{
 use crate::instructions::SpotFulfillmentType;
 use crate::math::casting::Cast;
 use crate::math::liquidation::is_user_being_liquidated;
+use crate::math::margin::meets_initial_margin_requirement;
 use crate::math::margin::{
     calculate_max_withdrawable_amount, meets_maintenance_margin_requirement,
     meets_place_order_margin_requirement, meets_withdraw_margin_requirement,
@@ -357,6 +358,14 @@ pub fn handle_deposit<'c: 'info, 'info>(
     let mut spot_market = spot_market_map.get_ref_mut(&market_index)?;
     let oracle_price_data = &oracle_map.get_price_data(&spot_market.oracle)?.clone();
 
+    validate!(
+        user.pool_id == spot_market.pool_id,
+        ErrorCode::InvalidPoolId,
+        "user pool id ({}) != market pool id ({})",
+        user.pool_id,
+        spot_market.pool_id
+    )?;
+
     validate!(
         !matches!(spot_market.status, MarketStatus::Initialized),
         ErrorCode::MarketBeingInitialized,
@@ -730,6 +739,14 @@ pub fn handle_transfer_deposit<'c: 'info, 'info>(
     {
         let spot_market = &mut spot_market_map.get_ref_mut(&market_index)?;
 
+        validate!(
+            from_user.pool_id == spot_market.pool_id,
+            ErrorCode::InvalidPoolId,
+            "user pool id ({}) != market pool id ({})",
+            from_user.pool_id,
+            spot_market.pool_id
+        )?;
+
         from_user.increment_total_withdraws(
             amount,
             oracle_price,
@@ -797,6 +814,14 @@ pub fn handle_transfer_deposit<'c: 'info, 'info>(
     {
         let spot_market = &mut spot_market_map.get_ref_mut(&market_index)?;
 
+        validate!(
+            to_user.pool_id == spot_market.pool_id,
+            ErrorCode::InvalidPoolId,
+            "user pool id ({}) != market pool id ({})",
+            to_user.pool_id,
+            spot_market.pool_id
+        )?;
+
         to_user.increment_total_deposits(
             amount,
             oracle_price,
@@ -2191,6 +2216,39 @@ pub fn handle_update_user_margin_trading_enabled<'c: 'info, 'info>(
     Ok(())
 }
 
+pub fn handle_update_user_pool_id<'c: 'info, 'info>(
+    ctx: Context<'_, '_, 'c, 'info, UpdateUser<'info>>,
+    _sub_account_id: u16,
+    pool_id: u8,
+) -> Result<()> {
+    #[cfg(all(feature = "mainnet-beta", not(feature = "anchor-test")))]
+    {
+        panic!("pools disabled on mainnet-beta");
+    }
+
+    let remaining_accounts_iter = &mut ctx.remaining_accounts.iter().peekable();
+    let AccountMaps {
+        perp_market_map,
+        spot_market_map,
+        mut oracle_map,
+        ..
+    } = load_maps(
+        remaining_accounts_iter,
+        &MarketSet::new(),
+        &MarketSet::new(),
+        Clock::get()?.slot,
+        None,
+    )?;
+
+    let mut user = load_mut!(ctx.accounts.user)?;
+    user.pool_id = pool_id;
+
+    // will throw if user has deposits/positions in other pools
+    meets_initial_margin_requirement(&user, &perp_market_map, &spot_market_map, &mut oracle_map)?;
+
+    Ok(())
+}
+
 pub fn handle_update_user_delegate(
     ctx: Context<UpdateUser>,
     _sub_account_id: u16,

programs/drift/src/lib.rs

@@ -319,6 +319,14 @@ pub mod drift {
         handle_update_user_margin_trading_enabled(ctx, _sub_account_id, margin_trading_enabled)
     }
 
+    pub fn update_user_pool_id<'c: 'info, 'info>(
+        ctx: Context<'_, '_, 'c, 'info, UpdateUser<'info>>,
+        _sub_account_id: u16,
+        pool_id: u8,
+    ) -> Result<()> {
+        handle_update_user_pool_id(ctx, _sub_account_id, pool_id)
+    }
+
     pub fn update_user_delegate(
         ctx: Context<UpdateUser>,
         _sub_account_id: u16,
@@ -1011,6 +1019,13 @@ pub mod drift {
         handle_update_insurance_fund_unstaking_period(ctx, insurance_fund_unstaking_period)
     }
 
+    pub fn update_spot_market_pool_id(
+        ctx: Context<AdminUpdateSpotMarket>,
+        pool_id: u8,
+    ) -> Result<()> {
+        handle_update_spot_market_pool_id(ctx, pool_id)
+    }
+
     pub fn update_spot_market_liquidation_fee(
         ctx: Context<AdminUpdateSpotMarket>,
         liquidator_fee: u32,