From db7eae97ff8decfd4e79026638519c897fa11a1c Mon Sep 17 00:00:00 2001
From: yyforyongyu <yy2452@columbia.edu>
Date: Wed, 17 Jan 2024 00:25:41 +0800
Subject: [PATCH] sweep: expand `InputSet` with more interface methods

This commit adds more interface methods to `InputSet` to prepare the
addition of budget-based aggregator.
---
 sweep/mock_test.go    | 16 ++++++++++++++++
 sweep/tx_input_set.go | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+)

diff --git a/sweep/mock_test.go b/sweep/mock_test.go
index fc7ff9c34..f908cf6db 100644
--- a/sweep/mock_test.go
+++ b/sweep/mock_test.go
@@ -5,8 +5,10 @@ import (
 	"testing"
 	"time"
 
+	"github.com/btcsuite/btcd/btcutil"
 	"github.com/btcsuite/btcd/chaincfg/chainhash"
 	"github.com/btcsuite/btcd/wire"
+	"github.com/lightningnetwork/lnd/fn"
 	"github.com/lightningnetwork/lnd/input"
 	"github.com/lightningnetwork/lnd/lnwallet"
 	"github.com/lightningnetwork/lnd/lnwallet/chainfee"
@@ -446,3 +448,17 @@ func (m *MockInputSet) NeedWalletInput() bool {
 
 	return args.Bool(0)
 }
+
+// DeadlineHeight returns the deadline height for the set.
+func (m *MockInputSet) DeadlineHeight() fn.Option[int32] {
+	args := m.Called()
+
+	return args.Get(0).(fn.Option[int32])
+}
+
+// Budget givens the total amount that can be used as fees by this input set.
+func (m *MockInputSet) Budget() btcutil.Amount {
+	args := m.Called()
+
+	return args.Get(0).(btcutil.Amount)
+}
diff --git a/sweep/tx_input_set.go b/sweep/tx_input_set.go
index b80ea2db0..789bb277b 100644
--- a/sweep/tx_input_set.go
+++ b/sweep/tx_input_set.go
@@ -8,6 +8,7 @@ import (
 	"github.com/btcsuite/btcd/btcutil"
 	"github.com/btcsuite/btcd/txscript"
 	"github.com/btcsuite/btcd/wire"
+	"github.com/lightningnetwork/lnd/fn"
 	"github.com/lightningnetwork/lnd/input"
 	"github.com/lightningnetwork/lnd/lnwallet"
 	"github.com/lightningnetwork/lnd/lnwallet/chainfee"
@@ -53,6 +54,24 @@ type InputSet interface {
 	// NeedWalletInput returns true if the input set needs more wallet
 	// inputs.
 	NeedWalletInput() bool
+
+	// DeadlineHeight returns an optional absolute block height to express
+	// the time-sensitivity of the input set. The outputs from a force
+	// close tx have different time preferences:
+	// - to_local: no time pressure as it can only be swept by us.
+	// - first level outgoing HTLC: must be swept before its corresponding
+	//   incoming HTLC's CLTV is reached.
+	// - first level incoming HTLC: must be swept before its CLTV is
+	//   reached.
+	// - second level HTLCs: no time pressure.
+	// - anchor: for CPFP-purpose anchor, it must be swept before any of
+	//   the above CLTVs is reached. For non-CPFP purpose anchor, there's
+	//   no time pressure.
+	DeadlineHeight() fn.Option[int32]
+
+	// Budget givens the total amount that can be used as fees by this
+	// input set.
+	Budget() btcutil.Amount
 }
 
 type txInputSetState struct {
@@ -167,6 +186,20 @@ func (t *txInputSet) Inputs() []input.Input {
 	return t.inputs
 }
 
+// Budget gives the total amount that can be used as fees by this input set.
+//
+// NOTE: this field is only used for `BudgetInputSet`.
+func (t *txInputSet) Budget() btcutil.Amount {
+	return t.totalOutput()
+}
+
+// DeadlineHeight gives the block height that this set must be confirmed by.
+//
+// NOTE: this field is only used for `BudgetInputSet`.
+func (t *txInputSet) DeadlineHeight() fn.Option[int32] {
+	return fn.None[int32]()
+}
+
 // FeeRate returns the fee rate that should be used for the tx.
 func (t *txInputSet) FeeRate() chainfee.SatPerKWeight {
 	return t.feeRate