sdk: add getMaxSwapAmount

sdk/src/math/utils.ts

@@ -34,6 +34,10 @@ export const divCeil = (a: BN, b: BN): BN => {
 	}
 };
 
+export const sigNum = (x: BN): BN => {
+	return x.isNeg() ? new BN(-1) : new BN(1);
+};
+
 /**
  * calculates the time remaining until the next update based on a rounded, "on-the-hour" update schedule
  * this schedule is used for Perpetual Funding Rate and Revenue -> Insurance Updates

sdk/src/user.ts

@@ -28,6 +28,8 @@ import {
 	OPEN_ORDER_MARGIN_REQUIREMENT,
 	FIVE_MINUTE,
 	BASE_PRECISION,
+	ONE,
+	TWO,
 } from './constants/numericConstants';
 import {
 	UserAccountSubscriber,
@@ -50,6 +52,8 @@ import {
 	calculateSpotMarketMarginRatio,
 	getSignedTokenAmount,
 	SpotBalanceType,
+	sigNum,
+	getBalance,
 } from '.';
 import {
 	getTokenAmount,
@@ -162,6 +166,18 @@ export class User {
 		);
 	}
 
+	getEmptySpotPosition(marketIndex: number): SpotPosition {
+		return {
+			marketIndex,
+			scaledBalance: ZERO,
+			balanceType: SpotBalanceType.DEPOSIT,
+			cumulativeDeposits: ZERO,
+			openAsks: ZERO,
+			openBids: ZERO,
+			openOrders: 0,
+		};
+	}
+
 	/**
 	 * Returns the token amount for a given market. The spot market precision is based on the token mint decimals.
 	 * Positive if it is a deposit, negative if it is a borrow.
@@ -436,7 +452,7 @@ export class User {
 	 * @returns : Precision QUOTE_PRECISION
 	 */
 	public getFreeCollateral(): BN {
-		const totalCollateral = this.getTotalCollateral();
+		const totalCollateral = this.getTotalCollateral('Initial', true);
 		const initialMarginRequirement = this.getInitialMarginRequirement();
 		const freeCollateral = totalCollateral.sub(initialMarginRequirement);
 		return freeCollateral.gte(ZERO) ? freeCollateral : ZERO;
@@ -1944,6 +1960,212 @@ export class User {
 		return tradeAmount;
 	}
 
+	/**
+	 * Calculates the max amount of token that can be swapped from inMarket to outMarket
+	 * Assumes swap happens at oracle price
+	 *
+	 * @param inMarketIndex
+	 * @param outMarketIndex
+	 * @param marginTradingEnabled
+	 */
+	public getMaxSwapAmount({
+		inMarketIndex,
+		outMarketIndex,
+	}: {
+		inMarketIndex: number;
+		outMarketIndex: number;
+	}): BN {
+		const inMarket = this.driftClient.getSpotMarketAccount(inMarketIndex);
+		const outMarket = this.driftClient.getSpotMarketAccount(outMarketIndex);
+
+		const inOraclePrice = this.getOracleDataForSpotMarket(inMarketIndex).price;
+		const outOraclePrice =
+			this.getOracleDataForSpotMarket(outMarketIndex).price;
+
+		const outSaferThanIn =
+			inMarket.initialAssetWeight < outMarket.initialAssetWeight;
+
+		const inSpotPosition =
+			this.getSpotPosition(inMarketIndex) ||
+			this.getEmptySpotPosition(inMarketIndex);
+		const outSpotPosition =
+			this.getSpotPosition(outMarketIndex) ||
+			this.getEmptySpotPosition(outMarketIndex);
+
+		const freeCollateral = this.getFreeCollateral();
+
+		const inTokenAmount = this.getSpotTokenAmount(inMarketIndex);
+		if (freeCollateral.lt(ONE)) {
+			if (outSaferThanIn) {
+				return inTokenAmount;
+			} else {
+				return ZERO;
+			}
+		}
+
+		const inContributionInitial =
+			this.calculateSpotPositionFreeCollateralContribution(inSpotPosition);
+		const outContributionInitial =
+			this.calculateSpotPositionFreeCollateralContribution(outSpotPosition);
+		const initialContribution = inContributionInitial.add(
+			outContributionInitial
+		);
+
+		const inPrecision = new BN(10 ** inMarket.decimals);
+		const outPrecision = new BN(10 ** outMarket.decimals);
+
+		const cloneAndUpdateSpotPosition = (
+			position: SpotPosition,
+			tokenAmount: BN,
+			market: SpotMarketAccount
+		) => {
+			const clonedPosition = Object.assign({}, position);
+			const preTokenAmount = getTokenAmount(
+				position.scaledBalance,
+				market,
+				position.balanceType
+			);
+
+			if (sigNum(preTokenAmount).eq(sigNum(tokenAmount))) {
+				const scaledBalanceDelta = getBalance(
+					tokenAmount,
+					market,
+					position.balanceType
+				);
+				clonedPosition.scaledBalance =
+					position.scaledBalance.add(scaledBalanceDelta);
+				return clonedPosition;
+			}
+
+			const updateDirection = isVariant(position.balanceType, 'deposit')
+				? SpotBalanceType.BORROW
+				: SpotBalanceType.DEPOSIT;
+
+			if (tokenAmount.abs().gt(preTokenAmount.abs())) {
+				clonedPosition.scaledBalance = getBalance(
+					tokenAmount.abs().sub(preTokenAmount.abs()),
+					market,
+					updateDirection
+				);
+				clonedPosition.balanceType = updateDirection;
+			} else {
+				const scaledBalanceDelta = getBalance(
+					tokenAmount,
+					market,
+					position.balanceType
+				);
+
+				clonedPosition.scaledBalance =
+					position.scaledBalance.sub(scaledBalanceDelta);
+			}
+			return clonedPosition;
+		};
+
+		let minSwap = ZERO;
+		let maxSwap = freeCollateral
+			.mul(inPrecision)
+			.mul(SPOT_MARKET_WEIGHT_PRECISION)
+			.div(SPOT_MARKET_WEIGHT_PRECISION.div(TEN))
+			.div(inOraclePrice); // just assume user can go 10x
+		let swap = maxSwap.div(TWO);
+		const error = QUOTE_PRECISION;
+
+		let freeCollateralAfter = freeCollateral;
+		while (freeCollateralAfter.abs().gt(error) || freeCollateralAfter.isNeg()) {
+			const inSwap = swap;
+			const outSwap = inSwap
+				.mul(outPrecision)
+				.mul(inOraclePrice)
+				.div(outOraclePrice)
+				.div(inPrecision);
+
+			const inPositionAfter = cloneAndUpdateSpotPosition(
+				inSpotPosition,
+				inSwap.neg(),
+				inMarket
+			);
+			const outPositionAfter = cloneAndUpdateSpotPosition(
+				outSpotPosition,
+				outSwap,
+				outMarket
+			);
+
+			const inContributionAfter =
+				this.calculateSpotPositionFreeCollateralContribution(inPositionAfter);
+			const outContributionAfter =
+				this.calculateSpotPositionFreeCollateralContribution(outPositionAfter);
+			const contributionAfter = inContributionAfter.add(outContributionAfter);
+
+			const contributionDelta = contributionAfter.sub(initialContribution);
+
+			freeCollateralAfter = freeCollateral.add(contributionDelta);
+
+			if (freeCollateralAfter.gt(error)) {
+				minSwap = swap;
+				swap = minSwap.add(maxSwap).div(TWO);
+			} else if (freeCollateralAfter.abs().gt(error)) {
+				maxSwap = swap;
+				swap = minSwap.add(maxSwap).div(TWO);
+			}
+		}
+
+		return swap;
+	}
+
+	calculateSpotPositionFreeCollateralContribution(
+		spotPosition: SpotPosition
+	): BN {
+		let freeCollateralContribution = ZERO;
+		const now = new BN(new Date().getTime() / 1000);
+		const strict = true;
+		const marginCategory = 'Initial';
+
+		const spotMarketAccount: SpotMarketAccount =
+			this.driftClient.getSpotMarketAccount(spotPosition.marketIndex);
+
+		const oraclePriceData = this.getOracleDataForSpotMarket(
+			spotPosition.marketIndex
+		);
+
+		const [worstCaseTokenAmount, worstCaseQuoteTokenAmount] =
+			getWorstCaseTokenAmounts(
+				spotPosition,
+				spotMarketAccount,
+				oraclePriceData
+			);
+
+		if (worstCaseTokenAmount.gt(ZERO)) {
+			const baseAssetValue = this.getSpotAssetValue(
+				worstCaseTokenAmount,
+				oraclePriceData,
+				spotMarketAccount,
+				marginCategory,
+				strict,
+				now
+			);
+
+			freeCollateralContribution =
+				freeCollateralContribution.add(baseAssetValue);
+		} else {
+			const baseLiabilityValue = this.getSpotLiabilityValue(
+				worstCaseTokenAmount,
+				oraclePriceData,
+				spotMarketAccount,
+				'Initial',
+				undefined,
+				strict,
+				now
+			).abs();
+
+			freeCollateralContribution =
+				freeCollateralContribution.sub(baseLiabilityValue);
+		}
+
+		freeCollateralContribution.add(worstCaseQuoteTokenAmount);
+
+		return freeCollateralContribution;
+	}
+
 	// TODO - should this take the price impact of the trade into account for strict accuracy?
 
 	/**