From f14913597be866279909d79232eff748d8abf820 Mon Sep 17 00:00:00 2001
From: antstalepresh <36045227+antstalepresh@users.noreply.github.com>
Date: Tue, 17 May 2022 02:31:37 +1000
Subject: [PATCH] lockup module accumulation store invariant (#1182)

---
 x/lockup/keeper/invariants.go | 48 ++++++++++++++++++++++++++++++++++-
 1 file changed, 47 insertions(+), 1 deletion(-)

diff --git a/x/lockup/keeper/invariants.go b/x/lockup/keeper/invariants.go
index 5e1241bf23a..a1a1248fed0 100644
--- a/x/lockup/keeper/invariants.go
+++ b/x/lockup/keeper/invariants.go
@@ -4,15 +4,17 @@ package keeper
 
 import (
 	"fmt"
+	"time"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
 
-	"github.com/osmosis-labs/osmosis/v7/x/gamm/types"
+	"github.com/osmosis-labs/osmosis/v7/x/lockup/types"
 )
 
 // RegisterInvariants registers all governance invariants.
 func RegisterInvariants(ir sdk.InvariantRegistry, keeper Keeper) {
 	ir.RegisterRoute(types.ModuleName, "synthetic-lockup-invariant", SyntheticLockupInvariant(keeper))
+	ir.RegisterRoute(types.ModuleName, "accumulation-store-invariant", AccumulationStoreInvariant(keeper))
 }
 
 func SyntheticLockupInvariant(keeper Keeper) sdk.Invariant {
@@ -34,3 +36,47 @@ func SyntheticLockupInvariant(keeper Keeper) sdk.Invariant {
 		return sdk.FormatInvariant(types.ModuleName, "synthetic-lockup-invariant", "All synthetic lockup invariant passed"), false
 	}
 }
+
+// AccumulationStoreInvariant ensures that the sum of all lockups at a given duration
+// is equal to the value stored within the accumulation store.
+func AccumulationStoreInvariant(keeper Keeper) sdk.Invariant {
+	return func(ctx sdk.Context) (string, bool) {
+		moduleAcc := keeper.ak.GetModuleAccount(ctx, types.ModuleName)
+		balances := keeper.bk.GetAllBalances(ctx, moduleAcc.GetAddress())
+
+		// check 1s, 1 day, 1 week, 2 weeks
+		durations := []time.Duration{
+			time.Second,
+			time.Hour * 24,
+			time.Hour * 24 * 7,
+			time.Hour * 24 * 14,
+		}
+
+		// loop all denoms on lockup module
+		for _, coin := range balances {
+			denom := coin.Denom
+			for _, duration := range durations {
+				accumulation := keeper.GetPeriodLocksAccumulation(ctx, types.QueryCondition{
+					LockQueryType: types.ByDuration,
+					Denom:         denom,
+					Duration:      duration,
+				})
+
+				locks := keeper.GetLocksLongerThanDurationDenom(ctx, denom, duration)
+				lockupSum := sdk.ZeroInt()
+				for _, lock := range locks {
+					lockupSum = lockupSum.Add(lock.Coins.AmountOf(denom))
+				}
+
+				if !accumulation.Equal(lockupSum) {
+					return sdk.FormatInvariant(types.ModuleName, "accumulation-store-invariant",
+						fmt.Sprintf("\taccumulation store value does not fit actual lockup sum: %s != %s\n",
+							accumulation.String(), lockupSum.String(),
+						)), true
+				}
+			}
+		}
+
+		return sdk.FormatInvariant(types.ModuleName, "accumulation-store-invariant", "All lockup accumulation invariant passed"), false
+	}
+}