sweep: extract positive input set struct

A refactoring that introduces no functional changes. This prepares for
the addition of wallet utxos to push the sweep tx above the dust limit.

It also enabled access to input-specific sweep parameters during tx
generation. This will be used in later commits to control the sweep
process.

sweep/sweeper.go

@@ -99,6 +99,13 @@ type pendingInput struct {
 	lastFeeRate chainfee.SatPerKWeight
 }
 
+// parameters returns the sweep parameters for this input.
+//
+// NOTE: Part of the txInput interface.
+func (p *pendingInput) parameters() Params {
+	return p.params
+}
+
 // pendingInputs is a type alias for a set of pending inputs.
 type pendingInputs = map[wire.OutPoint]*pendingInput
 
@@ -789,7 +796,7 @@ func (s *UtxoSweeper) getInputLists(cluster inputCluster,
 	// contain inputs that failed before. Therefore we also add sets
 	// consisting of only new inputs to the list, to make sure that new
 	// inputs are given a good, isolated chance of being published.
-	var newInputs, retryInputs []input.Input
+	var newInputs, retryInputs []txInput
 	for _, input := range cluster.inputs {
 		// Skip inputs that have a minimum publish height that is not
 		// yet reached.

sweep/tx_input_set.go

@@ -0,0 +1,132 @@
+package sweep
+
+import (
+	"github.com/btcsuite/btcutil"
+	"github.com/btcsuite/btcwallet/wallet/txrules"
+	"github.com/lightningnetwork/lnd/input"
+	"github.com/lightningnetwork/lnd/lnwallet/chainfee"
+)
+
+// txInputSet is an object that accumulates tx inputs and keeps running counters
+// on various properties of the tx.
+type txInputSet struct {
+	// weightEstimate is the (worst case) tx weight with the current set of
+	// inputs.
+	weightEstimate input.TxWeightEstimator
+
+	// inputTotal is the total value of all inputs.
+	inputTotal btcutil.Amount
+
+	// outputValue is the value of the tx output.
+	outputValue btcutil.Amount
+
+	// feePerKW is the fee rate used to calculate the tx fee.
+	feePerKW chainfee.SatPerKWeight
+
+	// inputs is the set of tx inputs.
+	inputs []input.Input
+
+	// dustLimit is the minimum output value of the tx.
+	dustLimit btcutil.Amount
+
+	// maxInputs is the maximum number of inputs that will be accepted in
+	// the set.
+	maxInputs int
+}
+
+// newTxInputSet constructs a new, empty input set.
+func newTxInputSet(feePerKW, relayFee chainfee.SatPerKWeight,
+	maxInputs int) *txInputSet {
+
+	dustLimit := txrules.GetDustThreshold(
+		input.P2WPKHSize,
+		btcutil.Amount(relayFee.FeePerKVByte()),
+	)
+
+	b := txInputSet{
+		feePerKW:  feePerKW,
+		dustLimit: dustLimit,
+		maxInputs: maxInputs,
+	}
+
+	// Add the sweep tx output to the weight estimate.
+	b.weightEstimate.AddP2WKHOutput()
+
+	return &b
+}
+
+// dustLimitReached returns true if we've accumulated enough inputs to meet the
+// dust limit.
+func (t *txInputSet) dustLimitReached() bool {
+	return t.outputValue >= t.dustLimit
+}
+
+// add adds a new input to the set. It returns a bool indicating whether the
+// input was added to the set. An input is rejected if it decreases the tx
+// output value after paying fees.
+func (t *txInputSet) add(input input.Input) bool {
+	// Stop if max inputs is reached.
+	if len(t.inputs) == t.maxInputs {
+		return false
+	}
+
+	// Can ignore error, because it has already been checked when
+	// calculating the yields.
+	size, isNestedP2SH, _ := input.WitnessType().SizeUpperBound()
+
+	// Add weight of this new candidate input to a copy of the weight
+	// estimator.
+	newWeightEstimate := t.weightEstimate
+	if isNestedP2SH {
+		newWeightEstimate.AddNestedP2WSHInput(size)
+	} else {
+		newWeightEstimate.AddWitnessInput(size)
+	}
+
+	value := btcutil.Amount(input.SignDesc().Output.Value)
+	newInputTotal := t.inputTotal + value
+
+	weight := newWeightEstimate.Weight()
+	fee := t.feePerKW.FeeForWeight(int64(weight))
+
+	// Calculate the output value if the current input would be
+	// added to the set.
+	newOutputValue := newInputTotal - fee
+
+	// If adding this input makes the total output value of the set
+	// decrease, this is a negative yield input. We don't add the input to
+	// the set and return the outcome.
+	if newOutputValue <= t.outputValue {
+		return false
+	}
+
+	// Update running values.
+	t.inputTotal = newInputTotal
+	t.outputValue = newOutputValue
+	t.inputs = append(t.inputs, input)
+	t.weightEstimate = newWeightEstimate
+
+	return true
+}
+
+// addPositiveYieldInputs adds sweepableInputs that have a positive yield to the
+// input set. This function assumes that the list of inputs is sorted descending
+// by yield.
+//
+// TODO(roasbeef): Consider including some negative yield inputs too to clean
+// up the utxo set even if it costs us some fees up front.  In the spirit of
+// minimizing any negative externalities we cause for the Bitcoin system as a
+// whole.
+func (t *txInputSet) addPositiveYieldInputs(sweepableInputs []txInput) {
+	for _, input := range sweepableInputs {
+		// Try to add the input to the transaction. If that doesn't
+		// succeed because it wouldn't increase the output value,
+		// return. Assuming inputs are sorted by yield, any further
+		// inputs wouldn't increase the output value either.
+		if !t.add(input) {
+			return
+		}
+	}
+
+	// We managed to add all inputs to the set.
+}

sweep/tx_input_set_test.go

@@ -0,0 +1,62 @@
+package sweep
+
+import (
+	"testing"
+
+	"github.com/btcsuite/btcutil"
+	"github.com/lightningnetwork/lnd/input"
+)
+
+// TestTxInputSet tests adding various sized inputs to the set.
+func TestTxInputSet(t *testing.T) {
+	const (
+		feeRate   = 1000
+		relayFee  = 300
+		maxInputs = 10
+	)
+	set := newTxInputSet(feeRate, relayFee, maxInputs)
+
+	if set.dustLimit != 537 {
+		t.Fatalf("incorrect dust limit")
+	}
+
+	// Create a 300 sat input. The fee to sweep this input to a P2WKH output
+	// is 439 sats. That means that this input yields -139 sats and we
+	// expect it not to be added.
+	if set.add(createP2WKHInput(300)) {
+		t.Fatal("expected add of negatively yielding input to fail")
+	}
+
+	// A 700 sat input should be accepted into the set, because it yields
+	// positively.
+	if !set.add(createP2WKHInput(700)) {
+		t.Fatal("expected add of positively yielding input to succeed")
+	}
+
+	// The tx output should now be 700-439 = 261 sats. The dust limit isn't
+	// reached yet.
+	if set.outputValue != 261 {
+		t.Fatal("unexpected output value")
+	}
+	if set.dustLimitReached() {
+		t.Fatal("expected dust limit not yet to be reached")
+	}
+
+	// Add a 1000 sat input. This increases the tx fee to 712 sats. The tx
+	// output should now be 1000+700 - 712 = 988 sats.
+	if !set.add(createP2WKHInput(1000)) {
+		t.Fatal("expected add of positively yielding input to succeed")
+	}
+	if set.outputValue != 988 {
+		t.Fatal("unexpected output value")
+	}
+	if !set.dustLimitReached() {
+		t.Fatal("expected dust limit to be reached")
+	}
+}
+
+// createP2WKHInput returns a P2WKH test input with the specified amount.
+func createP2WKHInput(amt btcutil.Amount) input.Input {
+	input := createTestInput(int64(amt), input.WitnessKeyHash)
+	return &input
+}

sweep/txgenerator.go

@@ -9,7 +9,6 @@ import (
 	"github.com/btcsuite/btcd/txscript"
 	"github.com/btcsuite/btcd/wire"
 	"github.com/btcsuite/btcutil"
-	"github.com/btcsuite/btcwallet/wallet/txrules"
 	"github.com/lightningnetwork/lnd/input"
 	"github.com/lightningnetwork/lnd/lnwallet/chainfee"
 )
@@ -21,6 +20,13 @@ var (
 	DefaultMaxInputsPerTx = 100
 )
 
+// txInput is an interface that provides the input data required for tx
+// generation.
+type txInput interface {
+	input.Input
+	parameters() Params
+}
+
 // inputSet is a set of inputs that can be used as the basis to generate a tx
 // on.
 type inputSet []input.Input
@@ -30,17 +36,10 @@ type inputSet []input.Input
 // contains up to the configured maximum number of inputs. Negative yield
 // inputs are skipped. No input sets with a total value after fees below the
 // dust limit are returned.
-func generateInputPartitionings(sweepableInputs []input.Input,
+func generateInputPartitionings(sweepableInputs []txInput,
 	relayFeePerKW, feePerKW chainfee.SatPerKWeight,
 	maxInputsPerTx int) ([]inputSet, error) {
 
-	// Calculate dust limit based on the P2WPKH output script of the sweep
-	// txes.
-	dustLimit := txrules.GetDustThreshold(
-		input.P2WPKHSize,
-		btcutil.Amount(relayFeePerKW.FeePerKVByte()),
-	)
-
 	// Sort input by yield. We will start constructing input sets starting
 	// with the highest yield inputs. This is to prevent the construction
 	// of a set with an output below the dust limit, causing the sweep
@@ -75,15 +74,21 @@ func generateInputPartitionings(sweepableInputs []input.Input,
 	// Select blocks of inputs up to the configured maximum number.
 	var sets []inputSet
 	for len(sweepableInputs) > 0 {
-		// Get the maximum number of inputs from sweepableInputs that
+		// Start building a set of positive-yield tx inputs under the
-		// we can use to create a positive yielding set from.
+		// condition that the tx will be published with the specified
-		count, outputValue := getPositiveYieldInputs(
+		// fee rate.
-			sweepableInputs, maxInputsPerTx, feePerKW,
+		txInputs := newTxInputSet(
+			feePerKW, relayFeePerKW, maxInputsPerTx,
 		)
 
-		// If there are no positive yield inputs left, we can stop
+		// From the set of sweepable inputs, keep adding inputs to the
-		// here.
+		// input set until the tx output value no longer goes up or the
-		if count == 0 {
+		// maximum number of inputs is reached.
+		txInputs.addPositiveYieldInputs(sweepableInputs)
+
+		// If there are no positive yield inputs, we can stop here.
+		inputCount := len(txInputs.inputs)
+		if inputCount == 0 {
 			return sets, nil
 		}
 
@@ -91,82 +96,22 @@ func generateInputPartitionings(sweepableInputs []input.Input,
 		// the dust limit, stop sweeping. Because of the sorting,
 		// continuing with the remaining inputs will only lead to sets
 		// with a even lower output value.
-		if outputValue < dustLimit {
+		if !txInputs.dustLimitReached() {
 			log.Debugf("Set value %v below dust limit of %v",
-				outputValue, dustLimit)
+				txInputs.outputValue, txInputs.dustLimit)
 			return sets, nil
 		}
 
 		log.Infof("Candidate sweep set of size=%v, has yield=%v",
-			count, outputValue)
+			inputCount, txInputs.outputValue)
 
-		sets = append(sets, sweepableInputs[:count])
+		sets = append(sets, txInputs.inputs)
-		sweepableInputs = sweepableInputs[count:]
+		sweepableInputs = sweepableInputs[inputCount:]
 	}
 
 	return sets, nil
 }
 
-// getPositiveYieldInputs returns the maximum of a number n for which holds
-// that the inputs [0,n) of sweepableInputs have a positive yield.
-// Additionally, the total values of these inputs minus the fee is returned.
-//
-// TODO(roasbeef): Consider including some negative yield inputs too to clean
-// up the utxo set even if it costs us some fees up front.  In the spirit of
-// minimizing any negative externalities we cause for the Bitcoin system as a
-// whole.
-func getPositiveYieldInputs(sweepableInputs []input.Input, maxInputs int,
-	feePerKW chainfee.SatPerKWeight) (int, btcutil.Amount) {
-
-	var weightEstimate input.TxWeightEstimator
-
-	// Add the sweep tx output to the weight estimate.
-	weightEstimate.AddP2WKHOutput()
-
-	var total, outputValue btcutil.Amount
-	for idx, input := range sweepableInputs {
-		// Can ignore error, because it has already been checked when
-		// calculating the yields.
-		size, isNestedP2SH, _ := input.WitnessType().SizeUpperBound()
-
-		// Keep a running weight estimate of the input set.
-		if isNestedP2SH {
-			weightEstimate.AddNestedP2WSHInput(size)
-		} else {
-			weightEstimate.AddWitnessInput(size)
-		}
-
-		newTotal := total + btcutil.Amount(input.SignDesc().Output.Value)
-
-		weight := weightEstimate.Weight()
-		fee := feePerKW.FeeForWeight(int64(weight))
-
-		// Calculate the output value if the current input would be
-		// added to the set.
-		newOutputValue := newTotal - fee
-
-		// If adding this input makes the total output value of the set
-		// decrease, this is a negative yield input. It shouldn't be
-		// added to the set. We return the current index as the number
-		// of inputs, so the current input is being excluded.
-		if newOutputValue <= outputValue {
-			return idx, outputValue
-		}
-
-		// Update running values.
-		total = newTotal
-		outputValue = newOutputValue
-
-		// Stop if max inputs is reached.
-		if idx == maxInputs-1 {
-			return maxInputs, outputValue
-		}
-	}
-
-	// We could add all inputs to the set, so return them all.
-	return len(sweepableInputs), outputValue
-}
-
 // createSweepTx builds a signed tx spending the inputs to a the output script.
 func createSweepTx(inputs []input.Input, outputPkScript []byte,
 	currentBlockHeight uint32, feePerKw chainfee.SatPerKWeight,