[TRA-446] include vault shares in genesis state (#1774)

indexer/packages/v4-protos/src/codegen/dydxprotocol/vault/genesis.ts

@@ -1,22 +1,55 @@
 import { Params, ParamsSDKType } from "./params";
+import { VaultId, VaultIdSDKType, NumShares, NumSharesSDKType } from "./vault";
+import { OwnerShare, OwnerShareSDKType } from "./query";
 import * as _m0 from "protobufjs/minimal";
 import { DeepPartial } from "../../helpers";
 /** GenesisState defines `x/vault`'s genesis state. */
 
 export interface GenesisState {
   /** The parameters of the module. */
   params?: Params;
+  /** The vaults. */
+
+  vaults: Vault[];
 }
 /** GenesisState defines `x/vault`'s genesis state. */
 
 export interface GenesisStateSDKType {
   /** The parameters of the module. */
   params?: ParamsSDKType;
+  /** The vaults. */
+
+  vaults: VaultSDKType[];
+}
+/** Vault defines the total shares and owner shares of a vault. */
+
+export interface Vault {
+  /** The ID of the vault. */
+  vaultId?: VaultId;
+  /** The total number of shares in the vault. */
+
+  totalShares?: NumShares;
+  /** The shares of each owner in the vault. */
+
+  ownerShares: OwnerShare[];
+}
+/** Vault defines the total shares and owner shares of a vault. */
+
+export interface VaultSDKType {
+  /** The ID of the vault. */
+  vault_id?: VaultIdSDKType;
+  /** The total number of shares in the vault. */
+
+  total_shares?: NumSharesSDKType;
+  /** The shares of each owner in the vault. */
+
+  owner_shares: OwnerShareSDKType[];
 }
 
 function createBaseGenesisState(): GenesisState {
   return {
-    params: undefined
+    params: undefined,
+    vaults: []
   };
 }
 
@@ -26,6 +59,10 @@ export const GenesisState = {
       Params.encode(message.params, writer.uint32(10).fork()).ldelim();
     }
 
+    for (const v of message.vaults) {
+      Vault.encode(v!, writer.uint32(18).fork()).ldelim();
+    }
+
     return writer;
   },
 
@@ -42,6 +79,10 @@ export const GenesisState = {
           message.params = Params.decode(reader, reader.uint32());
           break;
 
+        case 2:
+          message.vaults.push(Vault.decode(reader, reader.uint32()));
+          break;
+
         default:
           reader.skipType(tag & 7);
           break;
@@ -54,6 +95,72 @@ export const GenesisState = {
   fromPartial(object: DeepPartial<GenesisState>): GenesisState {
     const message = createBaseGenesisState();
     message.params = object.params !== undefined && object.params !== null ? Params.fromPartial(object.params) : undefined;
+    message.vaults = object.vaults?.map(e => Vault.fromPartial(e)) || [];
+    return message;
+  }
+
+};
+
+function createBaseVault(): Vault {
+  return {
+    vaultId: undefined,
+    totalShares: undefined,
+    ownerShares: []
+  };
+}
+
+export const Vault = {
+  encode(message: Vault, writer: _m0.Writer = _m0.Writer.create()): _m0.Writer {
+    if (message.vaultId !== undefined) {
+      VaultId.encode(message.vaultId, writer.uint32(10).fork()).ldelim();
+    }
+
+    if (message.totalShares !== undefined) {
+      NumShares.encode(message.totalShares, writer.uint32(18).fork()).ldelim();
+    }
+
+    for (const v of message.ownerShares) {
+      OwnerShare.encode(v!, writer.uint32(26).fork()).ldelim();
+    }
+
+    return writer;
+  },
+
+  decode(input: _m0.Reader | Uint8Array, length?: number): Vault {
+    const reader = input instanceof _m0.Reader ? input : new _m0.Reader(input);
+    let end = length === undefined ? reader.len : reader.pos + length;
+    const message = createBaseVault();
+
+    while (reader.pos < end) {
+      const tag = reader.uint32();
+
+      switch (tag >>> 3) {
+        case 1:
+          message.vaultId = VaultId.decode(reader, reader.uint32());
+          break;
+
+        case 2:
+          message.totalShares = NumShares.decode(reader, reader.uint32());
+          break;
+
+        case 3:
+          message.ownerShares.push(OwnerShare.decode(reader, reader.uint32()));
+          break;
+
+        default:
+          reader.skipType(tag & 7);
+          break;
+      }
+    }
+
+    return message;
+  },
+
+  fromPartial(object: DeepPartial<Vault>): Vault {
+    const message = createBaseVault();
+    message.vaultId = object.vaultId !== undefined && object.vaultId !== null ? VaultId.fromPartial(object.vaultId) : undefined;
+    message.totalShares = object.totalShares !== undefined && object.totalShares !== null ? NumShares.fromPartial(object.totalShares) : undefined;
+    message.ownerShares = object.ownerShares?.map(e => OwnerShare.fromPartial(e)) || [];
     return message;
   }
 

proto/dydxprotocol/vault/genesis.proto

@@ -3,11 +3,25 @@ package dydxprotocol.vault;
 
 import "gogoproto/gogo.proto";
 import "dydxprotocol/vault/params.proto";
+import "dydxprotocol/vault/query.proto";
+import "dydxprotocol/vault/vault.proto";
 
 option go_package = "github.com/dydxprotocol/v4-chain/protocol/x/vault/types";
 
 // GenesisState defines `x/vault`'s genesis state.
 message GenesisState {
   // The parameters of the module.
   Params params = 1 [ (gogoproto.nullable) = false ];
+  // The vaults.
+  repeated Vault vaults = 2;
+}
+
+// Vault defines the total shares and owner shares of a vault.
+message Vault {
+  // The ID of the vault.
+  VaultId vault_id = 1;
+  // The total number of shares in the vault.
+  NumShares total_shares = 2;
+  // The shares of each owner in the vault.
+  repeated OwnerShare owner_shares = 3;
 }

protocol/app/testdata/default_genesis_state.json

@@ -464,7 +464,8 @@
       "order_size_pct_ppm": 100000,
       "order_expiration_seconds": 2,
       "activation_threshold_quote_quantums": "1000000000"
-    }
+    },
+    "vaults": []
   },
   "vest": {
     "vest_entries": [

protocol/scripts/genesis/sample_pregenesis.json

@@ -1815,7 +1815,8 @@
         "skew_factor_ppm": 2000000,
         "spread_buffer_ppm": 1500,
         "spread_min_ppm": 10000
-      }
+      },
+      "vaults": []
     },
     "vest": {
       "vest_entries": [

protocol/testutil/constants/genesis.go

@@ -1501,7 +1501,8 @@ const GenesisState = `{
         "order_size_pct_ppm": 100000,
         "order_expiration_seconds": 2,
         "activation_threshold_quote_quantums": "1000000000"
-      }
+      },
+      "vaults": []
     },
     "vest": {
       "vest_entries": [

protocol/x/vault/genesis.go

@@ -10,15 +10,32 @@ import (
 func InitGenesis(ctx sdk.Context, k keeper.Keeper, genState types.GenesisState) {
 	k.InitializeForGenesis(ctx)
 
+	// Set params.
 	if err := k.SetParams(ctx, genState.Params); err != nil {
 		panic(err)
 	}
+	// Set total shares and owner shares of each vault.
+	for _, vault := range genState.Vaults {
+		if err := k.SetTotalShares(ctx, *vault.VaultId, *vault.TotalShares); err != nil {
+			panic(err)
+		}
+		for _, ownerShares := range vault.OwnerShares {
+			if err := k.SetOwnerShares(ctx, *vault.VaultId, ownerShares.Owner, *ownerShares.Shares); err != nil {
+				panic(err)
+			}
+		}
+	}
 }
 
 // ExportGenesis returns the module's exported genesis.
 func ExportGenesis(ctx sdk.Context, k keeper.Keeper) *types.GenesisState {
 	genesis := types.DefaultGenesis()
+
+	// Export params.
 	genesis.Params = k.GetParams(ctx)
 
+	// Export vaults.
+	genesis.Vaults = k.GetAllVaults(ctx)
+
 	return genesis
 }

protocol/x/vault/keeper/shares.go

@@ -101,6 +101,27 @@ func (k Keeper) getVaultOwnerSharesStore(
 	return prefix.NewStore(store, vaultId.ToStateKeyPrefix())
 }
 
+// GetAllOwnerShares gets all owner shares of a given vault.
+func (k Keeper) GetAllOwnerShares(
+	ctx sdk.Context,
+	vaultId types.VaultId,
+) []*types.OwnerShare {
+	allOwnerShares := []*types.OwnerShare{}
+	ownerSharesStore := k.getVaultOwnerSharesStore(ctx, vaultId)
+	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
+}
+
 // MintShares mints shares of a vault for `owner` based on `quantumsToDeposit` by:
 // 1. Increasing total shares of the vault.
 // 2. Increasing owner shares of the vault for given `owner`.

protocol/x/vault/keeper/shares_test.go

@@ -132,6 +132,43 @@ func TestGetSetOwnerShares(t *testing.T) {
 	require.Equal(t, numShares, got)
 }
 
+func TestGetAllOwnerShares(t *testing.T) {
+	tApp := testapp.NewTestAppBuilder(t).Build()
+	ctx := tApp.InitChain()
+	k := tApp.App.VaultKeeper
+
+	// Get all owner shares of a vault that has no owners.
+	allOwnerShares := k.GetAllOwnerShares(ctx, constants.Vault_Clob_0)
+	require.Equal(t, []*vaulttypes.OwnerShare{}, allOwnerShares)
+
+	// Set alice and bob as owners of a vault and get all owner shares.
+	alice := constants.AliceAccAddress.String()
+	aliceShares := vaulttypes.BigIntToNumShares(big.NewInt(7))
+	bob := constants.BobAccAddress.String()
+	bobShares := vaulttypes.BigIntToNumShares(big.NewInt(123))
+
+	err := k.SetOwnerShares(ctx, constants.Vault_Clob_0, alice, aliceShares)
+	require.NoError(t, err)
+	err = k.SetOwnerShares(ctx, constants.Vault_Clob_0, bob, bobShares)
+	require.NoError(t, err)
+
+	allOwnerShares = k.GetAllOwnerShares(ctx, constants.Vault_Clob_0)
+	require.ElementsMatch(
+		t,
+		[]*vaulttypes.OwnerShare{
+			{
+				Owner:  alice,
+				Shares: &aliceShares,
+			},
+			{
+				Owner:  bob,
+				Shares: &bobShares,
+			},
+		},
+		allOwnerShares,
+	)
+}
+
 func TestMintShares(t *testing.T) {
 	tests := map[string]struct {
 		/* --- Setup --- */

protocol/x/vault/keeper/vault.go

@@ -99,3 +99,28 @@ func (k Keeper) DecommissionVault(
 		ownerSharesStore.Delete(ownerSharesIterator.Key())
 	}
 }
+
+// GetAllVaults returns all vaults with their total shares and owner shares.
+func (k Keeper) GetAllVaults(ctx sdk.Context) []*types.Vault {
+	vaults := []*types.Vault{}
+	totalSharesIterator := k.getTotalSharesIterator(ctx)
+	defer totalSharesIterator.Close()
+	for ; totalSharesIterator.Valid(); totalSharesIterator.Next() {
+		vaultId, err := types.GetVaultIdFromStateKey(totalSharesIterator.Key())
+		if err != nil {
+			panic(err)
+		}
+
+		var totalShares types.NumShares
+		k.cdc.MustUnmarshal(totalSharesIterator.Value(), &totalShares)
+
+		allOwnerShares := k.GetAllOwnerShares(ctx, *vaultId)
+
+		vaults = append(vaults, &types.Vault{
+			VaultId:     vaultId,
+			TotalShares: &totalShares,
+			OwnerShares: allOwnerShares,
+		})
+	}
+	return vaults
+}

protocol/x/vault/types/errors.go

@@ -70,4 +70,19 @@ var (
 		13,
 		"ActivationThresholdQuoteQuantums must be non-negative",
 	)
+	ErrInvalidOrderSize = errorsmod.Register(
+		ModuleName,
+		14,
+		"OrderSize is invalid",
+	)
+	ErrInvalidOwner = errorsmod.Register(
+		ModuleName,
+		15,
+		"Owner is invalid",
+	)
+	ErrMismatchedTotalAndOwnerShares = errorsmod.Register(
+		ModuleName,
+		16,
+		"TotalShares does not match sum of OwnerShares",
+	)
 )

protocol/x/vault/types/genesis.go

@@ -10,5 +10,32 @@ func DefaultGenesis() *GenesisState {
 // Validate performs basic genesis state validation returning an error upon any
 // failure.
 func (gs GenesisState) Validate() error {
-	return gs.Params.Validate()
+	// Validate params.
+	if err := gs.Params.Validate(); err != nil {
+		return err
+	}
+
+	// Validate vaults, ensuring that for each vault:
+	// 1. TotalShares is non-negative.
+	// 2. OwnerShares is non-negative.
+	// 3. TotalShares is equal to the sum of OwnerShares.
+	// 4. Owner is not empty.
+	for _, vault := range gs.Vaults {
+		totalShares := vault.TotalShares.NumShares.BigInt()
+		if totalShares.Sign() == -1 {
+			return ErrNegativeShares
+		}
+		for _, ownerShares := range vault.OwnerShares {
+			if ownerShares.Owner == "" {
+				return ErrInvalidOwner
+			} else if ownerShares.Shares.NumShares.Sign() == -1 {
+				return ErrNegativeShares
+			}
+			totalShares.Sub(totalShares, ownerShares.Shares.NumShares.BigInt())
+		}
+		if totalShares.Sign() != 0 {
+			return ErrMismatchedTotalAndOwnerShares
+		}
+	}
+	return nil
 }