chore: refactory to remove delegations in a single loop over dels

covenant/covenant.go

@@ -439,18 +439,6 @@ func (ce *CovenantEmulator) delegationsToBatches(dels []*types.Delegation) [][]*
 	return batches
 }
 
-func RemoveNotInCommittee(paramCache ParamsGetter, covenantSerializedPk []byte, dels []*types.Delegation) []*types.Delegation {
-	sanitized := make([]*types.Delegation, 0, len(dels))
-
-	for _, del := range dels {
-		if !IsKeyInCommittee(paramCache, covenantSerializedPk, del) {
-			continue
-		}
-		sanitized = append(sanitized, del)
-	}
-	return sanitized
-}
-
 // IsKeyInCommittee returns true if the covenant serialized public key is in the covenant committee of the
 // parameter in which the BTC delegation was included.
 func IsKeyInCommittee(paramCache ParamsGetter, covenantSerializedPk []byte, del *types.Delegation) bool {
@@ -470,21 +458,6 @@ func IsKeyInCommittee(paramCache ParamsGetter, covenantSerializedPk []byte, del
 	return false
 }
 
-// RemoveAlreadySigned remove the delegations in which the serialized covenant public key
-// alredy signed.
-func RemoveAlreadySigned(covenantSerializedPk []byte, dels []*types.Delegation) []*types.Delegation {
-	sanitized := make([]*types.Delegation, 0, len(dels))
-
-	for _, del := range dels {
-		if CovenantAlreadySigned(covenantSerializedPk, del) {
-			continue
-		}
-		sanitized = append(sanitized, del)
-	}
-
-	return sanitized
-}
-
 // CovenantAlreadySigned returns true if the covenant already signed the BTC Delegation
 func CovenantAlreadySigned(covenantSerializedPk []byte, del *types.Delegation) bool {
 	for _, covSig := range del.CovenantSigs {
@@ -501,11 +474,29 @@ func CovenantAlreadySigned(covenantSerializedPk []byte, del *types.Delegation) b
 // sanitizeDelegations removes any delegations that have already been signed by the covenant and
 // remove delegations that were not constructed with this covenant public key
 func (ce *CovenantEmulator) sanitizeDelegations(dels []*types.Delegation) []*types.Delegation {
-	covenantSerializedPk := schnorr.SerializePubKey(ce.pk)
-	// 1. Remove delegations that do not need the covenant's signature
-	delsNotSigned := RemoveAlreadySigned(covenantSerializedPk, dels)
-	// 2. Remove delegations that were not constructed with this covenant public key
-	return RemoveNotInCommittee(ce.paramCache, covenantSerializedPk, delsNotSigned)
+	return SanitizeDelegations(ce.pk, ce.paramCache, dels)
+}
+
+// SanitizeDelegations remove the delegations in which the covenant public key already signed
+// or the delegation was not constructed with that covenant public key
+func SanitizeDelegations(pk *btcec.PublicKey, paramCache ParamsGetter, dels []*types.Delegation) []*types.Delegation {
+	covenantSerializedPk := schnorr.SerializePubKey(pk)
+
+	sanitized := make([]*types.Delegation, 0, len(dels))
+	for _, del := range dels {
+		// 1. Remove delegations that do not need the covenant's signature because
+		// this covenant already signed
+		if CovenantAlreadySigned(covenantSerializedPk, del) {
+			continue
+		}
+		// 2. Remove delegations that were not constructed with this covenant public key
+		if !IsKeyInCommittee(paramCache, covenantSerializedPk, del) {
+			continue
+		}
+		sanitized = append(sanitized, del)
+	}
+
+	return sanitized
 }
 
 // covenantSigSubmissionLoop is the reactor to submit Covenant signature for BTC delegations

covenant/covenant_test.go

@@ -9,6 +9,7 @@ import (
 	"github.com/btcsuite/btcd/btcec/v2"
 	"github.com/btcsuite/btcd/btcec/v2/schnorr"
 	"github.com/btcsuite/btcd/btcutil"
+	"github.com/decred/dcrd/dcrec/secp256k1/v4"
 
 	"github.com/babylonlabs-io/babylon/btcstaking"
 	asig "github.com/babylonlabs-io/babylon/crypto/schnorr-adaptor-signature"
@@ -227,6 +228,7 @@ func TestDeduplicationWithOddKey(t *testing.T) {
 	require.NoError(t, err)
 	pubKey := randomKey.PubKey()
 
+	paramVersion := uint32(2)
 	delegations := []*types.Delegation{
 		&types.Delegation{
 			CovenantSigs: []*types.CovenantAdaptorSigInfo{
@@ -235,18 +237,26 @@ func TestDeduplicationWithOddKey(t *testing.T) {
 					Pk: pubKeyFromSchnorr,
 				},
 			},
+			ParamsVersion: paramVersion,
 		},
 		&types.Delegation{
 			CovenantSigs: []*types.CovenantAdaptorSigInfo{
 				&types.CovenantAdaptorSigInfo{
 					Pk: pubKey,
 				},
 			},
+			ParamsVersion: paramVersion,
 		},
 	}
 
+	paramsGet := NewMockParam(map[uint32]*types.StakingParams{
+		paramVersion: &types.StakingParams{
+			CovenantPks: []*secp256k1.PublicKey{oddKeyPub, pubKeyFromSchnorr},
+		},
+	})
+
 	// 4. After removing the already signed delegation, the list should have only one element
-	sanitized := covenant.RemoveAlreadySigned(schnorr.SerializePubKey(oddKeyPub), delegations)
+	sanitized := covenant.SanitizeDelegations(oddKeyPub, paramsGet, delegations)
 	require.Equal(t, 1, len(sanitized))
 }
 
@@ -293,17 +303,17 @@ func TestIsKeyInCommittee(t *testing.T) {
 	// checks the case where the covenant is NOT in the committee
 	actual := covenant.IsKeyInCommittee(paramsGet, covenantSerializedPk, delNoCovenant)
 	require.False(t, actual)
-	emptyDels := covenant.RemoveNotInCommittee(paramsGet, covenantSerializedPk, []*types.Delegation{delNoCovenant, delNoCovenant})
+	emptyDels := covenant.SanitizeDelegations(covKeyPair.PublicKey, paramsGet, []*types.Delegation{delNoCovenant, delNoCovenant})
 	require.Len(t, emptyDels, 0)
 
 	// checks the case where the covenant is in the committee
 	actual = covenant.IsKeyInCommittee(paramsGet, covenantSerializedPk, delWithCovenant)
 	require.True(t, actual)
-	dels := covenant.RemoveNotInCommittee(paramsGet, covenantSerializedPk, []*types.Delegation{delWithCovenant, delNoCovenant})
+	dels := covenant.SanitizeDelegations(covKeyPair.PublicKey, paramsGet, []*types.Delegation{delWithCovenant, delNoCovenant})
 	require.Len(t, dels, 1)
-	dels = covenant.RemoveNotInCommittee(paramsGet, covenantSerializedPk, []*types.Delegation{delWithCovenant})
+	dels = covenant.SanitizeDelegations(covKeyPair.PublicKey, paramsGet, []*types.Delegation{delWithCovenant})
 	require.Len(t, dels, 1)
-	dels = covenant.RemoveNotInCommittee(paramsGet, covenantSerializedPk, []*types.Delegation{delWithCovenant, delWithCovenant, delNoCovenant})
+	dels = covenant.SanitizeDelegations(covKeyPair.PublicKey, paramsGet, []*types.Delegation{delWithCovenant, delWithCovenant, delNoCovenant})
 	require.Len(t, dels, 2)
 }