bigz/fix-imbalanced-base-asset-amount-with-lp (#604)

* bigz/fix-imbalanced_base_asset_amount_with_lp

* add assert in test.rs

* add test assert

* sleep longer in test

* add changelog

CHANGELOG.md

@@ -17,6 +17,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - program: let auction start/end be the same ([#597](https://github.com/drift-labs/protocol-v2/pull/597))
 - program: account for reduce only when checking margin in trigger order ([#583](https://github.com/drift-labs/protocol-v2/pull/583))
+- program: use per_lp_base_unit for calculating base imbalance for lp jit ([#604](https://github.com/drift-labs/protocol-v2/pull/604))
 
 ### Breaking
 

programs/drift/src/controller/position/tests.rs

@@ -5,7 +5,6 @@ use crate::controller::position::{
 use crate::controller::lp::{apply_lp_rebase_to_perp_market, settle_lp_position};
 
 use crate::controller::repeg::_update_amm;
-// use crate::instructions::handle_update_perp_market_per_lp_base;
 use crate::math::constants::{
     AMM_RESERVE_PRECISION, AMM_RESERVE_PRECISION_I128, BASE_PRECISION_I64, PRICE_PRECISION_I64,
     PRICE_PRECISION_U64, QUOTE_PRECISION_I128,
@@ -283,6 +282,16 @@ fn amm_split_large_k_with_rebase() {
 
     assert_eq!(existing_position.per_lp_base, perp_market.amm.per_lp_base);
 
+    assert_eq!(
+        perp_market
+            .amm
+            .imbalanced_base_asset_amount_with_lp()
+            .unwrap(),
+        -303686915482213
+    );
+
+    assert_eq!(perp_market.amm.target_base_asset_amount_per_lp, -565000000);
+
     // update base back
     let base_change = -2;
     apply_lp_rebase_to_perp_market(&mut perp_market, base_change).unwrap();
@@ -292,6 +301,15 @@ fn amm_split_large_k_with_rebase() {
         quote_asset_amount: 0,
     };
 
+    // unchanged with rebase (even when target!=0)
+    assert_eq!(
+        perp_market
+            .amm
+            .imbalanced_base_asset_amount_with_lp()
+            .unwrap(),
+        -303686915482213
+    );
+
     update_lp_market_position(&mut perp_market, &delta, 0, AMMLiquiditySplit::Shared).unwrap();
 
     assert_eq!(

programs/drift/src/state/perp_market.rs

@@ -11,9 +11,8 @@ use crate::math::constants::{
     AMM_RESERVE_PRECISION, MAX_CONCENTRATION_COEFFICIENT, PRICE_PRECISION_I64,
 };
 use crate::math::constants::{
-    AMM_RESERVE_PRECISION_I128, BASE_PRECISION, BID_ASK_SPREAD_PRECISION_U128,
-    LP_FEE_SLICE_DENOMINATOR, LP_FEE_SLICE_NUMERATOR, MARGIN_PRECISION_U128, SPOT_WEIGHT_PRECISION,
-    TWENTY_FOUR_HOUR,
+    AMM_RESERVE_PRECISION_I128, BID_ASK_SPREAD_PRECISION_U128, LP_FEE_SLICE_DENOMINATOR,
+    LP_FEE_SLICE_NUMERATOR, MARGIN_PRECISION_U128, SPOT_WEIGHT_PRECISION, TWENTY_FOUR_HOUR,
 };
 use crate::math::helpers::get_proportion_i128;
 
@@ -840,7 +839,9 @@ impl AMM {
             .base_asset_amount_per_lp
             .safe_sub(self.get_target_base_asset_amount_per_lp()?)?;
 
-        get_proportion_i128(target_lp_gap, self.user_lp_shares, BASE_PRECISION)
+        let base_unit = self.get_per_lp_base_unit()?.cast()?;
+
+        get_proportion_i128(target_lp_gap, self.user_lp_shares, base_unit)
     }
 
     pub fn amm_wants_to_jit_make(&self, taker_direction: PositionDirection) -> DriftResult<bool> {

test-scripts/single-anchor-test.sh

@@ -4,7 +4,7 @@ if [ "$1" != "--skip-build" ]
     cp target/idl/drift.json sdk/src/idl/
 fi
 
-test_files=(prepegMarketOrderBaseAssetAmount.ts)
+test_files=(liquidityProvider.ts)
 
 for test_file in ${test_files[@]}; do
   ANCHOR_TEST_FILE=${test_file} anchor test --skip-build || exit 1;

tests/liquidityProvider.ts

@@ -1247,7 +1247,7 @@ describe('liquidity providing', () => {
 		const txSig1 = await driftClient.updatePerpMarketPerLpBase(0, 1);
 		await _viewLogs(txSig1);
 
-		await sleep(500); // todo?
+		await sleep(1400); // todo?
 		await driftClient.fetchAccounts();
 		const marketAfter = driftClient.getPerpMarketAccount(0);