fix: refactor redelegations distribution logic

x/interchainstaking/keeper/delegation.go

@@ -202,7 +202,7 @@ func (k *Keeper) PrepareDelegationMessagesForShares(zone *types.Zone, coins sdk.
 }
 
 func (k *Keeper) DeterminePlanForDelegation(ctx sdk.Context, zone *types.Zone, amount sdk.Coins) (map[string]sdkmath.Int, error) {
-	currentAllocations, currentSum, _ := k.GetDelegationMap(ctx, zone)
+	currentAllocations, currentSum, _, _ := k.GetDelegationMap(ctx, zone)
 	targetAllocations, err := k.GetAggregateIntentOrDefault(ctx, zone)
 	if err != nil {
 		return nil, err
@@ -248,25 +248,23 @@ func (k *Keeper) WithdrawDelegationRewardsForResponse(ctx sdk.Context, zone *typ
 	return k.SubmitTx(ctx, msgs, zone.DelegationAddress, "", zone.MessagesPerTx)
 }
 
-func (k *Keeper) GetDelegationMap(ctx sdk.Context, zone *types.Zone) (out map[string]sdkmath.Int, sum sdkmath.Int, locked map[string]bool) {
+func (k *Keeper) GetDelegationMap(ctx sdk.Context, zone *types.Zone) (out map[string]sdkmath.Int, sum sdkmath.Int, locked map[string]bool, lockedSum sdkmath.Int) {
 	out = make(map[string]sdkmath.Int)
 	locked = make(map[string]bool)
 	sum = sdk.ZeroInt()
+	lockedSum = sdk.ZeroInt()
 
 	k.IterateAllDelegations(ctx, zone, func(delegation types.Delegation) bool {
-		existing, found := out[delegation.ValidatorAddress]
-		if !found {
-			out[delegation.ValidatorAddress] = delegation.Amount.Amount
-			locked[delegation.ValidatorAddress] = delegation.RedelegationEnd != 0 && delegation.RedelegationEnd >= ctx.BlockTime().Unix()
-		} else {
-			out[delegation.ValidatorAddress] = existing.Add(delegation.Amount.Amount)
-			locked[delegation.ValidatorAddress] = locked[delegation.ValidatorAddress] || (delegation.RedelegationEnd != 0 && delegation.RedelegationEnd >= ctx.BlockTime().Unix())
+		out[delegation.ValidatorAddress] = delegation.Amount.Amount
+		if delegation.RedelegationEnd >= ctx.BlockTime().Unix() {
+			locked[delegation.ValidatorAddress] = true
+			lockedSum = lockedSum.Add(delegation.Amount.Amount)
 		}
 		sum = sum.Add(delegation.Amount.Amount)
 		return false
 	})
 
-	return out, sum, locked
+	return out, sum, locked, lockedSum
 }
 
 func (k *Keeper) MakePerformanceDelegation(ctx sdk.Context, zone *types.Zone, validator string) error {

x/interchainstaking/keeper/keeper.go

@@ -630,12 +630,12 @@ func (k *Keeper) GetAggregateIntentOrDefault(ctx sdk.Context, z *types.Zone) (ty
 }
 
 func (k *Keeper) Rebalance(ctx sdk.Context, zone *types.Zone, epochNumber int64) error {
-	currentAllocations, currentSum, currentLocked := k.GetDelegationMap(ctx, zone)
+	currentAllocations, currentSum, currentLocked, lockedSum := k.GetDelegationMap(ctx, zone)
 	targetAllocations, err := k.GetAggregateIntentOrDefault(ctx, zone)
 	if err != nil {
 		return err
 	}
-	rebalances := types.DetermineAllocationsForRebalancing(currentAllocations, currentLocked, currentSum, targetAllocations, k.ZoneRedelegationRecords(ctx, zone.ChainId), k.Logger(ctx))
+	rebalances := types.DetermineAllocationsForRebalancing(currentAllocations, currentLocked, currentSum, lockedSum, targetAllocations, k.Logger(ctx))
 	msgs := make([]sdk.Msg, 0)
 	for _, rebalance := range rebalances {
 		msgs = append(msgs, &stakingtypes.MsgBeginRedelegate{DelegatorAddress: zone.DelegationAddress.Address, ValidatorSrcAddress: rebalance.Source, ValidatorDstAddress: rebalance.Target, Amount: sdk.NewCoin(zone.BaseDenom, rebalance.Amount)})

x/interchainstaking/keeper/redemptions.go

@@ -272,7 +272,7 @@ func (k *Keeper) GCCompletedUnbondings(ctx sdk.Context, zone *types.Zone) error
 }
 
 func (k *Keeper) DeterminePlanForUndelegation(ctx sdk.Context, zone *types.Zone, amount sdk.Coins) (map[string]math.Int, error) {
-	currentAllocations, currentSum, _ := k.GetDelegationMap(ctx, zone)
+	currentAllocations, currentSum, _, _ := k.GetDelegationMap(ctx, zone)
 	availablePerValidator, _, err := k.GetUnlockedTokensForZone(ctx, zone)
 	if err != nil {
 		return nil, err

x/interchainstaking/types/rebalance.go

@@ -57,140 +57,187 @@ func CalculateDeltas(currentAllocations map[string]sdkmath.Int, currentSum sdkma
 	return deltas
 }
 
-type RebalanceTarget struct {
-	Amount sdkmath.Int
-	Source string
-	Target string
-}
+// CalculateDeltas determines, for the current delegations, in delta between actual allocations and the target intent.
+// Positive delta represents current allocation is below target, and vice versa.
+func CalculateDeltasNew(currentAllocations map[string]sdkmath.Int, locked map[string]bool, currentSum sdkmath.Int, targetAllocations ValidatorIntents) (targets, sources Deltas) {
+	targets = make(Deltas, 0)
+	sources = make(Deltas, 0)
 
-func DetermineAllocationsForRebalancing(
-	currentAllocations map[string]sdkmath.Int,
-	currentLocked map[string]bool,
-	currentSum sdkmath.Int,
-	targetAllocations ValidatorIntents,
-	existingRedelegations []RedelegationRecord,
-	logger log.Logger,
-) []RebalanceTarget {
-	out := make([]RebalanceTarget, 0)
-	deltas := CalculateDeltas(currentAllocations, currentSum, targetAllocations)
-
-	wantToRebalance := sdk.ZeroInt()
-	canRebalanceFrom := sdk.ZeroInt()
-
-	totalLocked := int64(0)
-	lockedPerValidator := map[string]int64{}
-	for _, redelegation := range existingRedelegations {
-		totalLocked += redelegation.Amount
-		thisLocked, found := lockedPerValidator[redelegation.Destination]
-		if !found {
-			thisLocked = 0
+	targetValopers := func(in ValidatorIntents) []string {
+		out := make([]string, 0, len(in))
+		for _, i := range in {
+			out = append(out, i.ValoperAddress)
 		}
-		lockedPerValidator[redelegation.Destination] = thisLocked + redelegation.Amount
-	}
-	for _, valoper := range utils.Keys(currentAllocations) {
-		// if validator already has a redelegation _to_ it, we can no longer redelegate _from_ it (transitive redelegations)
-		// remove _locked_ amount from lpv and total locked for purposes of rebalancing.
-		if currentLocked[valoper] {
-			thisLocked, found := lockedPerValidator[valoper]
-			if !found {
-				thisLocked = 0
-			}
-			totalLocked = totalLocked - thisLocked + currentAllocations[valoper].Int64()
-			lockedPerValidator[valoper] = currentAllocations[valoper].Int64()
+		return out
+	}(targetAllocations)
+
+	keySet := utils.Unique(append(targetValopers, utils.Keys(currentAllocations)...))
+	sort.Strings(keySet)
+	// for target allocations, raise the intent weight by the total delegated value to get target amount
+	for _, valoper := range keySet {
+		current, ok := currentAllocations[valoper]
+		if !ok {
+			current = sdk.ZeroInt()
 		}
-	}
 
-	// TODO: make these params
-	maxCanRebalanceTotal := currentSum.Sub(sdkmath.NewInt(totalLocked)).Quo(sdk.NewInt(2))
-	maxCanRebalance := sdkmath.MinInt(maxCanRebalanceTotal, currentSum.Quo(sdk.NewInt(7)))
-	if logger != nil {
-		logger.Debug("Rebalancing", "totalLocked", totalLocked, "lockedPerValidator", lockedPerValidator, "canRebalanceTotal", maxCanRebalanceTotal, "canRebalanceEpoch", maxCanRebalance)
-	}
+		target, ok := targetAllocations.GetForValoper(valoper)
+		if !ok {
+			target = &ValidatorIntent{ValoperAddress: valoper, Weight: sdk.ZeroDec()}
+		}
+		targetAmount := target.Weight.MulInt(currentSum).TruncateInt()
+		// diff between target and current allocations
+		// positive == below target, negative == above target
+		delta := targetAmount.Sub(current)
 
-	// deltas are sorted in CalculateDeltas; don't re-sort.
-	for _, delta := range deltas {
-		switch {
-		case delta.Weight.IsZero():
-			// do nothing
-		case delta.Weight.IsPositive():
-			// if delta > current value - locked value, truncate, as we cannot rebalance locked tokens.
-			wantToRebalance = wantToRebalance.Add(delta.Weight.TruncateInt())
-		case delta.Weight.IsNegative():
-			if delta.Weight.Abs().GT(sdk.NewDecFromInt(currentAllocations[delta.ValoperAddress].Sub(sdkmath.NewInt(lockedPerValidator[delta.ValoperAddress])))) {
-				delta.Weight = sdk.NewDecFromInt(currentAllocations[delta.ValoperAddress].Sub(sdkmath.NewInt(lockedPerValidator[delta.ValoperAddress]))).Neg()
-				if logger != nil {
-					logger.Debug("Truncated delta due to locked tokens", "valoper", delta.ValoperAddress, "delta", delta.Weight.Abs())
-				}
+		if delta.IsPositive() {
+			targets = append(targets, &Delta{Amount: delta, ValoperAddress: valoper})
+		} else {
+			if _, found := locked[valoper]; !found {
+				// only append to sources if the delegation is not locked - i.e. it doesn't have an incoming redelegation.
+				sources = append(sources, &Delta{Amount: delta.Abs(), ValoperAddress: valoper})
 			}
-			canRebalanceFrom = canRebalanceFrom.Add(delta.Weight.Abs().TruncateInt())
 		}
 	}
 
-	toRebalance := sdk.MinInt(sdk.MinInt(wantToRebalance, canRebalanceFrom), maxCanRebalance)
+	targets.Sort()
+	sources.Sort()
 
-	if toRebalance.Equal(sdkmath.ZeroInt()) {
-		if logger != nil {
-			logger.Debug("No rebalancing this epoch")
+	return targets, sources
+}
+
+type Delta struct {
+	ValoperAddress string
+	Amount         sdkmath.Int
+}
+
+type Deltas []*Delta
+
+func (d Deltas) Sort() {
+	// filter zeros
+	newDeltas := make(Deltas, 0)
+	for _, delta := range d {
+		if !delta.Amount.IsZero() {
+			newDeltas = append(newDeltas, delta)
 		}
-		return []RebalanceTarget{}
-	}
-	if logger != nil {
-		logger.Debug("Will rebalance this epoch", "amount", toRebalance)
 	}
+	d = newDeltas
 
-	tgtIdx := 0
-	srcIdx := len(deltas) - 1
-	for i := 0; toRebalance.GT(sdk.ZeroInt()); {
-		i++
-		if i > 20 {
-			break
-		}
-		src := deltas[srcIdx]
-		tgt := deltas[tgtIdx]
-		if src.ValoperAddress == tgt.ValoperAddress {
-			break
-		}
-		var amount sdkmath.Int
-		if src.Weight.Abs().TruncateInt().IsZero() { //nolint:gocritic
-			srcIdx--
-			continue
-		} else if src.Weight.Abs().TruncateInt().GT(toRebalance) { // amount == rebalance
-			amount = toRebalance
-		} else {
-			amount = src.Weight.Abs().TruncateInt()
-		}
+	// sort keys by relative value of delta
+	sort.SliceStable(d, func(i, j int) bool {
+		// < sorts alphabetically.
+		return d[i].ValoperAddress < d[j].ValoperAddress
+	})
 
-		if tgt.Weight.Abs().TruncateInt().IsZero() {
-			tgtIdx++
-			continue
-		} else if tgt.Weight.Abs().TruncateInt().LTE(toRebalance) {
-			amount = sdk.MinInt(amount, tgt.Weight.Abs().TruncateInt())
-		}
+	// sort keys by relative value of delta
+	sort.SliceStable(d, func(i, j int) bool {
+		return d[i].Amount.GT(d[j].Amount)
+	})
+}
 
-		out = append(out, RebalanceTarget{Amount: amount, Target: tgt.ValoperAddress, Source: src.ValoperAddress})
-		deltas[srcIdx].Weight = src.Weight.Add(sdk.NewDecFromInt(amount))
-		deltas[tgtIdx].Weight = tgt.Weight.Sub(sdk.NewDecFromInt(amount))
-		toRebalance = toRebalance.Sub(amount)
+type RebalanceTarget struct {
+	Amount sdkmath.Int
+	Source string
+	Target string
+}
 
-	}
+type RebalanceTargets []*RebalanceTarget
 
+func (t RebalanceTargets) Sort() {
 	// sort keys by relative value of delta
-	sort.SliceStable(out, func(i, j int) bool {
-		return out[i].Source < out[j].Source
+	sort.SliceStable(t, func(i, j int) bool {
+		// < sorts alphabetically.
+		return t[i].Source < t[j].Source
 	})
 
-	sort.SliceStable(out, func(i, j int) bool {
-		return out[i].Target < out[j].Target
+	// sort keys by relative value of delta
+	sort.SliceStable(t, func(i, j int) bool {
+		// < sorts alphabetically.
+		return t[i].Target < t[j].Target
 	})
 
 	// sort keys by relative value of delta
-	sort.SliceStable(out, func(i, j int) bool {
-		return out[i].Amount.GT(out[j].Amount)
+	sort.SliceStable(t, func(i, j int) bool {
+		return t[i].Amount.LT(t[j].Amount)
 	})
+}
+
+// DetermineAllocationsForRebalancing takes maps of current and locked delegations, and based upon the target allocations,
+// attempts to satisfy the target allocations in the fewest number of transformations. It returns a slice of RebalanceTargets.
+func DetermineAllocationsForRebalancing(
+	currentAllocations map[string]sdkmath.Int,
+	currentLocked map[string]bool,
+	currentSum sdkmath.Int,
+	lockedSum sdkmath.Int,
+	targetAllocations ValidatorIntents,
+	logger log.Logger,
+) RebalanceTargets {
+	out := make(RebalanceTargets, 0)
+	targets, sources := CalculateDeltasNew(currentAllocations, currentLocked, currentSum, targetAllocations)
+
+	// rebalanceBudget = (total_delegations - locked)/2 == 50% of (total_delegations - locked)
+	// TODO: make this 2 (max_redelegation_factor) a param.
+	rebalanceBudget := currentSum.Sub(lockedSum).Quo(sdk.NewInt(2))
+
+	if logger != nil {
+		logger.Debug("Rebalancing", "total", currentSum, "totalLocked", lockedSum, "rebalanceBudget", rebalanceBudget)
+	}
+
+TARGET:
+	// targets are validators with a delegation deficit, sorted in descending order.
+	// that is, those at the top should be satisfied first to maximise progress toward goal.
+	for _, target := range targets {
+		// amount is amount we should try to redelegate toward target. This may be constrained by the remaining redelegateBudget.
+		// if it is zero (i.e. we hit the redelegation budget) break out of the loop.
+		amount := sdkmath.MinInt(target.Amount, rebalanceBudget)
+		if amount.IsZero() {
+			break
+		}
+		sources.Sort()
+		// sources are validators with available balance to redelegate, sorted in desc order.
+		for _, source := range sources {
+			switch {
+			case source.Amount.IsZero():
+				// if source is zero, skip.
+				continue
+			case source.Amount.GTE(amount):
+				// if source >= amount, fully satisfy target.
+				out = append(out, &RebalanceTarget{Amount: amount, Target: target.ValoperAddress, Source: source.ValoperAddress})
+				source.Amount = source.Amount.Sub(amount)
+				target.Amount = target.Amount.Sub(amount)
+				rebalanceBudget = rebalanceBudget.Sub(amount)
+				continue TARGET
+			case source.Amount.LT(amount):
+				// if source < amount, partially satisfy amount.
+				out = append(out, &RebalanceTarget{Amount: source.Amount, Target: target.ValoperAddress, Source: source.ValoperAddress})
+				amount = amount.Sub(source.Amount)
+				target.Amount = target.Amount.Sub(source.Amount)
+				rebalanceBudget = rebalanceBudget.Sub(source.Amount)
+				source.Amount = source.Amount.Sub(source.Amount)
+				if amount.IsZero() || rebalanceBudget.IsZero() {
+					// if the amount is fully satisfied or the rebalanceBudget is zero, skip to next target.
+					continue TARGET
+				}
+				// otherwise, try next source.
+			}
+		}
+		// we only get here if we are unable to satisfy targets due to rebalanceBudget depletion.
+		if logger != nil {
+			logger.Info("unable to satisfy targets with available sources.")
+		}
+	}
+
+	out.Sort()
 
 	return out
 }
 
+// func (d Deltas) Render() (out string) {
+// 	for _, delta := range d {
+// 		out = out + fmt.Sprintf("%s:\t%d\n", delta.ValoperAddress, delta.Amount.Int64())
+// 	}
+// 	return
+// }
+
 // MinDeltas returns the lowest value in a slice of Deltas.
 func MinDeltas(deltas ValidatorIntents) sdkmath.Int {
 	minValue := sdk.NewInt(math.MaxInt64)