From 72e208260285f2f78f60fc9c4ae8044c58beee20 Mon Sep 17 00:00:00 2001
From: Vincent Chau <99756290+vincentwschau@users.noreply.github.com>
Date: Thu, 21 Mar 2024 18:56:59 -0400
Subject: [PATCH] Address nits.

---
 .../subaccounts/keeper/isolated_subaccount.go | 39 ++++++++++++++++---
 protocol/x/subaccounts/keeper/subaccount.go   | 14 ++-----
 2 files changed, 38 insertions(+), 15 deletions(-)

diff --git a/protocol/x/subaccounts/keeper/isolated_subaccount.go b/protocol/x/subaccounts/keeper/isolated_subaccount.go
index 26899a1dc9..d302817788 100644
--- a/protocol/x/subaccounts/keeper/isolated_subaccount.go
+++ b/protocol/x/subaccounts/keeper/isolated_subaccount.go
@@ -248,6 +248,11 @@ func (k *Keeper) transferCollateralForIsolatedPerpetual(
 		return nil
 	}
 
+	// If there are zero quantums to transfer, don't transfer collateral.
+	if stateTransition.QuoteQuantums.Sign() == 0 {
+		return nil
+	}
+
 	isolatedCollateralPoolAddr, err := k.GetCollateralPoolFromPerpetualId(ctx, stateTransition.PerpetualId)
 	if err != nil {
 		return err
@@ -276,11 +281,6 @@ func (k *Keeper) transferCollateralForIsolatedPerpetual(
 		)
 	}
 
-	// If there are zero quantums to transfer, don't transfer collateral.
-	if stateTransition.QuoteQuantums.Sign() == 0 {
-		return nil
-	}
-
 	// Invalid to transfer negative quantums. This should already be caught by collateralization
 	// checks as well.
 	if stateTransition.QuoteQuantums.Sign() == -1 {
@@ -317,6 +317,35 @@ func (k *Keeper) transferCollateralForIsolatedPerpetual(
 	return nil
 }
 
+// computeAndExecuteCollateralTransfer computes collateral transfers resulting from updates to
+// a subaccount and executes the collateral transfer using `x/bank`.`
+// The input `settledUpdate` must have an updated subaccount (`settledUpdate.SettledSubaccount`),
+// so all the updates must have been applied already to the subaccount.
+// Note: This uses the `x/bank` keeper and modifies `x/bank` state.
+func (k *Keeper) computeAndExecuteCollateralTransfer(
+	ctx sdk.Context,
+	settledUpdateWithUpdatedSubaccount SettledUpdate,
+	perpetuals []perptypes.Perpetual,
+) error {
+	// The subaccount in `settledUpdateWithUpdatedSubaccount` already has the perpetual updates
+	// and asset updates applied to it.
+	stateTransition, err := GetIsolatedPerpetualStateTransition(
+		settledUpdateWithUpdatedSubaccount,
+		perpetuals,
+	)
+	if err != nil {
+		return err
+	}
+	if err := k.transferCollateralForIsolatedPerpetual(
+		ctx,
+		stateTransition,
+	); err != nil {
+		return err
+	}
+
+	return nil
+}
+
 func getPerpIdToMarketTypeMap(
 	perpetuals []perptypes.Perpetual,
 ) map[uint32]perptypes.PerpetualMarketType {
diff --git a/protocol/x/subaccounts/keeper/subaccount.go b/protocol/x/subaccounts/keeper/subaccount.go
index 4ff5439290..55a1b5fb14 100644
--- a/protocol/x/subaccounts/keeper/subaccount.go
+++ b/protocol/x/subaccounts/keeper/subaccount.go
@@ -313,18 +313,12 @@ func (k Keeper) UpdateSubaccounts(
 	// Transfer collateral between collateral pools for any isolated perpetual positions that changed
 	// state due to an update.
 	for _, settledUpdateWithUpdatedSubaccount := range settledUpdates {
-		// The subaccount in `settledUpdateWithUpdatedSubaccount` already has the perpetual updates
-		// and asset updates applied to it.
-		stateTransition, err := GetIsolatedPerpetualStateTransition(
+		if err := k.computeAndExecuteCollateralTransfer(
+			ctx,
+			// The subaccount in `settledUpdateWithUpdatedSubaccount` already has the perpetual updates
+			// and asset updates applied to it.
 			settledUpdateWithUpdatedSubaccount,
 			allPerps,
-		)
-		if err != nil {
-			return false, nil, err
-		}
-		if err := k.transferCollateralForIsolatedPerpetual(
-			ctx,
-			stateTransition,
 		); err != nil {
 			return false, nil, err
 		}