Merge pull request #40 from initia-labs/fix/distribution-iterator

fix: distribution iterator to do single iter

x/distribution/keeper/delegation.go

@@ -128,35 +128,32 @@ func (k Keeper) CalculateDelegationRewards(ctx context.Context, val stakingtypes
 	// for them for the stake sanity check below.
 	endingHeight := uint64(sdkCtx.BlockHeight())
 	if endingHeight > startingHeight {
+		err = k.ValidatorSlashEvents.Walk(ctx, new(collections.Range[collections.Triple[[]byte, uint64, uint64]]).
+			StartInclusive(collections.Join3[[]byte, uint64, uint64](valAddr, startingHeight, 0)).
+			EndExclusive(collections.Join3[[]byte, uint64, uint64](valAddr, endingHeight+1, 0)),
+			func(key collections.Triple[[]byte, uint64, uint64], event customtypes.ValidatorSlashEvent) (stop bool, err error) {
+				endingPeriod := event.ValidatorPeriod
+				if endingPeriod > startingPeriod {
+					rewardsBetween, err := k.calculateDelegationRewardsBetween(ctx, val, startingPeriod, endingPeriod, stakes)
+					if err != nil {
+						return false, err
+					}
+
+					rewards = rewards.Add(rewardsBetween...)
 
-		// TODO - change height iteration to raw iteration
-		// k.ValidatorSlashEvents.IterateRaw()
-		for height := startingHeight; height < endingHeight+1; height++ {
-			err = k.ValidatorSlashEvents.Walk(ctx, collections.NewSuperPrefixedTripleRange[[]byte, uint64, uint64](valAddr, height),
-				func(key collections.Triple[[]byte, uint64, uint64], event customtypes.ValidatorSlashEvent) (stop bool, err error) {
-					endingPeriod := event.ValidatorPeriod
-					if endingPeriod > startingPeriod {
-						rewardsBetween, err := k.calculateDelegationRewardsBetween(ctx, val, startingPeriod, endingPeriod, stakes)
-						if err != nil {
-							return false, err
-						}
-
-						rewards = rewards.Add(rewardsBetween...)
-
-						// Note: It is necessary to truncate so we don't allow withdrawing
-						// more rewards than owed.
-						for i, stake := range stakes {
-							stakes[i].Amount = stake.Amount.MulTruncate(math.LegacyOneDec().Sub(event.Fractions.AmountOf(stake.Denom)))
-						}
-						startingPeriod = endingPeriod
+					// Note: It is necessary to truncate so we don't allow withdrawing
+					// more rewards than owed.
+					for i, stake := range stakes {
+						stakes[i].Amount = stake.Amount.MulTruncate(math.LegacyOneDec().Sub(event.Fractions.AmountOf(stake.Denom)))
 					}
+					startingPeriod = endingPeriod
+				}
 
-					return false, nil
-				},
-			)
-			if err != nil {
-				return
-			}
+				return false, nil
+			},
+		)
+		if err != nil {
+			return
 		}
 	}
 

x/distribution/keeper/store.go

@@ -47,23 +47,21 @@ func (k Keeper) GetValidatorSlashEvent(ctx context.Context, val sdk.ValAddress,
 }
 
 // get validator outstanding rewards
+// ValidatorOutstandingRewards should be initialized at initValidator()
 func (k Keeper) GetValidatorOutstandingRewards(ctx context.Context, val sdk.ValAddress) (rewards types.ValidatorOutstandingRewards, err error) {
 	rewards, err = k.ValidatorOutstandingRewards.Get(ctx, val)
-	if err != nil && errors.Is(err, collections.ErrNotFound) {
-		return types.ValidatorOutstandingRewards{}, nil
-	} else if err != nil {
+	if err != nil {
 		return types.ValidatorOutstandingRewards{}, err
 	}
 
 	return
 }
 
 // get current rewards for a validator
+// ValidatorCurrentRewards rewards should be initialized at initValidator()
 func (k Keeper) GetValidatorCurrentRewards(ctx context.Context, val sdk.ValAddress) (rewards types.ValidatorCurrentRewards, err error) {
 	rewards, err = k.ValidatorCurrentRewards.Get(ctx, val)
-	if err != nil && errors.Is(err, collections.ErrNotFound) {
-		return types.ValidatorCurrentRewards{}, nil
-	} else if err != nil {
+	if err != nil {
 		return types.ValidatorCurrentRewards{}, err
 	}
 

x/distribution/keeper/validator.go

@@ -163,7 +163,7 @@ func (k Keeper) decrementReferenceCount(ctx context.Context, valAddr sdk.ValAddr
 func (k Keeper) updateValidatorSlashFraction(ctx context.Context, valAddr sdk.ValAddress, fractions sdk.DecCoins) error {
 	for _, fraction := range fractions {
 		if fraction.Amount.GT(math.LegacyOneDec()) || fraction.Amount.IsNegative() {
-			panic(fmt.Sprintf("fraction must be >=0 and <=1, current fraction: %v", fraction))
+			return fmt.Errorf("fraction must be >=0 and <=1, current fraction: %v", fraction)
 		}
 	}