github.com/prysmaticlabs/prysm@v1.4.4/shared/aggregation/testing/bitlistutils.go (about)

     1  package testing
     2  
     3  import (
     4  	"math/rand"
     5  	"testing"
     6  
     7  	types "github.com/prysmaticlabs/eth2-types"
     8  	"github.com/prysmaticlabs/go-bitfield"
     9  	ethpb "github.com/prysmaticlabs/prysm/proto/eth/v1alpha1"
    10  	prysmv2 "github.com/prysmaticlabs/prysm/proto/prysm/v2"
    11  	"github.com/prysmaticlabs/prysm/shared/bls"
    12  	"github.com/prysmaticlabs/prysm/shared/timeutils"
    13  )
    14  
    15  // BitlistWithAllBitsSet creates list of bitlists with all bits set.
    16  func BitlistWithAllBitsSet(length uint64) bitfield.Bitlist {
    17  	b := bitfield.NewBitlist(length)
    18  	for i := uint64(0); i < length; i++ {
    19  		b.SetBitAt(i, true)
    20  	}
    21  	return b
    22  }
    23  
    24  // BitlistsWithSingleBitSet creates list of bitlists with a single bit set in each.
    25  func BitlistsWithSingleBitSet(n, length uint64) []bitfield.Bitlist {
    26  	lists := make([]bitfield.Bitlist, n)
    27  	for i := uint64(0); i < n; i++ {
    28  		b := bitfield.NewBitlist(length)
    29  		b.SetBitAt(i%length, true)
    30  		lists[i] = b
    31  	}
    32  	return lists
    33  }
    34  
    35  // Bitlists64WithSingleBitSet creates list of bitlists with a single bit set in each.
    36  func Bitlists64WithSingleBitSet(n, length uint64) []*bitfield.Bitlist64 {
    37  	lists := make([]*bitfield.Bitlist64, n)
    38  	for i := uint64(0); i < n; i++ {
    39  		b := bitfield.NewBitlist64(length)
    40  		b.SetBitAt(i%length, true)
    41  		lists[i] = b
    42  	}
    43  	return lists
    44  }
    45  
    46  // BitlistsWithMultipleBitSet creates list of bitlists with random n bits set.
    47  func BitlistsWithMultipleBitSet(t testing.TB, n, length, count uint64) []bitfield.Bitlist {
    48  	seed := timeutils.Now().UnixNano()
    49  	t.Logf("bitlistsWithMultipleBitSet random seed: %v", seed)
    50  	rand.Seed(seed)
    51  	lists := make([]bitfield.Bitlist, n)
    52  	for i := uint64(0); i < n; i++ {
    53  		b := bitfield.NewBitlist(length)
    54  		keys := rand.Perm(int(length))
    55  		for _, key := range keys[:count] {
    56  			b.SetBitAt(uint64(key), true)
    57  		}
    58  		lists[i] = b
    59  	}
    60  	return lists
    61  }
    62  
    63  // Bitlists64WithMultipleBitSet creates list of bitlists with random n bits set.
    64  func Bitlists64WithMultipleBitSet(t testing.TB, n, length, count uint64) []*bitfield.Bitlist64 {
    65  	seed := timeutils.Now().UnixNano()
    66  	t.Logf("Bitlists64WithMultipleBitSet random seed: %v", seed)
    67  	rand.Seed(seed)
    68  	lists := make([]*bitfield.Bitlist64, n)
    69  	for i := uint64(0); i < n; i++ {
    70  		b := bitfield.NewBitlist64(length)
    71  		keys := rand.Perm(int(length))
    72  		for _, key := range keys[:count] {
    73  			b.SetBitAt(uint64(key), true)
    74  		}
    75  		lists[i] = b
    76  	}
    77  	return lists
    78  }
    79  
    80  // MakeAttestationsFromBitlists creates list of attestations from list of bitlist.
    81  func MakeAttestationsFromBitlists(bl []bitfield.Bitlist) []*ethpb.Attestation {
    82  	atts := make([]*ethpb.Attestation, len(bl))
    83  	for i, b := range bl {
    84  		atts[i] = &ethpb.Attestation{
    85  			AggregationBits: b,
    86  			Data: &ethpb.AttestationData{
    87  				Slot:           42,
    88  				CommitteeIndex: 1,
    89  			},
    90  			Signature: bls.NewAggregateSignature().Marshal(),
    91  		}
    92  	}
    93  	return atts
    94  }
    95  
    96  // MakeSyncContributionsFromBitVector creates list of sync contributions from list of bitvector.
    97  func MakeSyncContributionsFromBitVector(bl []bitfield.Bitvector128) []*prysmv2.SyncCommitteeContribution {
    98  	c := make([]*prysmv2.SyncCommitteeContribution, len(bl))
    99  	for i, b := range bl {
   100  		c[i] = &prysmv2.SyncCommitteeContribution{
   101  			Slot:              types.Slot(1),
   102  			SubcommitteeIndex: 2,
   103  			AggregationBits:   b,
   104  			Signature:         bls.NewAggregateSignature().Marshal(),
   105  		}
   106  	}
   107  	return c
   108  }