R4R: distribution accum invariants

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,111 @@ func (k Keeper) RemoveDelegatorWithdrawAddr(ctx sdk.Context, delAddr, withdrawAd
 
 //___________________________________________________________________________________________
 
-// Withdraw all the 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,
-	valAddr sdk.ValAddress) sdk.Error {
+// 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) {
 
-	if !k.HasDelegationDistInfo(ctx, delegatorAddr, valAddr) {
-		return types.ErrNoDelegationDistInfo(k.codespace)
-	}
+	wc := k.GetWithdrawContext(ctx, valAddr)
 
-	// TODO: Reconcile with duplicate code in getDelegatorRewardsAll.
-	height := ctx.BlockHeight()
-	lastTotalPower := sdk.NewDecFromInt(k.stakeKeeper.GetLastTotalPower(ctx))
-	lastValPower := k.stakeKeeper.GetLastValidatorPower(ctx, valAddr)
-	feePool := k.GetFeePool(ctx)
-	delInfo := k.GetDelegationDistInfo(ctx, delegatorAddr, valAddr)
 	valInfo := k.GetValidatorDistInfo(ctx, valAddr)
+	delInfo := k.GetDelegationDistInfo(ctx, delAddr, valAddr)
 	validator := k.stakeKeeper.Validator(ctx, valAddr)
-	delegation := k.stakeKeeper.Delegation(ctx, delegatorAddr, valAddr)
+	delegation := k.stakeKeeper.Delegation(ctx, delAddr, valAddr)
 
-	delInfo, valInfo, feePool, withdraw := delInfo.WithdrawRewards(feePool, valInfo, height, lastTotalPower,
-		lastValPower, validator.GetDelegatorShares(), delegation.GetShares(), validator.GetCommission())
+	delInfo, valInfo, feePool, withdraw := delInfo.WithdrawRewards(wc, valInfo,
+		validator.GetDelegatorShares(), delegation.GetShares())
 
-	k.SetValidatorDistInfo(ctx, valInfo)
-	k.SetDelegationDistInfo(ctx, delInfo)
-	withdrawAddr := k.GetDelegatorWithdrawAddr(ctx, delegatorAddr)
-	coinsToAdd, change := withdraw.TruncateDecimal()
+	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) WithdrawToDelegator(ctx sdk.Context, feePool types.FeePool,
+	delAddr sdk.AccAddress, amount types.DecCoins) {
+
+	withdrawAddr := k.GetDelegatorWithdrawAddr(ctx, delAddr)
+	coinsToAdd, change := amount.TruncateDecimal()
 	feePool.CommunityPool = feePool.CommunityPool.Plus(change)
 	k.SetFeePool(ctx, feePool)
 	_, _, err := k.bankKeeper.AddCoins(ctx, withdrawAddr, coinsToAdd)
 	if err != nil {
 		panic(err)
 	}
-	return nil
 }
 
 //___________________________________________________________________________________________
 
-// return all rewards for all delegations of a delegator
-func (k Keeper) WithdrawDelegationRewardsAll(ctx sdk.Context, delegatorAddr sdk.AccAddress) {
-	height := ctx.BlockHeight()
-	withdraw := k.getDelegatorRewardsAll(ctx, delegatorAddr, height)
-	feePool := k.GetFeePool(ctx)
-	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) getDelegatorRewardsAll(ctx sdk.Context, delAddr sdk.AccAddress, height int64) types.DecCoins {
+func (k Keeper) WithdrawDelegationRewardsAll(ctx sdk.Context, delAddr sdk.AccAddress) {
+	withdraw := k.withdrawDelegationRewardsAll(ctx, delAddr)
+	feePool := k.GetFeePool(ctx)
+	k.WithdrawToDelegator(ctx, feePool, delAddr, withdraw)
+}
 
-	withdraw := types.DecCoins{}
-	lastTotalPower := sdk.NewDecFromInt(k.stakeKeeper.GetLastTotalPower(ctx))
+func (k Keeper) withdrawDelegationRewardsAll(ctx sdk.Context,
+	delAddr sdk.AccAddress) types.DecCoins {
 
 	// 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 := 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,
-			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 +183,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

@@ -14,13 +14,20 @@ func (k Keeper) onValidatorCreated(ctx sdk.Context, addr sdk.ValAddress) {
 	vdi := types.ValidatorDistInfo{
 		OperatorAddr:            addr,
 		FeePoolWithdrawalHeight: height,
-		Pool:                    types.DecCoins{},
-		PoolCommission:          types.DecCoins{},
-		DelAccum:                types.NewTotalAccum(height),
+		Pool:           types.DecCoins{},
+		PoolCommission: types.DecCoins{},
+		DelAccum:       types.NewTotalAccum(height),
 	}
 	k.SetValidatorDistInfo(ctx, vdi)
 }
 
+// Reset the height to the current height when bonded
+func (k Keeper) onValidatorBonded(ctx sdk.Context, addr sdk.ValAddress) {
+	vdi := k.GetValidatorDistInfo(ctx, addr)
+	vdi.FeePoolWithdrawalHeight = ctx.BlockHeight()
+	k.SetValidatorDistInfo(ctx, vdi)
+}
+
 // Withdrawal all validator rewards
 func (k Keeper) onValidatorModified(ctx sdk.Context, addr sdk.ValAddress) {
 	// This doesn't need to be run at genesis
@@ -104,5 +111,5 @@ func (h Hooks) OnValidatorBeginUnbonding(ctx sdk.Context, _ sdk.ConsAddress, add
 	h.k.onValidatorModified(ctx, addr)
 }
 func (h Hooks) OnValidatorBonded(ctx sdk.Context, _ sdk.ConsAddress, addr sdk.ValAddress) {
-	h.k.onValidatorModified(ctx, addr)
+	h.k.onValidatorBonded(ctx, addr)
 }

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, 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++
+	}
+}