Merge PR #2597: Add distribution accum invariants

PENDING.md

@@ -35,6 +35,7 @@ IMPROVEMENTS
 * Gaia
 
 * SDK
+ - #2573 [x/distribution] add accum invariance
 
 * Tendermint
 
@@ -48,5 +49,7 @@ BUG FIXES
 * Gaia
 
 * SDK
+ - #2573 [x/distribution] accum invariance bugfix
+ - #2573 [x/slashing] unbonding-delegation slashing invariance bugfix
 
 * Tendermint

cmd/gaia/app/sim_test.go

@@ -127,7 +127,9 @@ func invariants(app *GaiaApp) []simulation.Invariant {
 	return []simulation.Invariant{
 		banksim.NonnegativeBalanceInvariant(app.accountKeeper),
 		govsim.AllInvariants(),
-		stakesim.AllInvariants(app.bankKeeper, app.stakeKeeper, app.feeCollectionKeeper, app.distrKeeper, app.accountKeeper),
+		distrsim.AllInvariants(app.distrKeeper, app.stakeKeeper),
+		stakesim.AllInvariants(app.bankKeeper, app.stakeKeeper,
+			app.feeCollectionKeeper, app.distrKeeper, app.accountKeeper),
 		slashingsim.AllInvariants(),
 	}
 }

x/bank/simulation/invariants.go

@@ -14,7 +14,7 @@ import (
 
 // NonnegativeBalanceInvariant checks that all accounts in the application have non-negative balances
 func NonnegativeBalanceInvariant(mapper auth.AccountKeeper) simulation.Invariant {
-	return func(app *baseapp.BaseApp) error {
+	return func(app *baseapp.BaseApp, _ abci.Header) error {
 		ctx := app.NewContext(false, abci.Header{})
 		accts := mock.GetAllAccounts(mapper, ctx)
 		for _, acc := range accts {
@@ -32,7 +32,7 @@ func NonnegativeBalanceInvariant(mapper auth.AccountKeeper) simulation.Invariant
 // TotalCoinsInvariant checks that the sum of the coins across all accounts
 // is what is expected
 func TotalCoinsInvariant(mapper auth.AccountKeeper, totalSupplyFn func() sdk.Coins) simulation.Invariant {
-	return func(app *baseapp.BaseApp) error {
+	return func(app *baseapp.BaseApp, _ abci.Header) error {
 		ctx := app.NewContext(false, abci.Header{})
 		totalCoins := sdk.Coins{}
 

x/distribution/alias.go

@@ -23,6 +23,11 @@ type (
 	MsgWithdrawValidatorRewardsAll = types.MsgWithdrawValidatorRewardsAll
 
 	GenesisState = types.GenesisState
+
+	// expected keepers
+	StakeKeeper         = types.StakeKeeper
+	BankKeeper          = types.BankKeeper
+	FeeCollectionKeeper = types.FeeCollectionKeeper
 )
 
 var (
@@ -62,9 +67,7 @@ var (
 	ErrNilDelegatorAddr = types.ErrNilDelegatorAddr
 	ErrNilWithdrawAddr  = types.ErrNilWithdrawAddr
 	ErrNilValidatorAddr = types.ErrNilValidatorAddr
-)
 
-var (
 	ActionModifyWithdrawAddress       = tags.ActionModifyWithdrawAddress
 	ActionWithdrawDelegatorRewardsAll = tags.ActionWithdrawDelegatorRewardsAll
 	ActionWithdrawDelegatorReward     = tags.ActionWithdrawDelegatorReward

x/distribution/keeper/delegation.go

@@ -69,78 +69,110 @@ func (k Keeper) RemoveDelegatorWithdrawAddr(ctx sdk.Context, delAddr, withdrawAd
 
 //___________________________________________________________________________________________
 
-// Withdraw all the rewards for a single delegation
+// return all rewards for a delegation
+func (k Keeper) withdrawDelegationReward(ctx sdk.Context,
+	delAddr sdk.AccAddress, valAddr sdk.ValAddress) (types.FeePool,
+	types.ValidatorDistInfo, types.DelegationDistInfo, types.DecCoins) {
+
+	wc := k.GetWithdrawContext(ctx, valAddr)
+	valInfo := k.GetValidatorDistInfo(ctx, valAddr)
+	delInfo := k.GetDelegationDistInfo(ctx, delAddr, valAddr)
+	validator := k.stakeKeeper.Validator(ctx, valAddr)
+	delegation := k.stakeKeeper.Delegation(ctx, delAddr, valAddr)
+
+	delInfo, valInfo, feePool, withdraw := delInfo.WithdrawRewards(wc, valInfo,
+		validator.GetDelegatorShares(), delegation.GetShares())
+
+	return feePool, valInfo, delInfo, withdraw
+}
+
+// get all rewards for all delegations of a delegator
+func (k Keeper) currentDelegationReward(ctx sdk.Context, delAddr sdk.AccAddress,
+	valAddr sdk.ValAddress) types.DecCoins {
+
+	wc := k.GetWithdrawContext(ctx, valAddr)
+
+	valInfo := k.GetValidatorDistInfo(ctx, valAddr)
+	delInfo := k.GetDelegationDistInfo(ctx, delAddr, valAddr)
+	validator := k.stakeKeeper.Validator(ctx, valAddr)
+	delegation := k.stakeKeeper.Delegation(ctx, delAddr, valAddr)
+
+	estimation := delInfo.CurrentRewards(wc, valInfo,
+		validator.GetDelegatorShares(), delegation.GetShares())
+
+	return estimation
+}
+
+//___________________________________________________________________________________________
+
+// withdraw all rewards for a single delegation
 // NOTE: This gets called "onDelegationSharesModified",
 // meaning any changes to bonded coins
-func (k Keeper) WithdrawDelegationReward(ctx sdk.Context, delegatorAddr sdk.AccAddress,
+func (k Keeper) WithdrawToDelegator(ctx sdk.Context, feePool types.FeePool,
-	valAddr sdk.ValAddress) sdk.Error {
+	delAddr sdk.AccAddress, amount types.DecCoins) {
 
-	if !k.HasDelegationDistInfo(ctx, delegatorAddr, valAddr) {
+	withdrawAddr := k.GetDelegatorWithdrawAddr(ctx, delAddr)
-		return types.ErrNoDelegationDistInfo(k.codespace)
+	coinsToAdd, change := amount.TruncateDecimal()
-	}
-
-	// TODO: Reconcile with duplicate code in getDelegatorRewardsAll.
-	height := ctx.BlockHeight()
-	lastTotalPower := sdk.NewDecFromInt(k.stakeKeeper.GetLastTotalPower(ctx))
-	lastValPower := sdk.NewDecFromInt(k.stakeKeeper.GetLastValidatorPower(ctx, valAddr))
-	feePool := k.GetFeePool(ctx)
-	delInfo := k.GetDelegationDistInfo(ctx, delegatorAddr, valAddr)
-	valInfo := k.GetValidatorDistInfo(ctx, valAddr)
-	validator := k.stakeKeeper.Validator(ctx, valAddr)
-	delegation := k.stakeKeeper.Delegation(ctx, delegatorAddr, valAddr)
-
-	delInfo, valInfo, feePool, withdraw := delInfo.WithdrawRewards(feePool, valInfo, height, lastTotalPower,
-		lastValPower, validator.GetDelegatorShares(), delegation.GetShares(), validator.GetCommission())
-
-	k.SetValidatorDistInfo(ctx, valInfo)
-	k.SetDelegationDistInfo(ctx, delInfo)
-	withdrawAddr := k.GetDelegatorWithdrawAddr(ctx, delegatorAddr)
-	coinsToAdd, change := withdraw.TruncateDecimal()
 	feePool.CommunityPool = feePool.CommunityPool.Plus(change)
 	k.SetFeePool(ctx, feePool)
 	_, _, err := k.bankKeeper.AddCoins(ctx, withdrawAddr, coinsToAdd)
 	if err != nil {
 		panic(err)
 	}
+}
+
+//___________________________________________________________________________________________
+
+// withdraw all rewards for a single delegation
+// NOTE: This gets called "onDelegationSharesModified",
+// meaning any changes to bonded coins
+func (k Keeper) WithdrawDelegationReward(ctx sdk.Context, delAddr sdk.AccAddress,
+	valAddr sdk.ValAddress) sdk.Error {
+
+	if !k.HasDelegationDistInfo(ctx, delAddr, valAddr) {
+		return types.ErrNoDelegationDistInfo(k.codespace)
+	}
+
+	feePool, valInfo, delInfo, withdraw :=
+		k.withdrawDelegationReward(ctx, delAddr, valAddr)
+
+	k.SetValidatorDistInfo(ctx, valInfo)
+	k.SetDelegationDistInfo(ctx, delInfo)
+	k.WithdrawToDelegator(ctx, feePool, delAddr, withdraw)
 	return nil
 }
 
+// current rewards for a single delegation
+func (k Keeper) CurrentDelegationReward(ctx sdk.Context, delAddr sdk.AccAddress,
+	valAddr sdk.ValAddress) (sdk.Coins, sdk.Error) {
+
+	if !k.HasDelegationDistInfo(ctx, delAddr, valAddr) {
+		return sdk.Coins{}, types.ErrNoDelegationDistInfo(k.codespace)
+	}
+	estCoins := k.currentDelegationReward(ctx, delAddr, valAddr)
+	trucate, _ := estCoins.TruncateDecimal()
+	return trucate, nil
+}
+
 //___________________________________________________________________________________________
 
 // return all rewards for all delegations of a delegator
-func (k Keeper) WithdrawDelegationRewardsAll(ctx sdk.Context, delegatorAddr sdk.AccAddress) {
+func (k Keeper) WithdrawDelegationRewardsAll(ctx sdk.Context, delAddr sdk.AccAddress) {
-	height := ctx.BlockHeight()
+	withdraw := k.withdrawDelegationRewardsAll(ctx, delAddr)
-	withdraw := k.getDelegatorRewardsAll(ctx, delegatorAddr, height)
 	feePool := k.GetFeePool(ctx)
-	withdrawAddr := k.GetDelegatorWithdrawAddr(ctx, delegatorAddr)
+	k.WithdrawToDelegator(ctx, feePool, delAddr, withdraw)
-	coinsToAdd, change := withdraw.TruncateDecimal()
-	feePool.CommunityPool = feePool.CommunityPool.Plus(change)
-	k.SetFeePool(ctx, feePool)
-	_, _, err := k.bankKeeper.AddCoins(ctx, withdrawAddr, coinsToAdd)
-	if err != nil {
-		panic(err)
-	}
 }
 
-// return all rewards for all delegations of a delegator
+func (k Keeper) withdrawDelegationRewardsAll(ctx sdk.Context,
-func (k Keeper) getDelegatorRewardsAll(ctx sdk.Context, delAddr sdk.AccAddress, height int64) types.DecCoins {
+	delAddr sdk.AccAddress) types.DecCoins {
-
-	withdraw := types.DecCoins{}
-	lastTotalPower := sdk.NewDecFromInt(k.stakeKeeper.GetLastTotalPower(ctx))
 
 	// iterate over all the delegations
-	// TODO: Reconcile with duplicate code in WithdrawDelegationReward.
+	withdraw := types.DecCoins{}
 	operationAtDelegation := func(_ int64, del sdk.Delegation) (stop bool) {
-		feePool := k.GetFeePool(ctx)
-		valAddr := del.GetValidatorAddr()
-		lastValPower := sdk.NewDecFromInt(k.stakeKeeper.GetLastValidatorPower(ctx, valAddr))
-		delInfo := k.GetDelegationDistInfo(ctx, delAddr, valAddr)
-		valInfo := k.GetValidatorDistInfo(ctx, valAddr)
-		validator := k.stakeKeeper.Validator(ctx, valAddr)
-		delegation := k.stakeKeeper.Delegation(ctx, delAddr, valAddr)
 
-		delInfo, valInfo, feePool, diWithdraw := delInfo.WithdrawRewards(feePool, valInfo, height, lastTotalPower,
+		valAddr := del.GetValidatorAddr()
-			lastValPower, validator.GetDelegatorShares(), delegation.GetShares(), validator.GetCommission())
+		feePool, valInfo, delInfo, diWithdraw :=
+			k.withdrawDelegationReward(ctx, delAddr, valAddr)
 		withdraw = withdraw.Plus(diWithdraw)
 		k.SetFeePool(ctx, feePool)
 		k.SetValidatorDistInfo(ctx, valInfo)
@@ -150,3 +182,19 @@ func (k Keeper) getDelegatorRewardsAll(ctx sdk.Context, delAddr sdk.AccAddress,
 	k.stakeKeeper.IterateDelegations(ctx, delAddr, operationAtDelegation)
 	return withdraw
 }
+
+// get all rewards for all delegations of a delegator
+func (k Keeper) CurrentDelegationRewardsAll(ctx sdk.Context,
+	delAddr sdk.AccAddress) types.DecCoins {
+
+	// iterate over all the delegations
+	total := types.DecCoins{}
+	operationAtDelegation := func(_ int64, del sdk.Delegation) (stop bool) {
+		valAddr := del.GetValidatorAddr()
+		est := k.currentDelegationReward(ctx, delAddr, valAddr)
+		total = total.Plus(est)
+		return false
+	}
+	k.stakeKeeper.IterateDelegations(ctx, delAddr, operationAtDelegation)
+	return total
+}

x/distribution/keeper/hooks.go

@@ -29,6 +29,15 @@ func (k Keeper) onValidatorModified(ctx sdk.Context, valAddr sdk.ValAddress) {
 	}
 }
 
+// Withdraw all validator rewards
+func (k Keeper) onValidatorBonded(ctx sdk.Context, valAddr sdk.ValAddress) {
+	lastPower := k.stakeKeeper.GetLastValidatorPower(ctx, valAddr)
+	if !lastPower.Equal(sdk.ZeroInt()) {
+		panic("expected last power to be 0 for validator entering bonded state")
+	}
+	k.onValidatorModified(ctx, valAddr)
+}
+
 // XXX Consider removing this after debugging.
 func (k Keeper) onValidatorPowerDidChange(ctx sdk.Context, valAddr sdk.ValAddress) {
 	vi := k.GetValidatorDistInfo(ctx, valAddr)
@@ -109,7 +118,7 @@ func (h Hooks) OnValidatorBeginUnbonding(ctx sdk.Context, _ sdk.ConsAddress, val
 	h.k.onValidatorModified(ctx, valAddr)
 }
 func (h Hooks) OnValidatorBonded(ctx sdk.Context, _ sdk.ConsAddress, valAddr sdk.ValAddress) {
-	h.k.onValidatorModified(ctx, valAddr)
+	h.k.onValidatorBonded(ctx, valAddr)
 }
 func (h Hooks) OnValidatorPowerDidChange(ctx sdk.Context, _ sdk.ConsAddress, valAddr sdk.ValAddress) {
 	h.k.onValidatorPowerDidChange(ctx, valAddr)

x/distribution/keeper/keeper.go

@@ -55,6 +55,17 @@ func (k Keeper) SetFeePool(ctx sdk.Context, feePool types.FeePool) {
 	store.Set(FeePoolKey, b)
 }
 
+// get the total validator accum for the ctx height
+// in the fee pool
+func (k Keeper) GetFeePoolValAccum(ctx sdk.Context) sdk.Dec {
+
+	// withdraw self-delegation
+	height := ctx.BlockHeight()
+	totalPower := sdk.NewDecFromInt(k.stakeKeeper.GetLastTotalPower(ctx))
+	fp := k.GetFeePool(ctx)
+	return fp.GetTotalValAccum(height, totalPower)
+}
+
 //______________________________________________________________________
 
 // set the proposer public key for this block
@@ -77,6 +88,23 @@ func (k Keeper) SetPreviousProposerConsAddr(ctx sdk.Context, consAddr sdk.ConsAd
 	store.Set(ProposerKey, b)
 }
 
+//______________________________________________________________________
+
+// get context required for withdraw operations
+func (k Keeper) GetWithdrawContext(ctx sdk.Context,
+	valOperatorAddr sdk.ValAddress) types.WithdrawContext {
+
+	feePool := k.GetFeePool(ctx)
+	height := ctx.BlockHeight()
+	validator := k.stakeKeeper.Validator(ctx, valOperatorAddr)
+	lastValPower := k.stakeKeeper.GetLastValidatorPower(ctx, valOperatorAddr)
+	lastTotalPower := sdk.NewDecFromInt(k.stakeKeeper.GetLastTotalPower(ctx))
+
+	return types.NewWithdrawContext(
+		feePool, height, lastTotalPower, sdk.NewDecFromInt(lastValPower),
+		validator.GetCommission())
+}
+
 //______________________________________________________________________
 // PARAM STORE
 

x/distribution/keeper/test_common.go

@@ -158,3 +158,24 @@ func (fck DummyFeeCollectionKeeper) SetCollectedFees(in sdk.Coins) {
 func (fck DummyFeeCollectionKeeper) ClearCollectedFees(_ sdk.Context) {
 	heldFees = sdk.Coins{}
 }
+
+//__________________________________________________________________________________
+// used in simulation
+
+// iterate over all the validator distribution infos (inefficient, just used to check invariants)
+func (k Keeper) IterateValidatorDistInfos(ctx sdk.Context,
+	fn func(index int64, distInfo types.ValidatorDistInfo) (stop bool)) {
+
+	store := ctx.KVStore(k.storeKey)
+	iter := sdk.KVStorePrefixIterator(store, ValidatorDistInfoKey)
+	defer iter.Close()
+	index := int64(0)
+	for ; iter.Valid(); iter.Next() {
+		var vdi types.ValidatorDistInfo
+		k.cdc.MustUnmarshalBinary(iter.Value(), &vdi)
+		if fn(index, vdi) {
+			return
+		}
+		index++
+	}
+}

x/distribution/keeper/validator.go

@@ -40,6 +40,22 @@ func (k Keeper) RemoveValidatorDistInfo(ctx sdk.Context, valAddr sdk.ValAddress)
 	store.Delete(GetValidatorDistInfoKey(valAddr))
 }
 
+// Get the calculated accum of a validator at the current block
+// without affecting the state.
+func (k Keeper) GetValidatorAccum(ctx sdk.Context, operatorAddr sdk.ValAddress) (sdk.Dec, sdk.Error) {
+	if !k.HasValidatorDistInfo(ctx, operatorAddr) {
+		return sdk.Dec{}, types.ErrNoValidatorDistInfo(k.codespace)
+	}
+
+	// withdraw self-delegation
+	height := ctx.BlockHeight()
+	lastValPower := k.stakeKeeper.GetLastValidatorPower(ctx, operatorAddr)
+	valInfo := k.GetValidatorDistInfo(ctx, operatorAddr)
+	accum := valInfo.GetValAccum(height, sdk.NewDecFromInt(lastValPower))
+
+	return accum, nil
+}
+
 // withdrawal all the validator rewards including the commission
 func (k Keeper) WithdrawValidatorRewardsAll(ctx sdk.Context, operatorAddr sdk.ValAddress) sdk.Error {
 
@@ -48,45 +64,37 @@ func (k Keeper) WithdrawValidatorRewardsAll(ctx sdk.Context, operatorAddr sdk.Va
 	}
 
 	// withdraw self-delegation
-	height := ctx.BlockHeight()
-	validator := k.stakeKeeper.Validator(ctx, operatorAddr)
-	lastValPower := sdk.NewDecFromInt(k.stakeKeeper.GetLastValidatorPower(ctx, operatorAddr))
 	accAddr := sdk.AccAddress(operatorAddr.Bytes())
-	withdraw := k.getDelegatorRewardsAll(ctx, accAddr, height)
+	withdraw := k.withdrawDelegationRewardsAll(ctx, accAddr)
 
 	// withdrawal validator commission rewards
-	lastTotalPower := sdk.NewDecFromInt(k.stakeKeeper.GetLastTotalPower(ctx))
 	valInfo := k.GetValidatorDistInfo(ctx, operatorAddr)
-	feePool := k.GetFeePool(ctx)
+	wc := k.GetWithdrawContext(ctx, operatorAddr)
-	valInfo, feePool, commission := valInfo.WithdrawCommission(feePool, height, lastTotalPower,
+	valInfo, feePool, commission := valInfo.WithdrawCommission(wc)
-		lastValPower, validator.GetCommission())
 	withdraw = withdraw.Plus(commission)
 	k.SetValidatorDistInfo(ctx, valInfo)
 
-	withdrawAddr := k.GetDelegatorWithdrawAddr(ctx, accAddr)
+	k.WithdrawToDelegator(ctx, feePool, accAddr, withdraw)
-	truncated, change := withdraw.TruncateDecimal()
-	feePool.CommunityPool = feePool.CommunityPool.Plus(change)
-	k.SetFeePool(ctx, feePool)
-	_, _, err := k.bankKeeper.AddCoins(ctx, withdrawAddr, truncated)
-	if err != nil {
-		panic(err)
-	}
-
 	return nil
 }
 
-// iterate over all the validator distribution infos (inefficient, just used to check invariants)
+// get all the validator rewards including the commission
-func (k Keeper) IterateValidatorDistInfos(ctx sdk.Context, fn func(index int64, distInfo types.ValidatorDistInfo) (stop bool)) {
+func (k Keeper) CurrentValidatorRewardsAll(ctx sdk.Context, operatorAddr sdk.ValAddress) (sdk.Coins, sdk.Error) {
-	store := ctx.KVStore(k.storeKey)
+
-	iter := sdk.KVStorePrefixIterator(store, ValidatorDistInfoKey)
+	if !k.HasValidatorDistInfo(ctx, operatorAddr) {
-	defer iter.Close()
+		return sdk.Coins{}, types.ErrNoValidatorDistInfo(k.codespace)
-	index := int64(0)
-	for ; iter.Valid(); iter.Next() {
-		var vdi types.ValidatorDistInfo
-		k.cdc.MustUnmarshalBinary(iter.Value(), &vdi)
-		if fn(index, vdi) {
-			return
-		}
-		index++
 	}
+
+	// withdraw self-delegation
+	accAddr := sdk.AccAddress(operatorAddr.Bytes())
+	withdraw := k.CurrentDelegationRewardsAll(ctx, accAddr)
+
+	// withdrawal validator commission rewards
+	valInfo := k.GetValidatorDistInfo(ctx, operatorAddr)
+
+	wc := k.GetWithdrawContext(ctx, operatorAddr)
+	commission := valInfo.CurrentCommissionRewards(wc)
+	withdraw = withdraw.Plus(commission)
+	truncated, _ := withdraw.TruncateDecimal()
+	return truncated, nil
 }

x/distribution/simulation/invariants.go

@@ -0,0 +1,50 @@
+package simulation
+
+import (
+	"fmt"
+
+	"github.com/cosmos/cosmos-sdk/baseapp"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	distr "github.com/cosmos/cosmos-sdk/x/distribution"
+	"github.com/cosmos/cosmos-sdk/x/mock/simulation"
+	abci "github.com/tendermint/tendermint/abci/types"
+)
+
+// AllInvariants runs all invariants of the distribution module
+// Currently: total supply, positive power
+func AllInvariants(d distr.Keeper, sk distr.StakeKeeper) simulation.Invariant {
+
+	return func(app *baseapp.BaseApp, header abci.Header) error {
+		err := ValAccumInvariants(d, sk)(app, header)
+		if err != nil {
+			return err
+		}
+		return nil
+	}
+}
+
+// ValAccumInvariants checks that the fee pool accum == sum all validators' accum
+func ValAccumInvariants(k distr.Keeper, sk distr.StakeKeeper) simulation.Invariant {
+
+	return func(app *baseapp.BaseApp, header abci.Header) error {
+		ctx := app.NewContext(false, header)
+		height := ctx.BlockHeight()
+
+		valAccum := sdk.ZeroDec()
+		k.IterateValidatorDistInfos(ctx, func(_ int64, vdi distr.ValidatorDistInfo) bool {
+			lastValPower := sk.GetLastValidatorPower(ctx, vdi.OperatorAddr)
+			valAccum = valAccum.Add(vdi.GetValAccum(height, sdk.NewDecFromInt(lastValPower)))
+			return false
+		})
+
+		lastTotalPower := sdk.NewDecFromInt(sk.GetLastTotalPower(ctx))
+		totalAccum := k.GetFeePool(ctx).GetTotalValAccum(height, lastTotalPower)
+
+		if !totalAccum.Equal(valAccum) {
+			return fmt.Errorf("validator accum invariance: \n\tfee pool totalAccum: %v"+
+				"\n\tvalidator accum \t%v\n", totalAccum.String(), valAccum.String())
+		}
+
+		return nil
+	}
+}

x/distribution/types/delegator_info.go

@@ -21,6 +21,12 @@ func NewDelegationDistInfo(delegatorAddr sdk.AccAddress, valOperatorAddr sdk.Val
 	}
 }
 
+// Get the calculated accum of this delegator at the provided height
+func (di DelegationDistInfo) GetDelAccum(height int64, delegatorShares sdk.Dec) sdk.Dec {
+	blocks := height - di.DelPoolWithdrawalHeight
+	return delegatorShares.MulInt(sdk.NewInt(blocks))
+}
+
 // Withdraw rewards from delegator.
 // Among many things, it does:
 // * updates validator info's total del accum
@@ -28,21 +34,21 @@ func NewDelegationDistInfo(delegatorAddr sdk.AccAddress, valOperatorAddr sdk.Val
 //   * updates validator info's FeePoolWithdrawalHeight, thus setting accum to 0
 //   * updates fee pool to latest height and total val accum w/ given totalBonded
 //   (see comment on TakeFeePoolRewards for more info)
-func (di DelegationDistInfo) WithdrawRewards(fp FeePool, vi ValidatorDistInfo,
+func (di DelegationDistInfo) WithdrawRewards(wc WithdrawContext, vi ValidatorDistInfo,
-	height int64, totalBonded, vdTokens, totalDelShares, delegatorShares,
+	totalDelShares, delegatorShares sdk.Dec) (
-	commissionRate sdk.Dec) (DelegationDistInfo, ValidatorDistInfo, FeePool, DecCoins) {
+	DelegationDistInfo, ValidatorDistInfo, FeePool, DecCoins) {
 
-	vi = vi.UpdateTotalDelAccum(height, totalDelShares)
+	fp := wc.FeePool
+	vi = vi.UpdateTotalDelAccum(wc.Height, totalDelShares)
 
 	if vi.DelAccum.Accum.IsZero() {
 		return di, vi, fp, DecCoins{}
 	}
 
-	vi, fp = vi.TakeFeePoolRewards(fp, height, totalBonded, vdTokens, commissionRate)
+	vi, fp = vi.TakeFeePoolRewards(wc)
 
-	blocks := height - di.DelPoolWithdrawalHeight
+	accum := di.GetDelAccum(wc.Height, delegatorShares)
-	di.DelPoolWithdrawalHeight = height
+	di.DelPoolWithdrawalHeight = wc.Height
-	accum := delegatorShares.MulInt(sdk.NewInt(blocks))
 	withdrawalTokens := vi.DelPool.MulDec(accum).QuoDec(vi.DelAccum.Accum)
 	remDelPool := vi.DelPool.Minus(withdrawalTokens)
 
@@ -51,3 +57,19 @@ func (di DelegationDistInfo) WithdrawRewards(fp FeePool, vi ValidatorDistInfo,
 
 	return di, vi, fp, withdrawalTokens
 }
+
+// get the delegators rewards at this current state,
+func (di DelegationDistInfo) CurrentRewards(wc WithdrawContext, vi ValidatorDistInfo,
+	totalDelShares, delegatorShares sdk.Dec) DecCoins {
+
+	totalDelAccum := vi.GetTotalDelAccum(wc.Height, totalDelShares)
+
+	if vi.DelAccum.Accum.IsZero() {
+		return DecCoins{}
+	}
+
+	rewards := vi.CurrentPoolRewards(wc)
+	accum := di.GetDelAccum(wc.Height, delegatorShares)
+	tokens := rewards.MulDec(accum).QuoDec(totalDelAccum)
+	return tokens
+}

x/distribution/types/delegator_info_test.go

@@ -29,8 +29,10 @@ func TestWithdrawRewards(t *testing.T) {
 	fp.ValPool = DecCoins{NewDecCoin("stake", 1000)}
 
 	// withdraw rewards
-	di1, vi, fp, rewardRecv1 := di1.WithdrawRewards(fp, vi, height, totalBondedTokens,
+	wc := NewWithdrawContext(fp, height,
-		validatorTokens, validatorDelShares, di1Shares, commissionRate)
+		totalBondedTokens, validatorTokens, commissionRate)
+	di1, vi, fp, rewardRecv1 := di1.WithdrawRewards(wc, vi,
+		validatorDelShares, di1Shares)
 
 	assert.Equal(t, height, di1.DelPoolWithdrawalHeight)
 	assert.True(sdk.DecEq(t, sdk.NewDec(900), fp.TotalValAccum.Accum))
@@ -44,8 +46,10 @@ func TestWithdrawRewards(t *testing.T) {
 	fp.ValPool[0].Amount = fp.ValPool[0].Amount.Add(sdk.NewDec(1000))
 
 	// withdraw rewards
-	di2, vi, fp, rewardRecv2 := di2.WithdrawRewards(fp, vi, height, totalBondedTokens,
+	wc = NewWithdrawContext(fp, height,
-		validatorTokens, validatorDelShares, di2Shares, commissionRate)
+		totalBondedTokens, validatorTokens, commissionRate)
+	di2, vi, fp, rewardRecv2 := di2.WithdrawRewards(wc, vi,
+		validatorDelShares, di2Shares)
 
 	assert.Equal(t, height, di2.DelPoolWithdrawalHeight)
 	assert.True(sdk.DecEq(t, sdk.NewDec(1800), fp.TotalValAccum.Accum))

x/distribution/types/fee_pool.go

@@ -4,33 +4,6 @@ import (
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )
 
-// total accumulation tracker
-type TotalAccum struct {
-	UpdateHeight int64   `json:"update_height"`
-	Accum        sdk.Dec `json:"accum"`
-}
-
-func NewTotalAccum(height int64) TotalAccum {
-	return TotalAccum{
-		UpdateHeight: height,
-		Accum:        sdk.ZeroDec(),
-	}
-}
-
-// update total validator accumulation factor for the new height
-// CONTRACT: height should be greater than the old height
-func (ta TotalAccum) UpdateForNewHeight(height int64, accumCreatedPerBlock sdk.Dec) TotalAccum {
-	blocks := height - ta.UpdateHeight
-	if blocks < 0 {
-		panic("reverse updated for new height")
-	}
-	ta.Accum = ta.Accum.Add(accumCreatedPerBlock.MulInt(sdk.NewInt(blocks)))
-	ta.UpdateHeight = height
-	return ta
-}
-
-//___________________________________________________________________________________________
-
 // global fee pool for distribution
 type FeePool struct {
 	TotalValAccum TotalAccum `json:"val_accum"`      // total valdator accum held by validators
@@ -45,6 +18,11 @@ func (f FeePool) UpdateTotalValAccum(height int64, totalBondedTokens sdk.Dec) Fe
 	return f
 }
 
+// get the total validator accum for the fee pool without modifying the state
+func (f FeePool) GetTotalValAccum(height int64, totalBondedTokens sdk.Dec) sdk.Dec {
+	return f.TotalValAccum.GetAccum(height, totalBondedTokens)
+}
+
 // zero fee pool
 func InitialFeePool() FeePool {
 	return FeePool{

x/distribution/types/fee_pool_test.go

@@ -7,17 +7,6 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func TestTotalAccumUpdateForNewHeight(t *testing.T) {
-
-	ta := NewTotalAccum(0)
-
-	ta = ta.UpdateForNewHeight(5, sdk.NewDec(3))
-	require.True(sdk.DecEq(t, sdk.NewDec(15), ta.Accum))
-
-	ta = ta.UpdateForNewHeight(8, sdk.NewDec(2))
-	require.True(sdk.DecEq(t, sdk.NewDec(21), ta.Accum))
-}
-
 func TestUpdateTotalValAccum(t *testing.T) {
 
 	fp := InitialFeePool()

x/distribution/types/total_accum.go

@@ -0,0 +1,40 @@
+package types
+
+import (
+	sdk "github.com/cosmos/cosmos-sdk/types"
+)
+
+// total accumulation tracker
+type TotalAccum struct {
+	UpdateHeight int64   `json:"update_height"`
+	Accum        sdk.Dec `json:"accum"`
+}
+
+func NewTotalAccum(height int64) TotalAccum {
+	return TotalAccum{
+		UpdateHeight: height,
+		Accum:        sdk.ZeroDec(),
+	}
+}
+
+// update total accumulation factor for the new height
+// CONTRACT: height should be greater than the old height
+func (ta TotalAccum) UpdateForNewHeight(height int64, accumCreatedPerBlock sdk.Dec) TotalAccum {
+	blocks := height - ta.UpdateHeight
+	if blocks < 0 {
+		panic("reverse updated for new height")
+	}
+	ta.Accum = ta.Accum.Add(accumCreatedPerBlock.MulInt(sdk.NewInt(blocks)))
+	ta.UpdateHeight = height
+	return ta
+}
+
+// get total accumulation factor for the given height
+// CONTRACT: height should be greater than the old height
+func (ta TotalAccum) GetAccum(height int64, accumCreatedPerBlock sdk.Dec) sdk.Dec {
+	blocks := height - ta.UpdateHeight
+	if blocks < 0 {
+		panic("reverse updated for new height")
+	}
+	return ta.Accum.Add(accumCreatedPerBlock.MulInt(sdk.NewInt(blocks)))
+}

x/distribution/types/total_accum_test.go

@@ -0,0 +1,19 @@
+package types
+
+import (
+	"testing"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/stretchr/testify/require"
+)
+
+func TestTotalAccumUpdateForNewHeight(t *testing.T) {
+
+	ta := NewTotalAccum(0)
+
+	ta = ta.UpdateForNewHeight(5, sdk.NewDec(3))
+	require.True(sdk.DecEq(t, sdk.NewDec(15), ta.Accum))
+
+	ta = ta.UpdateForNewHeight(8, sdk.NewDec(2))
+	require.True(sdk.DecEq(t, sdk.NewDec(21), ta.Accum))
+}

x/distribution/types/validator_info.go

@@ -4,6 +4,29 @@ import (
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )
 
+// common parameters used in withdraws from validators
+type WithdrawContext struct {
+	FeePool        FeePool
+	Height         int64   // block height
+	TotalPower     sdk.Dec // total bonded tokens in the network
+	ValPower       sdk.Dec // validator's bonded tokens
+	CommissionRate sdk.Dec // validator commission rate
+}
+
+func NewWithdrawContext(feePool FeePool, height int64, totalPower,
+	valPower, commissionRate sdk.Dec) WithdrawContext {
+
+	return WithdrawContext{
+		FeePool:        feePool,
+		Height:         height,
+		TotalPower:     totalPower,
+		ValPower:       valPower,
+		CommissionRate: commissionRate,
+	}
+}
+
+//_____________________________________________________________________________
+
 // distribution info for a particular validator
 type ValidatorDistInfo struct {
 	OperatorAddr sdk.ValAddress `json:"operator_addr"`
@@ -31,6 +54,17 @@ func (vi ValidatorDistInfo) UpdateTotalDelAccum(height int64, totalDelShares sdk
 	return vi
 }
 
+// Get the total delegator accum within this validator at the provided height
+func (vi ValidatorDistInfo) GetTotalDelAccum(height int64, totalDelShares sdk.Dec) sdk.Dec {
+	return vi.DelAccum.GetAccum(height, totalDelShares)
+}
+
+// Get the validator accum at the provided height
+func (vi ValidatorDistInfo) GetValAccum(height int64, valTokens sdk.Dec) sdk.Dec {
+	blocks := height - vi.FeePoolWithdrawalHeight
+	return valTokens.MulInt(sdk.NewInt(blocks))
+}
+
 // Move any available accumulated fees in the FeePool to the validator's pool
 // - updates validator info's FeePoolWithdrawalHeight, thus setting accum to 0
 // - updates fee pool to latest height and total val accum w/ given totalBonded
@@ -40,20 +74,19 @@ func (vi ValidatorDistInfo) UpdateTotalDelAccum(height int64, totalDelShares sdk
 // - called in DelegationDistInfo.WithdrawRewards
 // NOTE: When a delegator unbonds, say, onDelegationSharesModified ->
 //       WithdrawDelegationReward -> WithdrawRewards
-func (vi ValidatorDistInfo) TakeFeePoolRewards(fp FeePool, height int64, totalBonded, vdTokens,
+func (vi ValidatorDistInfo) TakeFeePoolRewards(wc WithdrawContext) (
-	commissionRate sdk.Dec) (ValidatorDistInfo, FeePool) {
+	ValidatorDistInfo, FeePool) {
 
-	fp = fp.UpdateTotalValAccum(height, totalBonded)
+	fp := wc.FeePool.UpdateTotalValAccum(wc.Height, wc.TotalPower)
 
 	if fp.TotalValAccum.Accum.IsZero() {
-		vi.FeePoolWithdrawalHeight = height
+		vi.FeePoolWithdrawalHeight = wc.Height
 		return vi, fp
 	}
 
 	// update the validators pool
-	blocks := height - vi.FeePoolWithdrawalHeight
+	accum := vi.GetValAccum(wc.Height, wc.ValPower)
-	vi.FeePoolWithdrawalHeight = height
+	vi.FeePoolWithdrawalHeight = wc.Height
-	accum := vdTokens.MulInt(sdk.NewInt(blocks))
 
 	if accum.GT(fp.TotalValAccum.Accum) {
 		panic("individual accum should never be greater than the total")
@@ -61,7 +94,7 @@ func (vi ValidatorDistInfo) TakeFeePoolRewards(fp FeePool, height int64, totalBo
 	withdrawalTokens := fp.ValPool.MulDec(accum).QuoDec(fp.TotalValAccum.Accum)
 	remValPool := fp.ValPool.Minus(withdrawalTokens)
 
-	commission := withdrawalTokens.MulDec(commissionRate)
+	commission := withdrawalTokens.MulDec(wc.CommissionRate)
 	afterCommission := withdrawalTokens.Minus(commission)
 
 	fp.TotalValAccum.Accum = fp.TotalValAccum.Accum.Sub(accum)
@@ -73,13 +106,48 @@ func (vi ValidatorDistInfo) TakeFeePoolRewards(fp FeePool, height int64, totalBo
 }
 
 // withdraw commission rewards
-func (vi ValidatorDistInfo) WithdrawCommission(fp FeePool, height int64,
+func (vi ValidatorDistInfo) WithdrawCommission(wc WithdrawContext) (
-	totalBonded, vdTokens, commissionRate sdk.Dec) (vio ValidatorDistInfo, fpo FeePool, withdrawn DecCoins) {
+	vio ValidatorDistInfo, fpo FeePool, withdrawn DecCoins) {
 
-	vi, fp = vi.TakeFeePoolRewards(fp, height, totalBonded, vdTokens, commissionRate)
+	vi, fp := vi.TakeFeePoolRewards(wc)
 
 	withdrawalTokens := vi.ValCommission
 	vi.ValCommission = DecCoins{} // zero
 
 	return vi, fp, withdrawalTokens
 }
+
+// get the validator's pool rewards at this current state,
+func (vi ValidatorDistInfo) CurrentPoolRewards(
+	wc WithdrawContext) DecCoins {
+
+	fp := wc.FeePool
+	totalValAccum := fp.GetTotalValAccum(wc.Height, wc.TotalPower)
+	valAccum := vi.GetValAccum(wc.Height, wc.ValPower)
+
+	if valAccum.GT(totalValAccum) {
+		panic("individual accum should never be greater than the total")
+	}
+	withdrawalTokens := fp.ValPool.MulDec(valAccum).QuoDec(totalValAccum)
+	commission := withdrawalTokens.MulDec(wc.CommissionRate)
+	afterCommission := withdrawalTokens.Minus(commission)
+	pool := vi.DelPool.Plus(afterCommission)
+	return pool
+}
+
+// get the validator's commission pool rewards at this current state,
+func (vi ValidatorDistInfo) CurrentCommissionRewards(
+	wc WithdrawContext) DecCoins {
+
+	fp := wc.FeePool
+	totalValAccum := fp.GetTotalValAccum(wc.Height, wc.TotalPower)
+	valAccum := vi.GetValAccum(wc.Height, wc.ValPower)
+
+	if valAccum.GT(totalValAccum) {
+		panic("individual accum should never be greater than the total")
+	}
+	withdrawalTokens := fp.ValPool.MulDec(valAccum).QuoDec(totalValAccum)
+	commission := withdrawalTokens.MulDec(wc.CommissionRate)
+	commissionPool := vi.ValCommission.Plus(commission)
+	return commissionPool
+}

x/distribution/types/validator_info_test.go

@@ -28,13 +28,15 @@ func TestTakeFeePoolRewards(t *testing.T) {
 	height = 10
 	fp.ValPool = DecCoins{NewDecCoin("stake", 1000)}
 
-	vi1, fp = vi1.TakeFeePoolRewards(fp, height, totalBondedTokens, validatorTokens1, commissionRate1)
+	vi1, fp = vi1.TakeFeePoolRewards(NewWithdrawContext(
+		fp, height, totalBondedTokens, validatorTokens1, commissionRate1))
 	require.True(sdk.DecEq(t, sdk.NewDec(900), fp.TotalValAccum.Accum))
 	assert.True(sdk.DecEq(t, sdk.NewDec(900), fp.ValPool[0].Amount))
 	assert.True(sdk.DecEq(t, sdk.NewDec(100-2), vi1.DelPool[0].Amount))
 	assert.True(sdk.DecEq(t, sdk.NewDec(2), vi1.ValCommission[0].Amount))
 
-	vi2, fp = vi2.TakeFeePoolRewards(fp, height, totalBondedTokens, validatorTokens2, commissionRate2)
+	vi2, fp = vi2.TakeFeePoolRewards(NewWithdrawContext(
+		fp, height, totalBondedTokens, validatorTokens2, commissionRate2))
 	require.True(sdk.DecEq(t, sdk.NewDec(500), fp.TotalValAccum.Accum))
 	assert.True(sdk.DecEq(t, sdk.NewDec(500), fp.ValPool[0].Amount))
 	assert.True(sdk.DecEq(t, sdk.NewDec(400-12), vi2.DelPool[0].Amount))
@@ -44,7 +46,8 @@ func TestTakeFeePoolRewards(t *testing.T) {
 	height = 20
 	fp.ValPool[0].Amount = fp.ValPool[0].Amount.Add(sdk.NewDec(1000))
 
-	vi3, fp = vi3.TakeFeePoolRewards(fp, height, totalBondedTokens, validatorTokens3, commissionRate3)
+	vi3, fp = vi3.TakeFeePoolRewards(NewWithdrawContext(
+		fp, height, totalBondedTokens, validatorTokens3, commissionRate3))
 	require.True(sdk.DecEq(t, sdk.NewDec(500), fp.TotalValAccum.Accum))
 	assert.True(sdk.DecEq(t, sdk.NewDec(500), fp.ValPool[0].Amount))
 	assert.True(sdk.DecEq(t, sdk.NewDec(1000-40), vi3.DelPool[0].Amount))
@@ -66,7 +69,8 @@ func TestWithdrawCommission(t *testing.T) {
 	fp.ValPool = DecCoins{NewDecCoin("stake", 1000)}
 
 	// for a more fun staring condition, have an non-withdraw update
-	vi, fp = vi.TakeFeePoolRewards(fp, height, totalBondedTokens, validatorTokens, commissionRate)
+	vi, fp = vi.TakeFeePoolRewards(NewWithdrawContext(
+		fp, height, totalBondedTokens, validatorTokens, commissionRate))
 	require.True(sdk.DecEq(t, sdk.NewDec(900), fp.TotalValAccum.Accum))
 	assert.True(sdk.DecEq(t, sdk.NewDec(900), fp.ValPool[0].Amount))
 	assert.True(sdk.DecEq(t, sdk.NewDec(100-2), vi.DelPool[0].Amount))
@@ -76,7 +80,8 @@ func TestWithdrawCommission(t *testing.T) {
 	height = 20
 	fp.ValPool[0].Amount = fp.ValPool[0].Amount.Add(sdk.NewDec(1000))
 
-	vi, fp, commissionRecv := vi.WithdrawCommission(fp, height, totalBondedTokens, validatorTokens, commissionRate)
+	vi, fp, commissionRecv := vi.WithdrawCommission(NewWithdrawContext(
+		fp, height, totalBondedTokens, validatorTokens, commissionRate))
 	require.True(sdk.DecEq(t, sdk.NewDec(1800), fp.TotalValAccum.Accum))
 	assert.True(sdk.DecEq(t, sdk.NewDec(1800), fp.ValPool[0].Amount))
 	assert.True(sdk.DecEq(t, sdk.NewDec(200-4), vi.DelPool[0].Amount))

x/gov/simulation/invariants.go

@@ -3,11 +3,12 @@ package simulation
 import (
 	"github.com/cosmos/cosmos-sdk/baseapp"
 	"github.com/cosmos/cosmos-sdk/x/mock/simulation"
+	abci "github.com/tendermint/tendermint/abci/types"
 )
 
 // AllInvariants tests all governance invariants
 func AllInvariants() simulation.Invariant {
-	return func(app *baseapp.BaseApp) error {
+	return func(app *baseapp.BaseApp, _ abci.Header) error {
 		// TODO Add some invariants!
 		// Checking proposal queues, no passed-but-unexecuted proposals, etc.
 		return nil

x/mock/simulation/random_simulate_blocks.go

@@ -138,7 +138,7 @@ func SimulateFromSeed(tb testing.TB, app *baseapp.BaseApp,
 
 		if testingMode {
 			// Make sure invariants hold at beginning of block
-			assertAllInvariants(t, app, invariants, "BeginBlock", displayLogs)
+			assertAllInvariants(t, app, header, invariants, "BeginBlock", displayLogs)
 		}
 
 		ctx := app.NewContext(false, header)
@@ -150,7 +150,7 @@ func SimulateFromSeed(tb testing.TB, app *baseapp.BaseApp,
 		numQueuedTimeOpsRan := runQueuedTimeOperations(timeOperationQueue, header.Time, tb, r, app, ctx, accs, logWriter, displayLogs, event)
 		if testingMode && onOperation {
 			// Make sure invariants hold at end of queued operations
-			assertAllInvariants(t, app, invariants, "QueuedOperations", displayLogs)
+			assertAllInvariants(t, app, header, invariants, "QueuedOperations", displayLogs)
 		}
 
 		thisBlockSize = thisBlockSize - numQueuedOpsRan - numQueuedTimeOpsRan
@@ -159,7 +159,7 @@ func SimulateFromSeed(tb testing.TB, app *baseapp.BaseApp,
 		opCount += operations + numQueuedOpsRan + numQueuedTimeOpsRan
 		if testingMode {
 			// Make sure invariants hold at end of block
-			assertAllInvariants(t, app, invariants, "StandardOperations", displayLogs)
+			assertAllInvariants(t, app, header, invariants, "StandardOperations", displayLogs)
 		}
 
 		res := app.EndBlock(abci.RequestEndBlock{})
@@ -170,7 +170,7 @@ func SimulateFromSeed(tb testing.TB, app *baseapp.BaseApp,
 
 		if testingMode {
 			// Make sure invariants hold at end of block
-			assertAllInvariants(t, app, invariants, "EndBlock", displayLogs)
+			assertAllInvariants(t, app, header, invariants, "EndBlock", displayLogs)
 		}
 		if commit {
 			app.Commit()
@@ -230,7 +230,7 @@ func createBlockSimulator(testingMode bool, tb testing.TB, t *testing.T, event f
 			queueOperations(operationQueue, timeOperationQueue, futureOps)
 			if testingMode {
 				if onOperation {
-					assertAllInvariants(t, app, invariants, fmt.Sprintf("operation: %v", logUpdate), displayLogs)
+					assertAllInvariants(t, app, header, invariants, fmt.Sprintf("operation: %v", logUpdate), displayLogs)
 				}
 				if opCount%50 == 0 {
 					fmt.Printf("\rSimulating... block %d/%d, operation %d/%d. ", header.Height, totalNumBlocks, opCount, blocksize)

x/mock/simulation/types.go

@@ -33,7 +33,7 @@ type (
 	// If the invariant has been broken, it should return an error
 	// containing a descriptive message about what happened.
 	// The simulator will then halt and print the logs.
-	Invariant func(app *baseapp.BaseApp) error
+	Invariant func(app *baseapp.BaseApp, header abci.Header) error
 
 	// Account contains a privkey, pubkey, address tuple
 	// eventually more useful data can be placed in here.
@@ -68,13 +68,14 @@ type (
 	}
 )
 
+// TODO remove? not being called anywhere
 // PeriodicInvariant returns an Invariant function closure that asserts
 // a given invariant if the mock application's last block modulo the given
 // period is congruent to the given offset.
 func PeriodicInvariant(invariant Invariant, period int, offset int) Invariant {
-	return func(app *baseapp.BaseApp) error {
+	return func(app *baseapp.BaseApp, header abci.Header) error {
 		if int(app.LastBlockHeight())%period == offset {
-			return invariant(app)
+			return invariant(app, header)
 		}
 		return nil
 	}

x/mock/simulation/util.go

@@ -9,6 +9,7 @@ import (
 	"testing"
 	"time"
 
+	abci "github.com/tendermint/tendermint/abci/types"
 	"github.com/tendermint/tendermint/crypto/ed25519"
 	"github.com/tendermint/tendermint/crypto/secp256k1"
 
@@ -102,9 +103,11 @@ func addLogMessage(testingmode bool, blockLogBuilders []*strings.Builder, height
 }
 
 // assertAllInvariants asserts a list of provided invariants against application state
-func assertAllInvariants(t *testing.T, app *baseapp.BaseApp, invariants []Invariant, where string, displayLogs func()) {
+func assertAllInvariants(t *testing.T, app *baseapp.BaseApp, header abci.Header,
+	invariants []Invariant, where string, displayLogs func()) {
+
 	for i := 0; i < len(invariants); i++ {
-		err := invariants[i](app)
+		err := invariants[i](app, header)
 		if err != nil {
 			fmt.Printf("Invariants broken after %s\n", where)
 			fmt.Println(err.Error())

x/slashing/simulation/invariants.go

@@ -3,12 +3,13 @@ package simulation
 import (
 	"github.com/cosmos/cosmos-sdk/baseapp"
 	"github.com/cosmos/cosmos-sdk/x/mock/simulation"
+	abci "github.com/tendermint/tendermint/abci/types"
 )
 
+// TODO Any invariants to check here?
 // AllInvariants tests all slashing invariants
 func AllInvariants() simulation.Invariant {
-	return func(app *baseapp.BaseApp) error {
+	return func(_ *baseapp.BaseApp, _ abci.Header) error {
-		// TODO Any invariants to check here?
 		return nil
 	}
 }

x/stake/simulation/invariants.go

@@ -15,25 +15,29 @@ import (
 
 // AllInvariants runs all invariants of the stake module.
 // Currently: total supply, positive power
-func AllInvariants(ck bank.Keeper, k stake.Keeper, f auth.FeeCollectionKeeper, d distribution.Keeper, am auth.AccountKeeper) simulation.Invariant {
+func AllInvariants(ck bank.Keeper, k stake.Keeper,
-	return func(app *baseapp.BaseApp) error {
+	f auth.FeeCollectionKeeper, d distribution.Keeper,
-		err := SupplyInvariants(ck, k, f, d, am)(app)
+	am auth.AccountKeeper) simulation.Invariant {
+
+	return func(app *baseapp.BaseApp, header abci.Header) error {
+		err := SupplyInvariants(ck, k, f, d, am)(app, header)
 		if err != nil {
 			return err
 		}
-		err = PositivePowerInvariant(k)(app)
+		err = PositivePowerInvariant(k)(app, header)
 		if err != nil {
 			return err
 		}
-		err = ValidatorSetInvariant(k)(app)
+		err = ValidatorSetInvariant(k)(app, header)
 		return err
 	}
 }
 
 // SupplyInvariants checks that the total supply reflects all held loose tokens, bonded tokens, and unbonding delegations
 // nolint: unparam
-func SupplyInvariants(ck bank.Keeper, k stake.Keeper, f auth.FeeCollectionKeeper, d distribution.Keeper, am auth.AccountKeeper) simulation.Invariant {
+func SupplyInvariants(ck bank.Keeper, k stake.Keeper,
-	return func(app *baseapp.BaseApp) error {
+	f auth.FeeCollectionKeeper, d distribution.Keeper, am auth.AccountKeeper) simulation.Invariant {
+	return func(app *baseapp.BaseApp, _ abci.Header) error {
 		ctx := app.NewContext(false, abci.Header{})
 		pool := k.GetPool(ctx)
 
@@ -71,20 +75,25 @@ func SupplyInvariants(ck bank.Keeper, k stake.Keeper, f auth.FeeCollectionKeeper
 		loose = loose.Add(feePool.ValPool.AmountOf("steak"))
 
 		// add validator distribution commission and yet-to-be-withdrawn-by-delegators
-		d.IterateValidatorDistInfos(ctx, func(_ int64, distInfo distribution.ValidatorDistInfo) (stop bool) {
+		d.IterateValidatorDistInfos(ctx,
-			loose = loose.Add(distInfo.ValCommission.AmountOf("steak"))
+			func(_ int64, distInfo distribution.ValidatorDistInfo) (stop bool) {
-			loose = loose.Add(distInfo.DelPool.AmountOf("steak"))
+				loose = loose.Add(distInfo.DelPool.AmountOf("steak"))
-			return false
+				loose = loose.Add(distInfo.ValCommission.AmountOf("steak"))
-		})
+				return false
+			},
+		)
 
-		// Loose tokens should equal coin supply plus unbonding delegations plus tokens on unbonded validators
+		// Loose tokens should equal coin supply plus unbonding delegations
+		// plus tokens on unbonded validators
 		if !pool.LooseTokens.Equal(loose) {
-			return fmt.Errorf("expected loose tokens to equal total steak held by accounts - pool.LooseTokens: %v, sum of account tokens: %v", pool.LooseTokens, loose)
+			return fmt.Errorf("loose token invariance:\n\tpool.LooseTokens: %v"+
+				"\n\tsum of account tokens: %v", pool.LooseTokens, loose)
 		}
 
 		// Bonded tokens should equal sum of tokens with bonded validators
 		if !pool.BondedTokens.Equal(bonded) {
-			return fmt.Errorf("expected bonded tokens to equal total steak held by bonded validators - pool.BondedTokens: %v, sum of bonded validator tokens: %v", pool.BondedTokens, bonded)
+			return fmt.Errorf("bonded token invariance:\n\tpool.BondedTokens: %v"+
+				"\n\tsum of account tokens: %v", pool.BondedTokens, bonded)
 		}
 
 		return nil
@@ -93,7 +102,7 @@ func SupplyInvariants(ck bank.Keeper, k stake.Keeper, f auth.FeeCollectionKeeper
 
 // PositivePowerInvariant checks that all stored validators have > 0 power
 func PositivePowerInvariant(k stake.Keeper) simulation.Invariant {
-	return func(app *baseapp.BaseApp) error {
+	return func(app *baseapp.BaseApp, _ abci.Header) error {
 		ctx := app.NewContext(false, abci.Header{})
 		var err error
 		k.IterateValidatorsBonded(ctx, func(_ int64, validator sdk.Validator) bool {
@@ -109,7 +118,7 @@ func PositivePowerInvariant(k stake.Keeper) simulation.Invariant {
 
 // ValidatorSetInvariant checks equivalence of Tendermint validator set and SDK validator set
 func ValidatorSetInvariant(k stake.Keeper) simulation.Invariant {
-	return func(app *baseapp.BaseApp) error {
+	return func(app *baseapp.BaseApp, _ abci.Header) error {
 		// TODO
 		return nil
 	}