Btcstaking begin block perf improvement

types/btc_schnorr_sig.go

@@ -12,8 +12,14 @@ type BIP340Signature []byte
 const BIP340SignatureLen = schnorr.SignatureSize
 
 func NewBIP340Signature(data []byte) (*BIP340Signature, error) {
+
 	var sig BIP340Signature
 	err := sig.Unmarshal(data)
+
+	if _, err := sig.ToBTCSig(); err != nil {
+		return nil, errors.New("bytes cannot be converted to a *schnorr.Signature object")
+	}
+
 	return &sig, err
 }
 
@@ -69,15 +75,6 @@ func (sig BIP340Signature) MarshalTo(data []byte) (int, error) {
 }
 
 func (sig *BIP340Signature) Unmarshal(data []byte) error {
-	newSig := BIP340Signature(data)
-
-	// ensure that the bytes can be transformed to a *schnorr.Signature object
-	// this includes all format checks
-	_, err := newSig.ToBTCSig()
-	if err != nil {
-		return errors.New("bytes cannot be converted to a *schnorr.Signature object")
-	}
-
 	*sig = data
 	return nil
 }

x/btcstaking/keeper/btc_delegators.go

@@ -77,6 +77,48 @@ func (k Keeper) IterateBTCDelegations(ctx context.Context, fpBTCPK *bbn.BIP340Pu
 	}
 }
 
+func (k Keeper) IterateBTCDelsKeys(ctx context.Context, handler func(key chainhash.Hash, delegation *types.BTCDelegation) bool) {
+	deldIter := k.btcDelegationStore(ctx).Iterator(nil, nil)
+	defer deldIter.Close()
+
+	for ; deldIter.Valid(); deldIter.Next() {
+		var deld types.BTCDelegation
+		k.cdc.MustUnmarshal(deldIter.Value(), &deld)
+		hash, err := chainhash.NewHash(deldIter.Key())
+		if err != nil {
+			panic(err)
+		}
+
+		shouldContinue := handler(*hash, &deld)
+		if !shouldContinue {
+			return
+		}
+	}
+}
+
+func (k Keeper) IterateBTCDelegationsHashes(ctx context.Context, fpBTCPK *bbn.BIP340PubKey, handler func(hash chainhash.Hash) bool) {
+	btcDelIter := k.btcDelegatorStore(ctx, fpBTCPK).Iterator(nil, nil)
+	defer btcDelIter.Close()
+	for ; btcDelIter.Valid(); btcDelIter.Next() {
+		// unmarshal delegator's delegation index
+		var btcDelIndex types.BTCDelegatorDelegationIndex
+		k.cdc.MustUnmarshal(btcDelIter.Value(), &btcDelIndex)
+		// retrieve and process each of the BTC delegation
+		for _, stakingTxHashBytes := range btcDelIndex.StakingTxHashList {
+			stakingTxHash, err := chainhash.NewHash(stakingTxHashBytes)
+			if err != nil {
+				panic(err) // only programming error is possible
+			}
+
+			shouldContinue := handler(*stakingTxHash)
+
+			if !shouldContinue {
+				return
+			}
+		}
+	}
+}
+
 // hasBTCDelegatorDelegations checks if the given BTC delegator has any BTC delegations under a given finality provider
 func (k Keeper) hasBTCDelegatorDelegations(ctx context.Context, fpBTCPK *bbn.BIP340PubKey, delBTCPK *bbn.BIP340PubKey) bool {
 	fpBTCPKBytes := fpBTCPK.MustMarshal()

x/btcstaking/keeper/keeper.go

@@ -9,6 +9,7 @@ import (
 	"cosmossdk.io/log"
 	"github.com/babylonchain/babylon/x/btcstaking/types"
 	"github.com/btcsuite/btcd/chaincfg"
+	"github.com/btcsuite/btcd/chaincfg/chainhash"
 	"github.com/cosmos/cosmos-sdk/codec"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )
@@ -60,7 +61,6 @@ func (k Keeper) Logger(ctx sdk.Context) log.Logger {
 // the reward distribution cache used for distributing rewards once the block
 // is finalised by finality providers.
 func (k Keeper) BeginBlocker(ctx context.Context) error {
-	// index BTC height at the current height
 	k.IndexBTCHeight(ctx)
 
 	covenantQuorum := k.GetParams(ctx).CovenantQuorum
@@ -70,24 +70,36 @@ func (k Keeper) BeginBlocker(ctx context.Context) error {
 	}
 	wValue := k.btccKeeper.GetParams(ctx).CheckpointFinalizationTimeout
 
-	// prepare for recording finality providers with positive voting power
+	distInfos := make(map[chainhash.Hash]*types.BTCDelDistInfo)
+
+	k.IterateBTCDelsKeys(ctx, func(key chainhash.Hash, btcDel *types.BTCDelegation) bool {
+		distInfo := &types.BTCDelDistInfo{
+			BabylonPk:   btcDel.BabylonPk,
+			VotingPower: btcDel.VotingPower(btcTipHeight, wValue, covenantQuorum),
+		}
+		if distInfo.VotingPower > 0 {
+			distInfos[key] = distInfo
+		}
+		return true
+	})
+
 	activeFps := []*types.FinalityProviderWithMeta{}
-	// prepare for recording finality providers and their BTC delegations
-	// for rewards
+
 	rdc := types.NewRewardDistCache()
 
-	// iterate over all finality providers to find out non-slashed ones that have
-	// positive voting power
 	k.IterateActiveFPs(
 		ctx,
 		func(fp *types.FinalityProvider) bool {
 			fpDistInfo := types.NewFinalityProviderDistInfo(fp)
 
-			// iterate over all BTC delegations under the finality provider
-			// in order to accumulate voting power and reward dist info for it
-			k.IterateBTCDelegations(ctx, fp.BtcPk, func(btcDel *types.BTCDelegation) bool {
-				// accumulate voting power and reward distribution cache
-				fpDistInfo.AddBTCDel(btcDel, btcTipHeight, wValue, covenantQuorum)
+			k.IterateBTCDelegationsHashes(ctx, fp.BtcPk, func(hash chainhash.Hash) bool {
+				distInfo, found := distInfos[hash]
+
+				if !found {
+					return true
+				}
+
+				fpDistInfo.AddBTCDistInfo(distInfo)
 				return true
 			})
 

x/btcstaking/types/btc_slashing_tx.go

@@ -63,12 +63,6 @@ func (tx BTCSlashingTx) MarshalTo(data []byte) (int, error) {
 
 func (tx *BTCSlashingTx) Unmarshal(data []byte) error {
 	*tx = data
-
-	// ensure data can be decoded to a tx
-	if _, err := tx.ToMsgTx(); err != nil {
-		return err
-	}
-
 	return nil
 }
 

x/btcstaking/types/incentive.go

@@ -67,6 +67,14 @@ func (v *FinalityProviderDistInfo) AddBTCDel(btcDel *BTCDelegation, btcHeight ui
 	}
 }
 
+func (v *FinalityProviderDistInfo) AddBTCDistInfo(info *BTCDelDistInfo) {
+	if info.VotingPower > 0 {
+		// if this BTC delegation has voting power, append it and accumulate voting power
+		v.BtcDels = append(v.BtcDels, info)
+		v.TotalVotingPower += info.VotingPower
+	}
+}
+
 // GetBTCDelPortion returns the portion of a BTC delegation's voting power out of
 // the finality provider's total voting power
 func (v *FinalityProviderDistInfo) GetBTCDelPortion(d *BTCDelDistInfo) sdkmath.LegacyDec {

x/btcstaking/types/msg.go

@@ -73,9 +73,19 @@ func (m *MsgCreateBTCDelegation) ValidateBasic() error {
 	if m.SlashingTx == nil {
 		return fmt.Errorf("empty slashing tx")
 	}
+
+	if _, err := m.SlashingTx.ToMsgTx(); err != nil {
+		return fmt.Errorf("invalid slashing tx: %w", err)
+	}
+
 	if m.DelegatorSlashingSig == nil {
 		return fmt.Errorf("empty delegator signature")
 	}
+
+	if _, err := m.DelegatorSlashingSig.ToBTCSig(); err != nil {
+		return fmt.Errorf("invalid delegator slashing signature: %w", err)
+	}
+
 	if _, err := sdk.AccAddressFromBech32(m.Signer); err != nil {
 		return err
 	}
@@ -111,6 +121,15 @@ func (m *MsgCreateBTCDelegation) ValidateBasic() error {
 	if m.DelegatorUnbondingSlashingSig == nil {
 		return fmt.Errorf("empty delegator signature")
 	}
+
+	if _, err := m.UnbondingSlashingTx.ToMsgTx(); err != nil {
+		return fmt.Errorf("invalid unbonding slashing tx: %w", err)
+	}
+
+	if _, err := m.DelegatorUnbondingSlashingSig.ToBTCSig(); err != nil {
+		return fmt.Errorf("invalid delegator unbonding slashing signature: %w", err)
+	}
+
 	unbondingTxMsg, err := bbn.NewBTCTxFromBytes(m.UnbondingTx)
 	if err != nil {
 		return err
@@ -145,6 +164,11 @@ func (m *MsgAddCovenantSigs) ValidateBasic() error {
 	if m.UnbondingTxSig == nil {
 		return fmt.Errorf("empty covenant signature")
 	}
+
+	if _, err := m.UnbondingTxSig.ToBTCSig(); err != nil {
+		return fmt.Errorf("invalid covenant unbonding signature: %w", err)
+	}
+
 	if m.SlashingUnbondingTxSigs == nil {
 		return fmt.Errorf("empty covenant signature")
 	}
@@ -161,5 +185,9 @@ func (m *MsgBTCUndelegate) ValidateBasic() error {
 		return fmt.Errorf("empty signature from the delegator")
 	}
 
+	if _, err := m.UnbondingTxSig.ToBTCSig(); err != nil {
+		return fmt.Errorf("invalid delegator unbonding signature: %w", err)
+	}
+
 	return nil
 }