From ddacbca26bd05ffcb9cabc37fa9e5105b29b7ab9 Mon Sep 17 00:00:00 2001
From: yyforyongyu <yy2452@columbia.edu>
Date: Wed, 27 Mar 2024 19:41:24 +0800
Subject: [PATCH] sweep: apply `fn.NewSet` and `fn.Map` in `validateInputs`

---
 sweep/aggregator_test.go   |  8 +++----
 sweep/tx_input_set.go      | 47 ++++++++++++++++++++++++--------------
 sweep/tx_input_set_test.go | 23 +++++++++++++++----
 3 files changed, 52 insertions(+), 26 deletions(-)

diff --git a/sweep/aggregator_test.go b/sweep/aggregator_test.go
index e1a5e093d7..537c02662e 100644
--- a/sweep/aggregator_test.go
+++ b/sweep/aggregator_test.go
@@ -657,19 +657,19 @@ func TestBudgetAggregatorCreateInputSets(t *testing.T) {
 	pi1 := SweeperInput{
 		Input: mockInput1,
 		params: Params{
-			DeadlineHeight: fn.Some(int32(1)),
+			DeadlineHeight: fn.Some(testHeight),
 		},
 	}
 	pi2 := SweeperInput{
 		Input: mockInput2,
 		params: Params{
-			DeadlineHeight: fn.Some(int32(1)),
+			DeadlineHeight: fn.Some(testHeight),
 		},
 	}
 	pi3 := SweeperInput{
 		Input: mockInput3,
 		params: Params{
-			DeadlineHeight: fn.Some(int32(1)),
+			DeadlineHeight: fn.Some(testHeight),
 		},
 	}
 	pi4 := SweeperInput{
@@ -678,7 +678,7 @@ func TestBudgetAggregatorCreateInputSets(t *testing.T) {
 			// This input has a deadline height that is different
 			// from the other inputs. When grouped with other
 			// inputs, it will cause an error to be returned.
-			DeadlineHeight: fn.Some(int32(2)),
+			DeadlineHeight: fn.Some(testHeight + 1),
 		},
 	}
 
diff --git a/sweep/tx_input_set.go b/sweep/tx_input_set.go
index 5ec29f8606..a4b4a746c5 100644
--- a/sweep/tx_input_set.go
+++ b/sweep/tx_input_set.go
@@ -561,30 +561,43 @@ var _ InputSet = (*BudgetInputSet)(nil)
 
 // validateInputs is used when creating new BudgetInputSet to ensure there are
 // no duplicate inputs and they all share the same deadline heights, if set.
-func validateInputs(inputs []SweeperInput) error {
+func validateInputs(inputs []SweeperInput, deadlineHeight int32) error {
 	// Sanity check the input slice to ensure it's non-empty.
 	if len(inputs) == 0 {
 		return fmt.Errorf("inputs slice is empty")
 	}
 
-	// dedupInputs is a map used to track unique outpoints of the inputs.
-	dedupInputs := make(map[wire.OutPoint]struct{})
-
-	// deadlineSet stores unique deadline heights.
-	deadlineSet := make(map[fn.Option[int32]]struct{})
-
-	for _, input := range inputs {
-		input.params.DeadlineHeight.WhenSome(func(h int32) {
-			deadlineSet[input.params.DeadlineHeight] = struct{}{}
-		})
-
-		dedupInputs[input.OutPoint()] = struct{}{}
-	}
+	// inputDeadline tracks the input's deadline height. It will be updated
+	// if the input has a different deadline than the specified
+	// deadlineHeight.
+	inputDeadline := deadlineHeight
+
+	// dedupInputs is a set used to track unique outpoints of the inputs.
+	dedupInputs := fn.NewSet(
+		// Iterate all the inputs and map the function.
+		fn.Map(func(inp SweeperInput) wire.OutPoint {
+			// If the input has a deadline height, we'll check if
+			// it's the same as the specified.
+			inp.params.DeadlineHeight.WhenSome(func(h int32) {
+				// Exit early if the deadlines matched.
+				if h == deadlineHeight {
+					return
+				}
+
+				// Update the deadline height if it's
+				// different.
+				inputDeadline = h
+			})
+
+			return inp.OutPoint()
+		}, inputs)...,
+	)
 
 	// Make sure the inputs share the same deadline height when there is
 	// one.
-	if len(deadlineSet) > 1 {
-		return fmt.Errorf("inputs have different deadline heights")
+	if inputDeadline != deadlineHeight {
+		return fmt.Errorf("input deadline height not matched: want "+
+			"%d, got %d", deadlineHeight, inputDeadline)
 	}
 
 	// Provide a defensive check to ensure that we don't have any duplicate
@@ -601,7 +614,7 @@ func NewBudgetInputSet(inputs []SweeperInput,
 	deadlineHeight int32) (*BudgetInputSet, error) {
 
 	// Validate the supplied inputs.
-	if err := validateInputs(inputs); err != nil {
+	if err := validateInputs(inputs, deadlineHeight); err != nil {
 		return nil, err
 	}
 
diff --git a/sweep/tx_input_set_test.go b/sweep/tx_input_set_test.go
index 6699886bad..2cf97dd0ba 100644
--- a/sweep/tx_input_set_test.go
+++ b/sweep/tx_input_set_test.go
@@ -282,21 +282,34 @@ func TestNewBudgetInputSet(t *testing.T) {
 			DeadlineHeight: fn.Some(int32(2)),
 		},
 	}
+	input3 := SweeperInput{
+		Input: inp2,
+		params: Params{
+			Budget:         100,
+			DeadlineHeight: fn.Some(testHeight),
+		},
+	}
 
 	// Pass a slice of inputs with different deadline heights.
 	set, err = NewBudgetInputSet([]SweeperInput{input1, input2}, testHeight)
-	rt.ErrorContains(err, "inputs have different deadline heights")
+	rt.ErrorContains(err, "input deadline height not matched")
 	rt.Nil(set)
 
-	// Pass a slice of inputs that only one input has the deadline height.
+	// Pass a slice of inputs that only one input has the deadline height,
+	// but it has a different value than the specified testHeight.
 	set, err = NewBudgetInputSet([]SweeperInput{input0, input2}, testHeight)
-	rt.NoError(err)
-	rt.NotNil(set)
+	rt.ErrorContains(err, "input deadline height not matched")
+	rt.Nil(set)
 
 	// Pass a slice of inputs that are duplicates.
-	set, err = NewBudgetInputSet([]SweeperInput{input1, input1}, testHeight)
+	set, err = NewBudgetInputSet([]SweeperInput{input3, input3}, testHeight)
 	rt.ErrorContains(err, "duplicate inputs")
 	rt.Nil(set)
+
+	// Pass a slice of inputs that only one input has the deadline height,
+	set, err = NewBudgetInputSet([]SweeperInput{input0, input3}, testHeight)
+	rt.NoError(err)
+	rt.NotNil(set)
 }
 
 // TestBudgetInputSetAddInput checks that `addInput` correctly updates the