Merge unlocked stake outputs (#3231)

Co-authored-by: Alberto Benegiamo <[email protected]>

wallet/chain/p/builder/builder.go

@@ -1052,6 +1052,22 @@ func (b *builder) spend(
 		}
 	}
 
+	for assetID, amount := range amountsToStake {
+		if amount == 0 {
+			continue
+		}
+
+		stakeOutputs = append(stakeOutputs, &avax.TransferableOutput{
+			Asset: avax.Asset{
+				ID: assetID,
+			},
+			Out: &secp256k1fx.TransferOutput{
+				Amt:          amount,
+				OutputOwners: *changeOwner,
+			},
+		})
+	}
+
 	// Iterate over the unlocked UTXOs
 	for _, utxo := range utxos {
 		assetID := utxo.AssetID()
@@ -1103,24 +1119,14 @@ func (b *builder) spend(
 		)
 		amountsToBurn[assetID] -= amountToBurn
 
-		amountAvalibleToStake := out.Amt - amountToBurn
+		amountAvailableToStake := out.Amt - amountToBurn
 		// Burn any value that should be burned
 		amountToStake := min(
 			remainingAmountToStake, // Amount we still need to stake
-			amountAvalibleToStake,  // Amount available to stake
+			amountAvailableToStake, // Amount available to stake
 		)
 		amountsToStake[assetID] -= amountToStake
-		if amountToStake > 0 {
-			// Some of this input was put for staking
-			stakeOutputs = append(stakeOutputs, &avax.TransferableOutput{
-				Asset: utxo.Asset,
-				Out: &secp256k1fx.TransferOutput{
-					Amt:          amountToStake,
-					OutputOwners: *changeOwner,
-				},
-			})
-		}
-		if remainingAmount := amountAvalibleToStake - amountToStake; remainingAmount > 0 {
+		if remainingAmount := amountAvailableToStake - amountToStake; remainingAmount > 0 {
 			// This input had extra value, so some of it must be returned
 			changeOutputs = append(changeOutputs, &avax.TransferableOutput{
 				Asset: utxo.Asset,

wallet/chain/p/builder_test.go

@@ -544,8 +544,35 @@ func TestAddPermissionlessValidatorTx(t *testing.T) {
 		require = require.New(t)
 
 		// backend
-		utxosKey   = testKeys[1]
-		utxos      = makeTestUTXOs(utxosKey)
+		utxosOffset uint64 = 2024
+		utxosKey           = testKeys[1]
+		utxosAddr          = utxosKey.Address()
+	)
+	makeUTXO := func(amount uint64) *avax.UTXO {
+		utxosOffset++
+		return &avax.UTXO{
+			UTXOID: avax.UTXOID{
+				TxID:        ids.Empty.Prefix(utxosOffset),
+				OutputIndex: uint32(utxosOffset),
+			},
+			Asset: avax.Asset{ID: avaxAssetID},
+			Out: &secp256k1fx.TransferOutput{
+				Amt: amount,
+				OutputOwners: secp256k1fx.OutputOwners{
+					Locktime:  0,
+					Addrs:     []ids.ShortID{utxosAddr},
+					Threshold: 1,
+				},
+			},
+		}
+	}
+
+	var (
+		utxos = []*avax.UTXO{
+			makeUTXO(testContext.AddPrimaryNetworkValidatorFee), // UTXO to pay the fee
+			makeUTXO(1 * units.NanoAvax),                        // small UTXO
+			makeUTXO(9 * units.Avax),                            // large UTXO
+		}
 		chainUTXOs = common.NewDeterministicChainUTXOs(require, map[ids.ID][]*avax.UTXO{
 			constants.PlatformChainID: utxos,
 		})
@@ -611,6 +638,11 @@ func TestAddPermissionlessValidatorTx(t *testing.T) {
 		),
 		addInputAmounts(utx.Ins),
 	)
+
+	// Outputs should be merged if possible. For example, if there are two
+	// unlocked inputs consumed for staking, this should only produce one staked
+	// output.
+	require.Len(utx.StakeOuts, 1)
 }
 
 func TestAddPermissionlessDelegatorTx(t *testing.T) {