add total superfluid delegation sum invariant

proto/osmosis/superfluid/genesis.proto

@@ -16,4 +16,6 @@ message GenesisState {
       [ (gogoproto.nullable) = false ];
   repeated SuperfluidIntermediaryAccount intermediary_accounts = 4
       [ (gogoproto.nullable) = false ];
+  repeated LockIdIntermediaryAccountConnection intemediary_account_connections = 5
+      [ (gogoproto.nullable) = false ];
 }

proto/osmosis/superfluid/superfluid.proto

@@ -60,3 +60,8 @@ message SuperfluidDelegationRecord {
     (gogoproto.castrepeated) = "github.com/cosmos/cosmos-sdk/types.Coin"
   ];
 }
+
+message LockIdIntermediaryAccountConnection {
+  uint64 lock_id = 1;
+  string intermediary_account = 2;
+}

x/superfluid/genesis.go

@@ -20,12 +20,22 @@ func InitGenesis(ctx sdk.Context, k keeper.Keeper, genState types.GenesisState)
 	for _, multiplierRecord := range genState.OsmoEquivalentMultipliers {
 		k.SetOsmoEquivalentMultiplier(ctx, multiplierRecord.EpochNumber, multiplierRecord.Denom, multiplierRecord.Multiplier)
 	}
+
+	// initialize lock id and intermediary connections
+	for _, connection := range genState.IntemediaryAccountConnections {
+		acc, err := sdk.AccAddressFromBech32(connection.IntermediaryAccount)
+		if err != nil {
+			panic(err)
+		}
+		k.SetLockIdIntermediaryAccountConnection(ctx, connection.LockId, acc)
+	}
 }
 
 // ExportGenesis returns the capability module's exported genesis.
 func ExportGenesis(ctx sdk.Context, k keeper.Keeper) *types.GenesisState {
 	return &types.GenesisState{
-		SuperfluidAssets:          k.GetAllSuperfluidAssets(ctx),
-		OsmoEquivalentMultipliers: k.GetAllOsmoEquivalentMultipliers(ctx),
+		SuperfluidAssets:              k.GetAllSuperfluidAssets(ctx),
+		OsmoEquivalentMultipliers:     k.GetAllOsmoEquivalentMultipliers(ctx),
+		IntemediaryAccountConnections: k.GetAllLockIdIntermediaryAccountConnections(ctx),
 	}
 }

x/superfluid/keeper/intermediary_account.go

@@ -113,11 +113,11 @@ func (k Keeper) DeleteIntermediaryAccount(ctx sdk.Context, address sdk.AccAddres
 	prefixStore.Delete(address)
 }
 
-func (k Keeper) SetLockIdIntermediaryAccountConnection(ctx sdk.Context, lockId uint64, acc types.SuperfluidIntermediaryAccount) {
+func (k Keeper) SetLockIdIntermediaryAccountConnection(ctx sdk.Context, lockId uint64, acc sdk.AccAddress) {
 	store := ctx.KVStore(k.storeKey)
 	prefixStore := prefix.NewStore(store, types.KeyPrefixLockIntermediaryAccAddr)
 
-	prefixStore.Set(sdk.Uint64ToBigEndian(lockId), acc.GetAccAddress())
+	prefixStore.Set(sdk.Uint64ToBigEndian(lockId), acc)
 }
 
 func (k Keeper) GetLockIdIntermediaryAccountConnection(ctx sdk.Context, lockId uint64) sdk.AccAddress {
@@ -127,6 +127,22 @@ func (k Keeper) GetLockIdIntermediaryAccountConnection(ctx sdk.Context, lockId u
 	return prefixStore.Get(sdk.Uint64ToBigEndian(lockId))
 }
 
+func (k Keeper) GetAllLockIdIntermediaryAccountConnections(ctx sdk.Context) []types.LockIdIntermediaryAccountConnection {
+	store := ctx.KVStore(k.storeKey)
+	prefixStore := prefix.NewStore(store, types.KeyPrefixLockIntermediaryAccAddr)
+
+	iterator := prefixStore.Iterator(nil, nil)
+
+	connections := []types.LockIdIntermediaryAccountConnection{}
+	for ; iterator.Valid(); iterator.Next() {
+		connections = append(connections, types.LockIdIntermediaryAccountConnection{
+			LockId:              sdk.BigEndianToUint64(iterator.Key()),
+			IntermediaryAccount: sdk.AccAddress(iterator.Value()).String(),
+		})
+	}
+	return connections
+}
+
 // Returns Superfluid Intermediate Account and a bool if found / not found
 func (k Keeper) GetIntermediaryAccountFromLockId(ctx sdk.Context, lockId uint64) (types.SuperfluidIntermediaryAccount, bool) {
 	addr := k.GetLockIdIntermediaryAccountConnection(ctx, lockId)

x/superfluid/keeper/intermediary_account_test.go

@@ -51,16 +51,21 @@ func (suite *KeeperTestSuite) TestLockIdIntermediaryAccountConnection() {
 	// set account
 	valAddr := sdk.ValAddress([]byte("addr1---------------"))
 	acc := types.NewSuperfluidIntermediaryAccount("gamm/pool/1", valAddr.String(), 1)
-	suite.app.SuperfluidKeeper.SetLockIdIntermediaryAccountConnection(suite.ctx, 1, acc)
+	suite.app.SuperfluidKeeper.SetLockIdIntermediaryAccountConnection(suite.ctx, 1, acc.GetAccAddress())
 
 	// get account
 	addr = suite.app.SuperfluidKeeper.GetLockIdIntermediaryAccountConnection(suite.ctx, 1)
 	suite.Require().Equal(addr.String(), acc.GetAccAddress().String())
 
+	// check get all
+	conns := suite.app.SuperfluidKeeper.GetAllLockIdIntermediaryAccountConnections(suite.ctx)
+	suite.Require().Len(conns, 1)
+
 	// delete account
 	suite.app.SuperfluidKeeper.DeleteLockIdIntermediaryAccountConnection(suite.ctx, 1)
 
 	// get account
 	addr = suite.app.SuperfluidKeeper.GetLockIdIntermediaryAccountConnection(suite.ctx, 1)
 	suite.Require().Equal(addr.String(), "")
+
 }

x/superfluid/keeper/invariants.go

@@ -0,0 +1,72 @@
+package keeper
+
+import (
+	sdk "github.com/cosmos/cosmos-sdk/types"
+
+	"github.com/osmosis-labs/osmosis/v7/x/superfluid/types"
+)
+
+const totalSuperfluidDelegationInvariantName = "total-superfluid-delegation-invariant-name"
+
+// RegisterInvariants registers all governance invariants
+func RegisterInvariants(ir sdk.InvariantRegistry, keeper Keeper) {
+	ir.RegisterRoute(types.ModuleName, totalSuperfluidDelegationInvariantName, TotalSuperfluidDelegationInvariant(keeper))
+}
+
+// AllInvariants runs all invariants of the gamm module
+func AllInvariants(keeper Keeper) sdk.Invariant {
+	return func(ctx sdk.Context) (string, bool) {
+		return TotalSuperfluidDelegationInvariant(keeper)(ctx)
+	}
+}
+
+// TotalSuperfluidDelegationInvariant checks the sum of intermediary account delegation is same as sum of individual lockup delegation
+func TotalSuperfluidDelegationInvariant(keeper Keeper) sdk.Invariant {
+	return func(ctx sdk.Context) (string, bool) {
+		accs := keeper.GetAllIntermediaryAccounts(ctx)
+		totalSuperfluidDelegationTokens := sdk.ZeroDec()
+
+		for _, acc := range accs {
+			valAddr, err := sdk.ValAddressFromBech32(acc.ValAddr)
+			if err != nil {
+				return sdk.FormatInvariant(types.ModuleName, totalSuperfluidDelegationInvariantName,
+					"\tinvalid validator address exists"), true
+			}
+			validator, found := keeper.sk.GetValidator(ctx, valAddr)
+			if !found {
+				return sdk.FormatInvariant(types.ModuleName, totalSuperfluidDelegationInvariantName,
+					"\tvalidator does not exists for specified validator address on intermediary account"), true
+			}
+			delegation, found := keeper.sk.GetDelegation(ctx, acc.GetAccAddress(), valAddr)
+			if found {
+				tokens := validator.TokensFromShares(delegation.Shares)
+				totalSuperfluidDelegationTokens = totalSuperfluidDelegationTokens.Add(tokens)
+			}
+		}
+
+		totalExpectedSuperfluidAmount := sdk.ZeroInt()
+		connections := keeper.GetAllLockIdIntermediaryAccountConnections(ctx)
+		for _, connection := range connections {
+			lockId := connection.LockId
+			lock, err := keeper.lk.GetLockByID(ctx, lockId)
+			if err != nil || lock == nil {
+				return sdk.FormatInvariant(types.ModuleName, totalSuperfluidDelegationInvariantName,
+					"\tinvalid superfluid lock id exists with no actual lockup"), true
+			}
+			if len(lock.Coins) != 1 {
+				return sdk.FormatInvariant(types.ModuleName, totalSuperfluidDelegationInvariantName,
+					"\tonly single coin lockup is eligible for superfluid staking"), true
+			}
+			amount := keeper.GetSuperfluidOSMOTokens(ctx, lock.Coins[0].Denom, lock.Coins[0].Amount)
+			totalExpectedSuperfluidAmount = totalExpectedSuperfluidAmount.Add(amount)
+		}
+
+		if !totalExpectedSuperfluidAmount.ToDec().Equal(totalSuperfluidDelegationTokens) {
+			return sdk.FormatInvariant(types.ModuleName, totalSuperfluidDelegationInvariantName,
+				"\ttotal superfluid intermediary account delegation amount does not match total sum of lockup delegations\n"), true
+		}
+
+		return sdk.FormatInvariant(types.ModuleName, totalSuperfluidDelegationInvariantName,
+			"\ttotal superfluid intermediary account delegation amount matches total sum of lockup delegations\n"), false
+	}
+}

x/superfluid/keeper/stake.go

@@ -177,7 +177,7 @@ func (k Keeper) SuperfluidDelegate(ctx sdk.Context, sender string, lockID uint64
 		return err
 	}
 	// create connection record between lock id and intermediary account
-	k.SetLockIdIntermediaryAccountConnection(ctx, lockID, acc)
+	k.SetLockIdIntermediaryAccountConnection(ctx, lockID, acc.GetAccAddress())
 
 	// Register a synthetic lockup for superfluid staking
 	err = k.createSyntheticLockup(ctx, lockID, acc, bondedStatus)

x/superfluid/module.go

@@ -157,7 +157,9 @@ func (am AppModule) RegisterServices(cfg module.Configurator) {
 }
 
 // RegisterInvariants registers the capability module's invariants.
-func (am AppModule) RegisterInvariants(_ sdk.InvariantRegistry) {}
+func (am AppModule) RegisterInvariants(ir sdk.InvariantRegistry) {
+	keeper.RegisterInvariants(ir, am.keeper)
+}
 
 // InitGenesis performs the capability module's genesis initialization It returns
 // no validator updates.