diff --git a/bit_array.go b/bit_array.go index dc006f0eb..84f87fd06 100644 --- a/bit_array.go +++ b/bit_array.go @@ -1,6 +1,7 @@ package common import ( + "encoding/binary" "fmt" "math/rand" "strings" @@ -273,3 +274,17 @@ func (bA *BitArray) stringIndented(indent string) string { } return fmt.Sprintf("BA{%v:%v}", bA.Bits, strings.Join(lines, indent)) } + +func (bA *BitArray) Bytes() []byte { + bA.mtx.Lock() + defer bA.mtx.Unlock() + + numBytes := (bA.Bits + 7) / 8 + bytes := make([]byte, numBytes) + for i := 0; i < len(bA.Elems); i++ { + elemBytes := [8]byte{} + binary.LittleEndian.PutUint64(elemBytes[:], bA.Elems[i]) + copy(bytes[i*8:], elemBytes[:]) + } + return bytes +} diff --git a/bit_array_test.go b/bit_array_test.go index 93274aab0..d8cbfb059 100644 --- a/bit_array_test.go +++ b/bit_array_test.go @@ -1,6 +1,7 @@ package common import ( + "bytes" "testing" ) @@ -118,3 +119,32 @@ func TestPickRandom(t *testing.T) { } } } + +func TestBytes(t *testing.T) { + bA := NewBitArray(4) + bA.SetIndex(0, true) + check := func(bA *BitArray, bz []byte) { + if !bytes.Equal(bA.Bytes(), bz) { + panic(Fmt("Expected %X but got %X", bz, bA.Bytes())) + } + } + check(bA, []byte{0x01}) + bA.SetIndex(3, true) + check(bA, []byte{0x09}) + + bA = NewBitArray(9) + check(bA, []byte{0x00, 0x00}) + bA.SetIndex(7, true) + check(bA, []byte{0x80, 0x00}) + bA.SetIndex(8, true) + check(bA, []byte{0x80, 0x01}) + + bA = NewBitArray(16) + check(bA, []byte{0x00, 0x00}) + bA.SetIndex(7, true) + check(bA, []byte{0x80, 0x00}) + bA.SetIndex(8, true) + check(bA, []byte{0x80, 0x01}) + bA.SetIndex(9, true) + check(bA, []byte{0x80, 0x03}) +} diff --git a/int.go b/int.go index 50e86a072..756e38cda 100644 --- a/int.go +++ b/int.go @@ -20,7 +20,7 @@ func SearchUint64s(a []uint64, x uint64) int { func (p Uint64Slice) Search(x uint64) int { return SearchUint64s(p, x) } -//----------------------------------------------------------------------------- +//-------------------------------------------------------------------------------- func PutUint64LE(dest []byte, i uint64) { binary.LittleEndian.PutUint64(dest, i)