migrate vault shares to megavault shares in v7 upgrade handler (#2379)

protocol/app/upgrades/v7.0.0/upgrade.go

@@ -3,6 +3,7 @@ package v_7_0_0
 import (
 	"context"
 	"fmt"
+	"math/big"
 
 	upgradetypes "cosmossdk.io/x/upgrade/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
@@ -14,6 +15,11 @@ import (
 	vaulttypes "github.com/dydxprotocol/v4-chain/protocol/x/vault/types"
 )
 
+const (
+	// Each megavault share is worth 1 USDC.
+	QUOTE_QUANTUMS_PER_MEGAVAULT_SHARE = 1_000_000
+)
+
 func initCurrencyPairIDCache(ctx sdk.Context, k pricestypes.PricesKeeper) {
 	marketParams := k.GetAllMarketParams(ctx)
 	for _, mp := range marketParams {
@@ -55,6 +61,68 @@ func migrateVaultQuotingParamsToVaultParams(ctx sdk.Context, k vaultkeeper.Keepe
 	}
 }
 
+// In 6.x,
+// Total shares store (key prefix `TotalShares:`) is `vaultId -> shares`
+// Owner shares store (key prefix `OwnerShares:`) is `vaultId -> owner -> shares`
+// In 7.x,
+// Total shares store is just `"TotalShares" -> shares`
+// Owner shares store (key prefix `OwnerShares:`) is `owner -> shares`
+// Thus, this function
+// 1. Calculate how much equity each owner owns
+// 2. Delete all keys in deprecated total shares and owner shares stores
+// 3. Grant each owner 1 megavault share per usdc of equity owned
+// 4. Set total megavault shares to sum of all owner shares granted
+func migrateVaultSharesToMegavaultShares(ctx sdk.Context, k vaultkeeper.Keeper) {
+	ctx.Logger().Info("Migrating vault shares to megavault shares")
+	quoteQuantumsPerShare := big.NewInt(QUOTE_QUANTUMS_PER_MEGAVAULT_SHARE)
+
+	ownerEquities := k.UnsafeGetAllOwnerEquities(ctx)
+	ctx.Logger().Info(fmt.Sprintf("Calculated owner equities %s", ownerEquities))
+	k.UnsafeDeleteAllVaultTotalShares(ctx)
+	ctx.Logger().Info("Deleted all keys in deprecated vault total shares store")
+	k.UnsafeDeleteAllVaultOwnerShares(ctx)
+	ctx.Logger().Info("Deleted all keys in deprecated vault owner shares store")
+
+	totalShares := big.NewInt(0)
+	for owner, equity := range ownerEquities {
+		ownerShares := new(big.Int).Quo(
+			equity.Num(),
+			equity.Denom(),
+		)
+		ownerShares.Quo(ownerShares, quoteQuantumsPerShare)
+
+		if ownerShares.Sign() <= 0 {
+			ctx.Logger().Warn(fmt.Sprintf(
+				"Owner %s has non-positive shares %s from %s quote quantums",
+				owner,
+				ownerShares,
+				equity,
+			))
+			continue
+		}
+
+		err := k.SetOwnerShares(ctx, owner, vaulttypes.BigIntToNumShares(ownerShares))
+		if err != nil {
+			panic(err)
+		}
+		ctx.Logger().Info(fmt.Sprintf(
+			"Set megavault owner shares of %s: shares=%s, equity=%s",
+			owner,
+			ownerShares,
+			equity,
+		))
+
+		totalShares.Add(totalShares, ownerShares)
+	}
+
+	err := k.SetTotalShares(ctx, vaulttypes.BigIntToNumShares(totalShares))
+	if err != nil {
+		panic(err)
+	}
+	ctx.Logger().Info(fmt.Sprintf("Set megavault total shares to: %s", totalShares))
+	ctx.Logger().Info("Successfully migrated vault shares to megavault shares")
+}
+
 func CreateUpgradeHandler(
 	mm *module.Manager,
 	configurator module.Configurator,
@@ -71,6 +139,9 @@ func CreateUpgradeHandler(
 		// Migrate vault quoting params to vault params.
 		migrateVaultQuotingParamsToVaultParams(sdkCtx, vaultKeeper)
 
+		// Migrate vault shares to megavault shares.
+		migrateVaultSharesToMegavaultShares(sdkCtx, vaultKeeper)
+
 		return mm.RunMigrations(ctx, configurator, vm)
 	}
 }

protocol/app/upgrades/v7.0.0/upgrade_container_test.go

@@ -3,6 +3,7 @@
 package v_7_0_0_test
 
 import (
+	"math/big"
 	"testing"
 
 	"github.com/cosmos/gogoproto/proto"
@@ -48,6 +49,7 @@ func preUpgradeChecks(node *containertest.Node, t *testing.T) {
 func postUpgradeChecks(node *containertest.Node, t *testing.T) {
 	// Add test for your upgrade handler logic below
 	postUpgradeVaultParamsCheck(node, t)
+	postUpgradeMegavaultSharesCheck(node, t)
 }
 
 func postUpgradeVaultParamsCheck(node *containertest.Node, t *testing.T) {
@@ -96,3 +98,67 @@ func checkVaultParams(
 	require.Equal(t, expectedStatus, vaultParamsResp.VaultParams.Status)
 	require.Equal(t, expectedQuotingParams, vaultParamsResp.VaultParams.QuotingParams)
 }
+
+func postUpgradeMegavaultSharesCheck(node *containertest.Node, t *testing.T) {
+	// Alice equity = vault_0_equity * 1 + vault_1_equity * 1/3 + vault_2_equity * 123_456/556_677
+	// = 1_000 + 2_000 * 1/3 + 3_000 * 123_456/556_677
+	// ~= 2331.99
+	// Bob equity = vault_1_equity * 1/3 + vault_2_equity * 433_221/556_677
+	// = 2_000 * 1/3 + 3_000 * 433_221/556_677
+	// ~= 3001.35
+	// Carl equity = vault_1_equity * 1/3
+	// = 2_000 * 1/3
+	// ~= 666.67
+	// 1 USDC in equity should be granted 1 megavault share and round down to nearest integer.
+	expectedOwnerShares := map[string]*big.Int{
+		constants.AliceAccAddress.String(): big.NewInt(2_331),
+		constants.BobAccAddress.String():   big.NewInt(3_001),
+		constants.CarlAccAddress.String():  big.NewInt(666),
+	}
+	// 2331 + 3001 + 666 = 5998
+	expectedTotalShares := big.NewInt(5_998)
+
+	// Check MegaVault total shares.
+	resp, err := containertest.Query(
+		node,
+		vaulttypes.NewQueryClient,
+		vaulttypes.QueryClient.MegavaultTotalShares,
+		&vaulttypes.QueryMegavaultTotalSharesRequest{},
+	)
+	require.NoError(t, err)
+	require.NotNil(t, resp)
+
+	totalSharesResp := vaulttypes.QueryMegavaultTotalSharesResponse{}
+	err = proto.UnmarshalText(resp.String(), &totalSharesResp)
+	require.NoError(t, err)
+
+	require.Equal(
+		t,
+		expectedTotalShares,
+		totalSharesResp.TotalShares.NumShares.BigInt(),
+	)
+
+	// Check MegaVault owner shares.
+	resp, err = containertest.Query(
+		node,
+		vaulttypes.NewQueryClient,
+		vaulttypes.QueryClient.MegavaultAllOwnerShares,
+		&vaulttypes.QueryMegavaultAllOwnerSharesRequest{},
+	)
+	require.NoError(t, err)
+	require.NotNil(t, resp)
+
+	allOwnerSharesResp := vaulttypes.QueryMegavaultAllOwnerSharesResponse{}
+	err = proto.UnmarshalText(resp.String(), &allOwnerSharesResp)
+	require.NoError(t, err)
+
+	require.Len(t, allOwnerSharesResp.OwnerShares, 3)
+	gotOwnerShares := make(map[string]*big.Int)
+	for _, ownerShare := range allOwnerSharesResp.OwnerShares {
+		gotOwnerShares[ownerShare.Owner] = ownerShare.Shares.NumShares.BigInt()
+	}
+	for owner, expectedShares := range expectedOwnerShares {
+		require.Contains(t, gotOwnerShares, owner)
+		require.Equal(t, expectedShares, gotOwnerShares[owner])
+	}
+}

protocol/testing/containertest/preupgrade_genesis.json

@@ -4354,7 +4354,21 @@
           "asset_positions": [
             {
               "asset_id": 0,
-              "quantums": "1000000000",
+              "quantums": "2000000000",
+              "index": 0
+            }
+          ]
+        },
+        {
+          "id": {
+            "owner": "dydx190te44zcctdgk0qmqtenve2m00g3r2dn7ntd72",
+            "number": 0
+          },
+          "margin_enabled": true,
+          "asset_positions": [
+            {
+              "asset_id": 0,
+              "quantums": "3000000000",
               "index": 0
             }
           ]
@@ -4405,14 +4419,49 @@
             "number": 1
           },
           "total_shares": {
-            "num_shares": "1000000000"
+            "num_shares": "3000000000"
           },
           "owner_shares": [
             {
               "owner": "dydx199tqg4wdlnu4qjlxchpd7seg454937hjrknju4",
               "shares": {
                 "num_shares": "1000000000"
               }
+            },
+            {
+              "owner": "dydx10fx7sy6ywd5senxae9dwytf8jxek3t2gcen2vs",
+              "shares": {
+                "num_shares": "1000000000"
+              }
+            },
+            {
+              "owner": "dydx1fjg6zp6vv8t9wvy4lps03r5l4g7tkjw9wvmh70",
+              "shares": {
+                "num_shares": "1000000000"
+              }
+            }
+          ]
+        },
+        {
+          "vault_id": {
+            "type": "VAULT_TYPE_CLOB",
+            "number": 2
+          },
+          "total_shares": {
+            "num_shares": "556677"
+          },
+          "owner_shares": [
+            {
+              "owner": "dydx199tqg4wdlnu4qjlxchpd7seg454937hjrknju4",
+              "shares": {
+                "num_shares": "123456"
+              }
+            },
+            {
+              "owner": "dydx10fx7sy6ywd5senxae9dwytf8jxek3t2gcen2vs",
+              "shares": {
+                "num_shares": "433221"
+              }
             }
           ]
         }

protocol/x/vault/keeper/deprecated_state.go

@@ -1,6 +1,8 @@
 package keeper
 
 import (
+	"math/big"
+
 	"cosmossdk.io/store/prefix"
 	storetypes "cosmossdk.io/store/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
@@ -20,6 +22,11 @@ const (
 	// Deprecated: For use by the v7.x upgrade handler
 	// TotalSharesKeyPrefix is the prefix to retrieve all TotalShares.
 	TotalSharesKeyPrefix = "TotalShares:"
+
+	// Deprecated: For use by the v7.x upgrade handler
+	// OwnerSharesKeyPrefix is the prefix to retrieve all OwnerShares.
+	// OwnerShares store: vaultId VaultId -> owner string -> shares NumShares.
+	OwnerSharesKeyPrefix = "OwnerShares:"
 )
 
 // v5.x state, used for v6.x upgrade.
@@ -116,3 +123,112 @@ func (k Keeper) UnsafeGetAllVaultIds(ctx sdk.Context) []types.VaultId {
 	}
 	return vaultIds
 }
+
+// GetTotalShares gets TotalShares for a vault.
+// Deprecated and used for v7.x upgrade handler
+func (k Keeper) UnsafeGetTotalShares(
+	ctx sdk.Context,
+	vaultId types.VaultId,
+) (val types.NumShares, exists bool) {
+	store := prefix.NewStore(ctx.KVStore(k.storeKey), []byte(TotalSharesKeyPrefix))
+
+	b := store.Get(vaultId.ToStateKey())
+	if b == nil {
+		return val, false
+	}
+
+	k.cdc.MustUnmarshal(b, &val)
+	return val, true
+}
+
+// UnsafeGetAllOwnerShares gets all owner shares of a given vault.
+// Deprecated and used for v7.x upgrade handler
+func (k Keeper) UnsafeGetAllOwnerShares(
+	ctx sdk.Context,
+	vaultId types.VaultId,
+) []*types.OwnerShare {
+	allOwnerShares := []*types.OwnerShare{}
+
+	store := prefix.NewStore(ctx.KVStore(k.storeKey), []byte(types.OwnerSharesKeyPrefix))
+	ownerSharesStore := prefix.NewStore(store, vaultId.ToStateKeyPrefix())
+	ownerSharesIterator := storetypes.KVStorePrefixIterator(ownerSharesStore, []byte{})
+	defer ownerSharesIterator.Close()
+	for ; ownerSharesIterator.Valid(); ownerSharesIterator.Next() {
+		owner := string(ownerSharesIterator.Key())
+		var ownerShares types.NumShares
+		k.cdc.MustUnmarshal(ownerSharesIterator.Value(), &ownerShares)
+		allOwnerShares = append(allOwnerShares, &types.OwnerShare{
+			Owner:  owner,
+			Shares: ownerShares,
+		})
+	}
+	return allOwnerShares
+}
+
+// UnsafeGetAllOwnerEquities returns equity that belongs to each owner across all vaults
+// using the deprecated owner shares and total shares state.
+// Deprecated and used for v7.x upgrade handler
+func (k Keeper) UnsafeGetAllOwnerEquities(ctx sdk.Context) map[string]*big.Rat {
+	ownerEquities := make(map[string]*big.Rat)
+	totalSharesStore := prefix.NewStore(ctx.KVStore(k.storeKey), []byte(TotalSharesKeyPrefix))
+	totalSharesIterator := storetypes.KVStorePrefixIterator(totalSharesStore, []byte{})
+	defer totalSharesIterator.Close()
+	for ; totalSharesIterator.Valid(); totalSharesIterator.Next() {
+		vaultId, err := types.GetVaultIdFromStateKey(totalSharesIterator.Key())
+		if err != nil {
+			panic(err)
+		}
+
+		var vaultTotalShares types.NumShares
+		k.cdc.MustUnmarshal(totalSharesIterator.Value(), &vaultTotalShares)
+		bigVaultTotalShares := vaultTotalShares.NumShares.BigInt()
+		vaultEquity, err := k.GetVaultEquity(ctx, *vaultId)
+		if err != nil {
+			panic(err)
+		}
+
+		ownerShares := k.UnsafeGetAllOwnerShares(ctx, *vaultId)
+		for _, ownerShare := range ownerShares {
+			// owner equity in this vault = vault equity * owner shares / vault total shares
+			ownerEquity := new(big.Rat).SetInt(vaultEquity)
+			ownerEquity.Mul(
+				ownerEquity,
+				new(big.Rat).SetInt(ownerShare.Shares.NumShares.BigInt()),
+			)
+			ownerEquity.Quo(
+				ownerEquity,
+				new(big.Rat).SetInt(bigVaultTotalShares),
+			)
+
+			if e, ok := ownerEquities[ownerShare.Owner]; ok {
+				ownerEquities[ownerShare.Owner] = e.Add(e, ownerEquity)
+			} else {
+				ownerEquities[ownerShare.Owner] = ownerEquity
+			}
+		}
+	}
+
+	return ownerEquities
+}
+
+// UnsafeDeleteVaultTotalShares deletes total shares of a given vault from state.
+// Used for v7.x upgrade handler
+func (k Keeper) UnsafeDeleteAllVaultTotalShares(ctx sdk.Context) {
+	store := prefix.NewStore(ctx.KVStore(k.storeKey), []byte(TotalSharesKeyPrefix))
+	iterator := storetypes.KVStorePrefixIterator(store, []byte{})
+	defer iterator.Close()
+	for ; iterator.Valid(); iterator.Next() {
+		store.Delete(iterator.Key())
+	}
+}
+
+// UnsafeDeleteVaultOwnerShares deletes all owner shares of a given vault from state.
+// Used for v7.x upgrade handler
+func (k Keeper) UnsafeDeleteAllVaultOwnerShares(ctx sdk.Context) {
+	store := prefix.NewStore(ctx.KVStore(k.storeKey), []byte(OwnerSharesKeyPrefix))
+	iterator := storetypes.KVStorePrefixIterator(store, []byte{})
+	defer iterator.Close()
+	for ; iterator.Valid(); iterator.Next() {
+		store.Delete(iterator.Key())
+	}
+}