sweep: wrap weight estimator

Preparation for a cpfp-aware weight estimator. For cpfp, a regular
weight estimator isn't sufficient, because it needs to take into account
the weight of the parent transaction(s) as well.

sweep/tx_input_set.go

@@ -33,7 +33,7 @@ const (
 type txInputSetState struct {
 	// weightEstimate is the (worst case) tx weight with the current set of
 	// inputs.
-	weightEstimate input.TxWeightEstimator
+	weightEstimate *weightEstimator
 
 	// inputTotal is the total value of all inputs.
 	inputTotal btcutil.Amount
@@ -54,7 +54,7 @@ type txInputSetState struct {
 
 func (t *txInputSetState) clone() txInputSetState {
 	s := txInputSetState{
-		weightEstimate:   t.weightEstimate,
+		weightEstimate:   t.weightEstimate.clone(),
 		inputTotal:       t.inputTotal,
 		outputValue:      t.outputValue,
 		walletInputTotal: t.walletInputTotal,
@@ -95,15 +95,20 @@ func newTxInputSet(wallet Wallet, feePerKW,
 		btcutil.Amount(relayFee.FeePerKVByte()),
 	)
 
+	state := txInputSetState{
+		weightEstimate: newWeightEstimator(),
+	}
+
 	b := txInputSet{
-		feePerKW:  feePerKW,
-		dustLimit: dustLimit,
-		maxInputs: maxInputs,
-		wallet:    wallet,
+		feePerKW:        feePerKW,
+		dustLimit:       dustLimit,
+		maxInputs:       maxInputs,
+		wallet:          wallet,
+		txInputSetState: state,
 	}
 
 	// Add the sweep tx output to the weight estimate.
-	b.weightEstimate.AddP2WKHOutput()
+	b.weightEstimate.addP2WKHOutput()
 
 	return &b
 }
@@ -126,29 +131,22 @@ func (t *txInputSet) addToState(inp input.Input, constraints addConstraints) *tx
 		return nil
 	}
 
-	// Can ignore error, because it has already been checked when
-	// calculating the yields.
-	size, isNestedP2SH, _ := inp.WitnessType().SizeUpperBound()
-
 	// Clone the current set state.
 	s := t.clone()
 
 	// Add the new input.
 	s.inputs = append(s.inputs, inp)
 
-	// Add weight of the new input.
-	if isNestedP2SH {
-		s.weightEstimate.AddNestedP2WSHInput(size)
-	} else {
-		s.weightEstimate.AddWitnessInput(size)
-	}
+	// Can ignore error, because it has already been checked when
+	// calculating the yields.
+	_ = s.weightEstimate.add(inp)
 
 	// Add the value of the new input.
 	value := btcutil.Amount(inp.SignDesc().Output.Value)
 	s.inputTotal += value
 
 	// Recalculate the tx fee.
-	weight := s.weightEstimate.Weight()
+	weight := s.weightEstimate.weight()
 	fee := t.feePerKW.FeeForWeight(int64(weight))
 
 	// Calculate the new output value.

sweep/txgenerator.go

@@ -120,7 +120,7 @@ func generateInputPartitionings(sweepableInputs []txInput,
 			"has yield=%v, weight=%v",
 			inputCount, len(txInputs.inputs)-inputCount,
 			txInputs.outputValue-txInputs.walletInputTotal,
-			txInputs.weightEstimate.Weight())
+			txInputs.weightEstimate.weight())
 
 		sets = append(sets, txInputs.inputs)
 		sweepableInputs = sweepableInputs[inputCount:]
@@ -222,11 +222,11 @@ func getWeightEstimate(inputs []input.Input) ([]input.Input, int64) {
 	//
 	// TODO(roasbeef): can be more intelligent about buffering outputs to
 	// be more efficient on-chain.
-	var weightEstimate input.TxWeightEstimator
+	weightEstimate := newWeightEstimator()
 
 	// Our sweep transaction will pay to a single segwit p2wkh address,
 	// ensure it contributes to our weight estimate.
-	weightEstimate.AddP2WKHOutput()
+	weightEstimate.addP2WKHOutput()
 
 	// For each output, use its witness type to determine the estimate
 	// weight of its witness, and add it to the proper set of spendable
@@ -235,8 +235,7 @@ func getWeightEstimate(inputs []input.Input) ([]input.Input, int64) {
 	for i := range inputs {
 		inp := inputs[i]
 
-		wt := inp.WitnessType()
-		err := wt.AddWeightEstimation(&weightEstimate)
+		err := weightEstimate.add(inp)
 		if err != nil {
 			log.Warn(err)
 
@@ -248,7 +247,7 @@ func getWeightEstimate(inputs []input.Input) ([]input.Input, int64) {
 		sweepInputs = append(sweepInputs, inp)
 	}
 
-	return sweepInputs, int64(weightEstimate.Weight())
+	return sweepInputs, int64(weightEstimate.weight())
 }
 
 // inputSummary returns a string containing a human readable summary about the

sweep/weight_estimator.go

@@ -0,0 +1,40 @@
+package sweep
+
+import (
+	"github.com/lightningnetwork/lnd/input"
+)
+
+// weightEstimator wraps a standard weight estimator instance.
+type weightEstimator struct {
+	estimator input.TxWeightEstimator
+}
+
+// newWeightEstimator instantiates a new sweeper weight estimator.
+func newWeightEstimator() *weightEstimator {
+	return &weightEstimator{}
+}
+
+// clone returns a copy of this weight estimator.
+func (w *weightEstimator) clone() *weightEstimator {
+	return &weightEstimator{
+		estimator: w.estimator,
+	}
+}
+
+// add adds the weight of the given input to the weight estimate.
+func (w *weightEstimator) add(inp input.Input) error {
+	wt := inp.WitnessType()
+
+	return wt.AddWeightEstimation(&w.estimator)
+}
+
+// addP2WKHOutput updates the weight estimate to account for an additional
+// native P2WKH output.
+func (w *weightEstimator) addP2WKHOutput() {
+	w.estimator.AddP2WKHOutput()
+}
+
+// weight gets the estimated weight of the transaction.
+func (w *weightEstimator) weight() int {
+	return w.estimator.Weight()
+}