chore: address the pr comments

CHANGELOG.md

@@ -79,7 +79,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 * (x/staking) [#12409](https://github.com/cosmos/cosmos-sdk/pull/12409) Migrate `x/staking` to self-managed parameters and deprecate it's usage of `x/params`.
 * (x/bank) [#11859](https://github.com/cosmos/cosmos-sdk/pull/11859) Move the SendEnabled information out of the Params and into the state store directly.
 * (x/gov) [#12771](https://github.com/cosmos/cosmos-sdk/pull/12771) Initial deposit requirement for proposals at submission time.
-* (x/staking) [#12967](https://github.com/cosmos/cosmos-sdk/pull/12967) `unbond` now creates only one unbonding delegation entry when multiple ubd msgs broadcast at a single height.
+* (x/staking) [#12967](https://github.com/cosmos/cosmos-sdk/pull/12967) `unbond` now creates only one unbonding delegation entry when multiple unbondings exist at a single height (e.g. through multiple messages in a transaction).
 
 ### API Breaking Changes
 

x/staking/migrations/v5/store.go

@@ -1,34 +1,42 @@
 package v5
 
 import (
+	"sort"
+
 	"github.com/cosmos/cosmos-sdk/codec"
 	storetypes "github.com/cosmos/cosmos-sdk/store/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/x/staking/types"
 )
 
 // MigrateStore performs in-place store migrations from v3 to v4.
-/*
-	this migrateStore will remove the ubdEntries with same creation_height
-	and create a new ubdEntry with updated balance and initial_balance
-*/
+// this migrateStore will remove the ubdEntries with same creation_height
+// and create a new ubdEntry with updated balance and initial_balance
 func MigrateStore(ctx sdk.Context, storeKey storetypes.StoreKey, cdc codec.BinaryCodec) error {
 	store := ctx.KVStore(storeKey)
 	iterator := sdk.KVStorePrefixIterator(store, types.UnbondingDelegationKey)
 	defer iterator.Close()
 
 	for ; iterator.Valid(); iterator.Next() {
 		ubd := types.MustUnmarshalUBD(cdc, iterator.Value())
+
 		entriesAtSameCreationHeight := make(map[int64][]types.UnbondingDelegationEntry)
+
+		var creationHeights []int64
+
 		for _, ubdEntry := range ubd.Entries {
 			entriesAtSameCreationHeight[ubdEntry.CreationHeight] = append(entriesAtSameCreationHeight[ubdEntry.CreationHeight], ubdEntry)
+			creationHeights = append(creationHeights, ubdEntry.CreationHeight)
 		}
+
 		// clear the old ubdEntries
 		ubd.Entries = nil
 
-		for _, entries := range entriesAtSameCreationHeight {
+		sort.Slice(creationHeights, func(i, j int) bool { return creationHeights[i] < creationHeights[j] })
+
+		for index := 0; index < len(creationHeights); index++ {
 			var ubdEntry types.UnbondingDelegationEntry
-			for _, entry := range entries {
+			for _, entry := range entriesAtSameCreationHeight[creationHeights[index]] {
 				ubdEntry.Balance = ubdEntry.Balance.Add(entry.Balance)
 				ubdEntry.InitialBalance = ubdEntry.InitialBalance.Add(entry.InitialBalance)
 				ubdEntry.CreationHeight = entry.CreationHeight
@@ -38,22 +46,23 @@ func MigrateStore(ctx sdk.Context, storeKey storetypes.StoreKey, cdc codec.Binar
 		}
 
 		// set the new ubd to the store
-		setUBDToStore(ctx, storeKey, cdc, ubd)
+		setUBDToStore(ctx, store, cdc, ubd)
 	}
 
 	return nil
 }
 
-func setUBDToStore(ctx sdk.Context, storeKey storetypes.StoreKey, cdc codec.BinaryCodec, ubd types.UnbondingDelegation) {
+func setUBDToStore(ctx sdk.Context, store storetypes.KVStore, cdc codec.BinaryCodec, ubd types.UnbondingDelegation) {
 	delegatorAddress := sdk.MustAccAddressFromBech32(ubd.DelegatorAddress)
 
-	store := ctx.KVStore(storeKey)
 	bz := types.MustMarshalUBD(cdc, ubd)
+
 	addr, err := sdk.ValAddressFromBech32(ubd.ValidatorAddress)
 	if err != nil {
 		panic(err)
 	}
+
 	key := types.GetUBDKey(delegatorAddress, addr)
+
 	store.Set(key, bz)
-	store.Set(types.GetUBDByValIndexKey(delegatorAddress, addr), []byte{}) // index, store empty bytes
 }