[TRA-78] Add function to retrieve collateral pool addr for a subaccou…

…nt (dydxprotocol#1142)

Signed-off-by: Shrenuj Bansal <shrenuj@dydx.exchange>
Signed-off-by: Eric <eric.warehime@gmail.com>

protocol/x/subaccounts/keeper/subaccount.go

@@ -7,6 +7,8 @@ import (
 	"math/rand"
 	"time"
 
+	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
+
 	"github.com/cosmos/gogoproto/proto"
 
 	storetypes "cosmossdk.io/store/types"
@@ -52,6 +54,49 @@ func (k Keeper) SetSubaccount(ctx sdk.Context, subaccount types.Subaccount) {
 	}
 }
 
+// GetCollateralPoolForSubaccount returns the collateral pool address for a subaccount
+// based on the subaccount's perpetual positions. If the subaccount holds a position in an isolated
+// market, the collateral pool address will be the isolated market's pool address. Otherwise, the
+// collateral pool address will be the module's pool address.
+func (k Keeper) GetCollateralPoolForSubaccount(ctx sdk.Context, subaccountId types.SubaccountId) (
+	sdk.AccAddress,
+	error,
+) {
+	poolName, err := k.GetCollateralPoolNameForSubaccount(ctx, subaccountId)
+	if err != nil {
+		return nil, err
+	}
+	return authtypes.NewModuleAddress(poolName), nil
+}
+
+func (k Keeper) GetCollateralPoolNameForSubaccount(ctx sdk.Context, subaccountId types.SubaccountId) (string, error) {
+	subaccount := k.GetSubaccount(ctx, subaccountId)
+	if len(subaccount.PerpetualPositions) == 0 {
+		return types.ModuleName, nil
+	}
+
+	// Get the first perpetual position and return the collateral pool name.
+	perpetual, err := k.perpetualsKeeper.GetPerpetual(ctx, subaccount.PerpetualPositions[0].PerpetualId)
+	if err != nil {
+		panic(fmt.Sprintf("GetCollateralPoolNameForSubaccount: %v", err))
+	}
+
+	if perpetual.Params.MarketType == perptypes.PerpetualMarketType_PERPETUAL_MARKET_TYPE_ISOLATED {
+		return types.ModuleName + ":" + lib.UintToString(perpetual.GetId()), nil
+	}
+
+	return types.ModuleName, nil
+}
+
+// IsIsolatedMarketSubaccount returns whether a subaccount is isolated to a specific market.
+func (k Keeper) IsIsolatedMarketSubaccount(ctx sdk.Context, subaccountId types.SubaccountId) (bool, error) {
+	poolName, err := k.GetCollateralPoolNameForSubaccount(ctx, subaccountId)
+	if err != nil {
+		panic(fmt.Sprintf("IsIsolatedMarketSubaccount: %v", err))
+	}
+	return poolName != types.ModuleName, nil
+}
+
 // GetSubaccount returns a subaccount from its index.
 //
 // Note that this function is getting called very frequently; metrics in this function

protocol/x/subaccounts/keeper/subaccount_test.go

@@ -6,6 +6,9 @@ import (
 	"strconv"
 	"testing"
 
+	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
+	"github.com/dydxprotocol/v4-chain/protocol/lib"
+
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/dydxprotocol/v4-chain/protocol/dtypes"
 	indexerevents "github.com/dydxprotocol/v4-chain/protocol/indexer/events"
@@ -99,6 +102,79 @@ func assertSubaccountUpdateEventsInIndexerBlock(
 	}
 }
 
+func TestGetCollateralPool(t *testing.T) {
+	tests := map[string]struct {
+		// state
+		perpetuals         []perptypes.Perpetual
+		perpetualPositions []*types.PerpetualPosition
+
+		expectedAddress sdk.AccAddress
+	}{
+		"collateral pool with cross margin markets": {
+			perpetuals: []perptypes.Perpetual{
+				constants.BtcUsd_SmallMarginRequirement,
+			},
+			perpetualPositions: []*types.PerpetualPosition{
+				&constants.PerpetualPosition_OneBTCLong,
+			},
+			expectedAddress: authtypes.NewModuleAddress(types.ModuleName),
+		},
+		"collateral pool with isolated margin markets": {
+			perpetuals: []perptypes.Perpetual{
+				constants.IsoUsd_IsolatedMarket,
+			},
+			perpetualPositions: []*types.PerpetualPosition{
+				{
+					PerpetualId: constants.IsoUsd_IsolatedMarket.GetId(),
+					Quantums:    dtypes.NewInt(100_000_000),
+				},
+			},
+			expectedAddress: authtypes.NewModuleAddress(
+				types.ModuleName + ":" + lib.UintToString(constants.IsoUsd_IsolatedMarket.GetId()),
+			),
+		},
+		"collateral pool with no positions": {
+			perpetualPositions: make([]*types.PerpetualPosition, 0),
+			expectedAddress:    authtypes.NewModuleAddress(types.ModuleName),
+		},
+	}
+	for name, tc := range tests {
+		t.Run(
+			name, func(t *testing.T) {
+				ctx, keeper, pricesKeeper, perpetualsKeeper, _, _, assetsKeeper, _, _ := testutil.SubaccountsKeepers(
+					t,
+					true,
+				)
+
+				testutil.CreateTestMarkets(t, ctx, pricesKeeper)
+				testutil.CreateTestLiquidityTiers(t, ctx, perpetualsKeeper)
+
+				require.NoError(t, testutil.CreateUsdcAsset(ctx, assetsKeeper))
+				for _, p := range tc.perpetuals {
+					_, err := perpetualsKeeper.CreatePerpetual(
+						ctx,
+						p.Params.Id,
+						p.Params.Ticker,
+						p.Params.MarketId,
+						p.Params.AtomicResolution,
+						p.Params.DefaultFundingPpm,
+						p.Params.LiquidityTier,
+						p.Params.MarketType,
+					)
+					require.NoError(t, err)
+				}
+
+				subaccount := createNSubaccount(keeper, ctx, 1, big.NewInt(1_000))[0]
+				subaccount.PerpetualPositions = tc.perpetualPositions
+				keeper.SetSubaccount(ctx, subaccount)
+				collateralPoolAddr, err := keeper.GetCollateralPoolForSubaccount(ctx, *subaccount.Id)
+				require.NoError(t, err)
+				require.Equal(t, tc.expectedAddress, collateralPoolAddr)
+			},
+		)
+	}
+}
+
 func TestSubaccountGet(t *testing.T) {
 	ctx, keeper, _, _, _, _, _, _, _ := testutil.SubaccountsKeepers(t, true)
 	items := createNSubaccount(keeper, ctx, 10, big.NewInt(1_000))