WIP: distribution refactor

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

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, 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 := 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)
-	valInfo, feePool, commission := valInfo.WithdrawCommission(feePool, height, lastTotalPower,
-		lastValPower, validator.GetCommission())
+	wc := k.GetWithdrawContext(ctx, operatorAddr)
+	valInfo, feePool, commission := valInfo.WithdrawCommission(wc)
 	withdraw = withdraw.Plus(commission)
 	k.SetValidatorDistInfo(ctx, valInfo)
 
-	withdrawAddr := k.GetDelegatorWithdrawAddr(ctx, accAddr)
-	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)
-	}
-
+	k.WithdrawToDelegator(ctx, feePool, accAddr, withdraw)
 	return nil
 }
 
-// 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++
+// get all the validator rewards including the commission
+func (k Keeper) CurrentValidatorRewardsAll(ctx sdk.Context, operatorAddr sdk.ValAddress) (sdk.Coins, sdk.Error) {
+
+	if !k.HasValidatorDistInfo(ctx, operatorAddr) {
+		return sdk.Coins{}, types.ErrNoValidatorDistInfo(k.codespace)
 	}
+
+	// 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/types/delegator_info.go

@@ -21,28 +21,34 @@ 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.WithdrawalHeight
+	return delegatorShares.MulInt(sdk.NewInt(blocks))
+}
+
 // Withdraw rewards from delegator.
 // Among many things, it does:
 // * updates validator info's total del accum
 // * calls vi.TakeFeePoolRewards, which:
 //   * 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,
-	height int64, totalBonded, vdTokens, totalDelShares, delegatorShares,
-	commissionRate sdk.Dec) (DelegationDistInfo, ValidatorDistInfo, FeePool, DecCoins) {
+func (di DelegationDistInfo) WithdrawRewards(wc WithdrawContext, vi ValidatorDistInfo,
+	totalDelShares, delegatorShares sdk.Dec) (
+	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.WithdrawalHeight
-	di.WithdrawalHeight = height
-	accum := delegatorShares.MulInt(sdk.NewInt(blocks))
+	accum := di.GetDelAccum(wc.Height, delegatorShares)
+	di.WithdrawalHeight = wc.Height
 	withdrawalTokens := vi.Pool.MulDec(accum).QuoDec(vi.DelAccum.Accum)
 	remainingTokens := vi.Pool.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
+}