@drift-labs/sdk 0.2.0-master.19 → 0.2.0-master.21

package/src/clearingHouse.ts

@@ -15,6 +15,7 @@ import {
 	OptionalOrderParams,
 	DefaultOrderParams,
 	OrderType,
+	ReferrerInfo,
 } from './types';
 import * as anchor from '@project-serum/anchor';
 import clearingHouseIDL from './idl/clearing_house.json';
@@ -51,7 +52,11 @@ import {
 } from './accounts/types';
 import { TxSender } from './tx/types';
 import { wrapInTx } from './tx/utils';
-import { QUOTE_ASSET_BANK_INDEX, ZERO } from './constants/numericConstants';
+import {
+	ONE,
+	QUOTE_ASSET_BANK_INDEX,
+	ZERO,
+} from './constants/numericConstants';
 import { findDirectionToClose, positionIsAvailable } from './math/position';
 import { getTokenAmount } from './math/bankBalance';
 import { DEFAULT_USER_NAME, encodeName } from './userName';
@@ -219,11 +224,13 @@ export class ClearingHouse {
 	 *	Forces the accountSubscriber to fetch account updates from rpc
 	 */
 	public async fetchAccounts(): Promise<void> {
-		await Promise.all(
-			[...this.users.values()]
-				.map((user) => user.fetchAccounts())
-				.concat(this.accountSubscriber.fetch())
-		);
+		const promises = [...this.users.values()]
+			.map((user) => user.fetchAccounts())
+			.concat(this.accountSubscriber.fetch());
+		if (this.userStats) {
+			promises.concat(this.userStats.fetchAccounts());
+		}
+		await Promise.all(promises);
 	}
 
 	public async unsubscribe(): Promise<void> {
@@ -341,10 +348,11 @@ export class ClearingHouse {
 
 	public async initializeUserAccount(
 		userId = 0,
-		name = DEFAULT_USER_NAME
+		name = DEFAULT_USER_NAME,
+		referrerInfo?: ReferrerInfo
 	): Promise<[TransactionSignature, PublicKey]> {
 		const [userAccountPublicKey, initializeUserAccountIx] =
-			await this.getInitializeUserInstructions(userId, name);
+			await this.getInitializeUserInstructions(userId, name, referrerInfo);
 
 		const tx = new Transaction();
 		if (userId === 0) {
@@ -358,7 +366,8 @@ export class ClearingHouse {
 
 	async getInitializeUserInstructions(
 		userId = 0,
-		name = DEFAULT_USER_NAME
+		name = DEFAULT_USER_NAME,
+		referrerInfo?: ReferrerInfo
 	): Promise<[PublicKey, TransactionInstruction]> {
 		const userAccountPublicKey = await getUserAccountPublicKey(
 			this.program.programId,
@@ -366,6 +375,20 @@ export class ClearingHouse {
 			userId
 		);
 
+		const remainingAccounts = new Array<AccountMeta>();
+		if (referrerInfo !== undefined) {
+			remainingAccounts.push({
+				pubkey: referrerInfo.referrer,
+				isWritable: true,
+				isSigner: false,
+			});
+			remainingAccounts.push({
+				pubkey: referrerInfo.referrerStats,
+				isWritable: true,
+				isSigner: false,
+			});
+		}
+
 		const nameBuffer = encodeName(name);
 		const initializeUserAccountIx =
 			await this.program.instruction.initializeUser(userId, nameBuffer, {
@@ -378,6 +401,7 @@ export class ClearingHouse {
 					systemProgram: anchor.web3.SystemProgram.programId,
 					state: await this.getStatePublicKey(),
 				},
+				remainingAccounts,
 			});
 
 		return [userAccountPublicKey, initializeUserAccountIx];
@@ -711,6 +735,7 @@ export class ClearingHouse {
 					bank: bank.pubkey,
 					bankVault: bank.vault,
 					user: userAccountPublicKey,
+					userStats: this.getUserStatsAccountPublicKey(),
 					userTokenAccount: userTokenAccount,
 					authority: this.wallet.publicKey,
 					tokenProgram: TOKEN_PROGRAM_ID,
@@ -842,10 +867,11 @@ export class ClearingHouse {
 		bankIndex = new BN(0),
 		userId = 0,
 		name = DEFAULT_USER_NAME,
-		fromUserId?: number
+		fromUserId?: number,
+		referrerInfo?: ReferrerInfo
 	): Promise<[TransactionSignature, PublicKey]> {
 		const [userAccountPublicKey, initializeUserAccountIx] =
-			await this.getInitializeUserInstructions(userId, name);
+			await this.getInitializeUserInstructions(userId, name, referrerInfo);
 
 		const additionalSigners: Array<Signer> = [];
 
@@ -920,7 +946,8 @@ export class ClearingHouse {
 		name = DEFAULT_USER_NAME,
 		bankIndex: BN,
 		tokenFaucet: TokenFaucet,
-		amount: BN
+		amount: BN,
+		referrerInfo?: ReferrerInfo
 	): Promise<[TransactionSignature, PublicKey]> {
 		const [associateTokenPublicKey, createAssociatedAccountIx, mintToIx] =
 			await tokenFaucet.createAssociatedTokenAccountAndMintToInstructions(
@@ -929,7 +956,7 @@ export class ClearingHouse {
 			);
 
 		const [userAccountPublicKey, initializeUserAccountIx] =
-			await this.getInitializeUserInstructions(userId, name);
+			await this.getInitializeUserInstructions(userId, name, referrerInfo);
 
 		const depositCollateralIx = await this.getDepositInstruction(
 			amount,
@@ -1038,6 +1065,7 @@ export class ClearingHouse {
 					bankVault: bank.vault,
 					bankVaultAuthority: bank.vaultAuthority,
 					user: userAccountPublicKey,
+					userStats: this.getUserStatsAccountPublicKey(),
 					userTokenAccount: userTokenAccount,
 					authority: this.wallet.publicKey,
 					tokenProgram: TOKEN_PROGRAM_ID,
@@ -1089,6 +1117,7 @@ export class ClearingHouse {
 				authority: this.wallet.publicKey,
 				fromUser,
 				toUser,
+				userStats: this.getUserStatsAccountPublicKey(),
 				state: await this.getStatePublicKey(),
 			},
 			remainingAccounts,
@@ -1401,11 +1430,18 @@ export class ClearingHouse {
 		userAccountPublicKey: PublicKey,
 		user: UserAccount,
 		order?: Order,
-		makerInfo?: MakerInfo
+		makerInfo?: MakerInfo,
+		referrerInfo?: ReferrerInfo
 	): Promise<TransactionSignature> {
 		const { txSig } = await this.txSender.send(
 			wrapInTx(
-				await this.getFillOrderIx(userAccountPublicKey, user, order, makerInfo)
+				await this.getFillOrderIx(
+					userAccountPublicKey,
+					user,
+					order,
+					makerInfo,
+					referrerInfo
+				)
 			),
 			[],
 			this.opts
@@ -1417,7 +1453,8 @@ export class ClearingHouse {
 		userAccountPublicKey: PublicKey,
 		userAccount: UserAccount,
 		order: Order,
-		makerInfo?: MakerInfo
+		makerInfo?: MakerInfo,
+		referrerInfo?: ReferrerInfo
 	): Promise<TransactionInstruction> {
 		const userStatsPublicKey = getUserStatsAccountPublicKey(
 			this.program.programId,
@@ -1427,7 +1464,11 @@ export class ClearingHouse {
 		const fillerPublicKey = await this.getUserAccountPublicKey();
 		const fillerStatsPublicKey = this.getUserStatsAccountPublicKey();
 
-		const marketIndex = order.marketIndex;
+		const marketIndex = order
+			? order.marketIndex
+			: userAccount.orders.find((order) =>
+					order.orderId.eq(userAccount.nextOrderId.sub(ONE))
+			  ).marketIndex;
 		const marketAccount = this.getMarketAccount(marketIndex);
 
 		const oracleAccountMap = new Map<string, AccountMeta>();
@@ -1502,6 +1543,19 @@ export class ClearingHouse {
 			});
 		}
 
+		if (referrerInfo) {
+			remainingAccounts.push({
+				pubkey: referrerInfo.referrer,
+				isWritable: true,
+				isSigner: false,
+			});
+			remainingAccounts.push({
+				pubkey: referrerInfo.referrerStats,
+				isWritable: true,
+				isSigner: false,
+			});
+		}
+
 		const orderId = order.orderId;
 		const makerOrderId = makerInfo ? makerInfo.order.orderId : null;
 
@@ -1614,10 +1668,13 @@ export class ClearingHouse {
 
 	public async placeAndTake(
 		orderParams: OptionalOrderParams,
-		makerInfo?: MakerInfo
+		makerInfo?: MakerInfo,
+		referrerInfo?: ReferrerInfo
 	): Promise<TransactionSignature> {
 		const { txSig, slot } = await this.txSender.send(
-			wrapInTx(await this.getPlaceAndTakeIx(orderParams, makerInfo)),
+			wrapInTx(
+				await this.getPlaceAndTakeIx(orderParams, makerInfo, referrerInfo)
+			),
 			[],
 			this.opts
 		);
@@ -1627,7 +1684,8 @@ export class ClearingHouse {
 
 	public async getPlaceAndTakeIx(
 		orderParams: OptionalOrderParams,
-		makerInfo?: MakerInfo
+		makerInfo?: MakerInfo,
+		referrerInfo?: ReferrerInfo
 	): Promise<TransactionInstruction> {
 		orderParams = this.getOrderParams(orderParams);
 		const userStatsPublicKey = await this.getUserStatsAccountPublicKey();
@@ -1653,6 +1711,19 @@ export class ClearingHouse {
 			});
 		}
 
+		if (referrerInfo) {
+			remainingAccounts.push({
+				pubkey: referrerInfo.referrer,
+				isWritable: true,
+				isSigner: false,
+			});
+			remainingAccounts.push({
+				pubkey: referrerInfo.referrerStats,
+				isWritable: true,
+				isSigner: false,
+			});
+		}
+
 		return await this.program.instruction.placeAndTake(
 			orderParams,
 			makerOrderId,
@@ -1670,10 +1741,13 @@ export class ClearingHouse {
 
 	public async placeAndMake(
 		orderParams: OptionalOrderParams,
-		takerInfo: TakerInfo
+		takerInfo: TakerInfo,
+		referrerInfo?: ReferrerInfo
 	): Promise<TransactionSignature> {
 		const { txSig, slot } = await this.txSender.send(
-			wrapInTx(await this.getPlaceAndMakeIx(orderParams, takerInfo)),
+			wrapInTx(
+				await this.getPlaceAndMakeIx(orderParams, takerInfo, referrerInfo)
+			),
 			[],
 			this.opts
 		);
@@ -1685,16 +1759,32 @@ export class ClearingHouse {
 
 	public async getPlaceAndMakeIx(
 		orderParams: OptionalOrderParams,
-		takerInfo: TakerInfo
+		takerInfo: TakerInfo,
+		referrerInfo?: ReferrerInfo
 	): Promise<TransactionInstruction> {
 		orderParams = this.getOrderParams(orderParams);
 		const userStatsPublicKey = this.getUserStatsAccountPublicKey();
 		const userAccountPublicKey = await this.getUserAccountPublicKey();
 
-		const remainingAccounts = this.getRemainingAccounts({
+		// todo merge this with getRemainingAccounts
+		const remainingAccounts = this.getRemainingAccountsWithCounterparty({
+			counterPartyUserAccount: takerInfo.takerUserAccount,
 			writableMarketIndex: orderParams.marketIndex,
 		});
 
+		if (referrerInfo) {
+			remainingAccounts.push({
+				pubkey: referrerInfo.referrer,
+				isWritable: true,
+				isSigner: false,
+			});
+			remainingAccounts.push({
+				pubkey: referrerInfo.referrerStats,
+				isWritable: true,
+				isSigner: false,
+			});
+		}
+
 		const takerOrderId = takerInfo!.order!.orderId;
 		return await this.program.instruction.placeAndMake(
 			orderParams,
@@ -1758,7 +1848,7 @@ export class ClearingHouse {
 		const tx = new Transaction()
 			.add(
 				ComputeBudgetProgram.requestUnits({
-					units: 1_200_000,
+					units: 1_000_000,
 					additionalFee: 0,
 				})
 			)
@@ -1899,9 +1989,9 @@ export class ClearingHouse {
 		const liquidatorPublicKey = await this.getUserAccountPublicKey();
 		const liquidatorStatsPublicKey = this.getUserStatsAccountPublicKey();
 
-		const remainingAccounts = this.getRemainingAccountsForLiquidation({
+		const remainingAccounts = this.getRemainingAccountsWithCounterparty({
 			writableMarketIndex: marketIndex,
-			userAccount,
+			counterPartyUserAccount: userAccount,
 		});
 
 		return await this.program.instruction.liquidatePerp(
@@ -1953,8 +2043,8 @@ export class ClearingHouse {
 	): Promise<TransactionInstruction> {
 		const liquidatorPublicKey = await this.getUserAccountPublicKey();
 
-		const remainingAccounts = this.getRemainingAccountsForLiquidation({
-			userAccount,
+		const remainingAccounts = this.getRemainingAccountsWithCounterparty({
+			counterPartyUserAccount: userAccount,
 			writableBankIndexes: [liabilityBankIndex, assetBankIndex],
 		});
 
@@ -2006,8 +2096,8 @@ export class ClearingHouse {
 	): Promise<TransactionInstruction> {
 		const liquidatorPublicKey = await this.getUserAccountPublicKey();
 
-		const remainingAccounts = this.getRemainingAccountsForLiquidation({
-			userAccount,
+		const remainingAccounts = this.getRemainingAccountsWithCounterparty({
+			counterPartyUserAccount: userAccount,
 			writableMarketIndex: perpMarketIndex,
 			writableBankIndexes: [liabilityBankIndex],
 		});
@@ -2060,8 +2150,8 @@ export class ClearingHouse {
 	): Promise<TransactionInstruction> {
 		const liquidatorPublicKey = await this.getUserAccountPublicKey();
 
-		const remainingAccounts = this.getRemainingAccountsForLiquidation({
-			userAccount,
+		const remainingAccounts = this.getRemainingAccountsWithCounterparty({
+			counterPartyUserAccount: userAccount,
 			writableMarketIndex: perpMarketIndex,
 			writableBankIndexes: [assetBankIndex],
 		});
@@ -2108,9 +2198,9 @@ export class ClearingHouse {
 	): Promise<TransactionInstruction> {
 		const liquidatorPublicKey = await this.getUserAccountPublicKey();
 
-		const remainingAccounts = this.getRemainingAccountsForLiquidation({
+		const remainingAccounts = this.getRemainingAccountsWithCounterparty({
 			writableMarketIndex: marketIndex,
-			userAccount,
+			counterPartyUserAccount: userAccount,
 		});
 
 		return await this.program.instruction.resolvePerpBankruptcy(marketIndex, {
@@ -2150,9 +2240,9 @@ export class ClearingHouse {
 	): Promise<TransactionInstruction> {
 		const liquidatorPublicKey = await this.getUserAccountPublicKey();
 
-		const remainingAccounts = this.getRemainingAccountsForLiquidation({
+		const remainingAccounts = this.getRemainingAccountsWithCounterparty({
 			writableBankIndexes: [bankIndex],
-			userAccount,
+			counterPartyUserAccount: userAccount,
 		});
 
 		return await this.program.instruction.resolveBorrowBankruptcy(bankIndex, {
@@ -2166,17 +2256,17 @@ export class ClearingHouse {
 		});
 	}
 
-	getRemainingAccountsForLiquidation(params: {
-		userAccount: UserAccount;
+	getRemainingAccountsWithCounterparty(params: {
+		counterPartyUserAccount: UserAccount;
 		writableMarketIndex?: BN;
 		writableBankIndexes?: BN[];
 	}): AccountMeta[] {
-		const liquidateeUserAccount = params.userAccount;
+		const counterPartyUserAccount = params.counterPartyUserAccount;
 
 		const oracleAccountMap = new Map<string, AccountMeta>();
 		const bankAccountMap = new Map<number, AccountMeta>();
 		const marketAccountMap = new Map<number, AccountMeta>();
-		for (const bankBalance of liquidateeUserAccount.bankBalances) {
+		for (const bankBalance of counterPartyUserAccount.bankBalances) {
 			if (!bankBalance.balance.eq(ZERO)) {
 				const bankAccount = this.getBankAccount(bankBalance.bankIndex);
 				bankAccountMap.set(bankBalance.bankIndex.toNumber(), {
@@ -2194,7 +2284,7 @@ export class ClearingHouse {
 				}
 			}
 		}
-		for (const position of liquidateeUserAccount.positions) {
+		for (const position of counterPartyUserAccount.positions) {
 			if (!positionIsAvailable(position)) {
 				const market = this.getMarketAccount(position.marketIndex);
 				marketAccountMap.set(position.marketIndex.toNumber(), {

package/src/clearingHouseUser.ts

@@ -33,7 +33,7 @@ import {
 	calculateBaseAssetValue,
 	calculatePositionFundingPNL,
 	calculatePositionPNL,
-	calculateUnsettledAssetWeight,
+	calculateUnrealizedAssetWeight,
 	calculateMarketMarginRatio,
 	PositionDirection,
 	calculateTradeSlippage,
@@ -144,6 +144,11 @@ export class ClearingHouseUser {
 		};
 	}
 
+	public getClonedPosition(position: UserPosition): UserPosition {
+		const clonedPosition = Object.assign({}, position);
+		return clonedPosition;
+	}
+
 	/**
 	 * @param orderId
 	 * @returns Order
@@ -180,8 +185,10 @@ export class ClearingHouseUser {
 	 * calculates the market position if the lp position was settled
 	 * @returns : userPosition
 	 */
-	public getSettledLPPosition(marketIndex: BN): [UserPosition, BN] {
-		const position = this.getUserPosition(marketIndex);
+	public getSettledLPPosition(marketIndex: BN): [UserPosition, BN, BN] {
+		const _position = this.getUserPosition(marketIndex);
+		const position = this.getClonedPosition(_position);
+
 		const market = this.clearingHouse.getMarketAccount(position.marketIndex);
 		const nShares = position.lpShares;
 
@@ -265,7 +272,7 @@ export class ClearingHouseUser {
 			position.lastCumulativeFundingRate = ZERO;
 		}
 
-		return [position, pnl];
+		return [position, remainderBaa, pnl];
 	}
 
 	/**
@@ -289,51 +296,55 @@ export class ClearingHouseUser {
 		return freeCollateral.gte(ZERO) ? freeCollateral : ZERO;
 	}
 
-	public getInitialMarginRequirement(): BN {
-		const postionMarginRequirement = this.getUserAccount().positions.reduce(
-			(marginRequirement, marketPosition) => {
+	/**
+	 * @returns The margin requirement of a certain type (Initial or Maintenance) in USDC. : QUOTE_PRECISION
+	 */
+	public getMarginRequirement(type: MarginCategory): BN {
+		return this.getUserAccount()
+			.positions.reduce((marginRequirement, marketPosition) => {
 				const market = this.clearingHouse.getMarketAccount(
 					marketPosition.marketIndex
 				);
-				const worstCaseBaseAssetAmount =
-					calculateWorstCaseBaseAssetAmount(marketPosition);
 
-				const worstCaseAssetValue = worstCaseBaseAssetAmount
-					.abs()
-					.mul(this.getOracleDataForMarket(market.marketIndex).price)
-					.div(AMM_TO_QUOTE_PRECISION_RATIO.mul(MARK_PRICE_PRECISION));
+				if (marketPosition.lpShares.gt(ZERO)) {
+					// is an lp
+					// clone so we dont mutate the position
+					marketPosition = this.getClonedPosition(marketPosition);
 
-				const marketMarginRatio = new BN(
-					calculateMarketMarginRatio(
-						market,
-						worstCaseBaseAssetAmount.abs(),
-						'Initial'
-					)
-				);
-				return marginRequirement.add(
-					worstCaseAssetValue.mul(marketMarginRatio).div(MARGIN_PRECISION)
-				);
-			},
-			ZERO
-		);
+					// settle position
+					const [settledPosition, dustBaa, _] = this.getSettledLPPosition(
+						market.marketIndex
+					);
+					marketPosition.baseAssetAmount =
+						settledPosition.baseAssetAmount.add(dustBaa);
+					marketPosition.quoteAssetAmount = settledPosition.quoteAssetAmount;
+
+					// open orders
+					let openAsks;
+					if (market.amm.maxBaseAssetReserve > market.amm.baseAssetReserve) {
+						openAsks = market.amm.maxBaseAssetReserve
+							.sub(market.amm.baseAssetReserve)
+							.mul(marketPosition.lpShares)
+							.div(market.amm.sqrtK)
+							.mul(new BN(-1));
+					} else {
+						openAsks = ZERO;
+					}
 
-		const bankMarginRequirement = this.getBankLiabilityValue(
-			undefined,
-			'Initial'
-		);
+					let openBids;
+					if (market.amm.minBaseAssetReserve < market.amm.baseAssetReserve) {
+						openBids = market.amm.baseAssetReserve
+							.sub(market.amm.minBaseAssetReserve)
+							.mul(marketPosition.lpShares)
+							.div(market.amm.sqrtK);
+					} else {
+						openBids = ZERO;
+					}
 
-		return bankMarginRequirement.add(postionMarginRequirement);
-	}
+					marketPosition.openAsks = marketPosition.openAsks.add(openAsks);
+					marketPosition.openBids = marketPosition.openBids.add(openBids);
+				}
 
-	/**
-	 * @returns The maintenance margin requirement in USDC. : QUOTE_PRECISION
-	 */
-	public getMaintenanceMarginRequirement(): BN {
-		return this.getUserAccount()
-			.positions.reduce((marginRequirement, marketPosition) => {
-				const market = this.clearingHouse.getMarketAccount(
-					marketPosition.marketIndex
-				);
 				const worstCaseBaseAssetAmount =
 					calculateWorstCaseBaseAssetAmount(marketPosition);
 
@@ -349,14 +360,28 @@ export class ClearingHouseUser {
 								calculateMarketMarginRatio(
 									market,
 									worstCaseBaseAssetAmount.abs(),
-									'Maintenance'
+									type
 								)
 							)
 						)
 						.div(MARGIN_PRECISION)
 				);
 			}, ZERO)
-			.add(this.getBankLiabilityValue(undefined, 'Maintenance'));
+			.add(this.getBankLiabilityValue(undefined, type));
+	}
+
+	/**
+	 * @returns The initial margin requirement in USDC. : QUOTE_PRECISION
+	 */
+	public getInitialMarginRequirement(): BN {
+		return this.getMarginRequirement('Initial');
+	}
+
+	/**
+	 * @returns The maintenance margin requirement in USDC. : QUOTE_PRECISION
+	 */
+	public getMaintenanceMarginRequirement(): BN {
+		return this.getMarginRequirement('Maintenance');
 	}
 
 	/**
@@ -372,11 +397,11 @@ export class ClearingHouseUser {
 			.positions.filter((pos) =>
 				marketIndex ? pos.marketIndex === marketIndex : true
 			)
-			.reduce((pnl, marketPosition) => {
+			.reduce((unrealizedPnl, marketPosition) => {
 				const market = this.clearingHouse.getMarketAccount(
 					marketPosition.marketIndex
 				);
-				let pnl0 = calculatePositionPNL(
+				let positionUnrealizedPnl = calculatePositionPNL(
 					market,
 					marketPosition,
 					withFunding,
@@ -384,12 +409,12 @@ export class ClearingHouseUser {
 				);
 
 				if (withWeightMarginCategory !== undefined) {
-					if (pnl0.gt(ZERO)) {
-						pnl0 = pnl0
+					if (positionUnrealizedPnl.gt(ZERO)) {
+						positionUnrealizedPnl = positionUnrealizedPnl
 							.mul(
-								calculateUnsettledAssetWeight(
+								calculateUnrealizedAssetWeight(
 									market,
-									pnl0,
+									positionUnrealizedPnl,
 									withWeightMarginCategory
 								)
 							)
@@ -397,7 +422,7 @@ export class ClearingHouseUser {
 					}
 				}
 
-				return pnl.add(pnl0);
+				return unrealizedPnl.add(positionUnrealizedPnl);
 			}, ZERO);
 	}
 
@@ -528,47 +553,9 @@ export class ClearingHouseUser {
 	 * @returns : Precision QUOTE_PRECISION
 	 */
 	public getTotalCollateral(marginCategory: MarginCategory = 'Initial'): BN {
-		return this.getUserAccount()
-			.bankBalances.reduce((totalAssetValue, bankBalance) => {
-				if (
-					bankBalance.balance.eq(ZERO) ||
-					isVariant(bankBalance.balanceType, 'borrow')
-				) {
-					return totalAssetValue;
-				}
-
-				// Todo this needs to account for whether it's based on initial or maintenance requirements
-				const bankAccount: BankAccount = this.clearingHouse.getBankAccount(
-					bankBalance.bankIndex
-				);
-
-				const tokenAmount = getTokenAmount(
-					bankBalance.balance,
-					bankAccount,
-					bankBalance.balanceType
-				);
-
-				const weight = calculateAssetWeight(
-					tokenAmount,
-					bankAccount,
-					marginCategory
-				);
-
-				return totalAssetValue.add(
-					tokenAmount
-						.mul(this.getOracleDataForBank(bankAccount.bankIndex).price)
-						.mul(weight)
-						.div(BANK_WEIGHT_PRECISION)
-						.div(MARK_PRICE_PRECISION)
-						// Adjust for decimals of bank account
-						.div(
-							new BN(10).pow(
-								new BN(bankAccount.decimals).sub(BANK_BALANCE_PRECISION_EXP)
-							)
-						)
-				);
-			}, ZERO)
-			.add(this.getUnrealizedPNL(true, undefined, marginCategory));
+		return this.getBankAssetValue(undefined, marginCategory).add(
+			this.getUnrealizedPNL(true, undefined, marginCategory)
+		);
 	}
 
 	/**

package/src/clearingHouseUserStats.ts

@@ -4,7 +4,11 @@ import { DataAndSlot, UserStatsAccountSubscriber } from './accounts/types';
 import { ClearingHouseUserStatsConfig } from './clearingHouseUserStatsConfig';
 import { PollingUserStatsAccountSubscriber } from './accounts/pollingUserStatsAccountSubscriber';
 import { WebSocketUserStatsAccountSubscriber } from './accounts/webSocketUserStatsAccountSubsriber';
-import { UserStatsAccount } from './types';
+import { ReferrerInfo, UserStatsAccount } from './types';
+import {
+	getUserAccountPublicKeySync,
+	getUserStatsAccountPublicKey,
+} from './addresses/pda';
 
 export class ClearingHouseUserStats {
 	clearingHouse: ClearingHouse;
@@ -50,4 +54,22 @@ export class ClearingHouseUserStats {
 	public getAccount(): UserStatsAccount {
 		return this.accountSubscriber.getUserStatsAccountAndSlot().data;
 	}
+
+	public getReferrerInfo(): ReferrerInfo | undefined {
+		if (this.getAccount().referrer.equals(PublicKey.default)) {
+			return undefined;
+		} else {
+			return {
+				referrer: getUserAccountPublicKeySync(
+					this.clearingHouse.program.programId,
+					this.getAccount().referrer,
+					0
+				),
+				referrerStats: getUserStatsAccountPublicKey(
+					this.clearingHouse.program.programId,
+					this.getAccount().referrer
+				),
+			};
+		}
+	}
 }