sweep: apply fn.NewSet and fn.Map in validateInputs

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),
 		},
 	}
 

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
 	}
 

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