Merge pull request #115 from omerfirmak/master

Make WriteTo respect BitSet length

bitset.go

@@ -272,8 +272,9 @@ func (b *BitSet) Shrink(lastbitindex uint) *BitSet {
 	copy(shrunk, b.set[:idx])
 	b.set = shrunk
 	b.length = length
-	if length < 64 {
-		b.set[idx-1] &= allBits >> uint64(64-wordsIndex(length))
+	lastWordUsedBits := length % 64
+	if lastWordUsedBits != 0 {
+		b.set[idx-1] &= allBits >> uint64(64-wordsIndex(lastWordUsedBits))
 	}
 	return b
 }
@@ -897,7 +898,8 @@ func (b *BitSet) DumpAsBits() string {
 
 // BinaryStorageSize returns the binary storage requirements
 func (b *BitSet) BinaryStorageSize() int {
-	return binary.Size(uint64(0)) + binary.Size(b.set)
+	nWords := wordsNeeded(b.length)
+	return binary.Size(uint64(0)) + binary.Size(b.set[:nWords])
 }
 
 // WriteTo writes a BitSet to a stream
@@ -915,7 +917,8 @@ func (b *BitSet) WriteTo(stream io.Writer) (int64, error) {
 	// binary.Write for large set
 	writer := bufio.NewWriter(stream)
 	var item = make([]byte, binary.Size(uint64(0))) // for serializing one uint64
-	for i := range b.set {
+	nWords := wordsNeeded(uint(length))
+	for i := range b.set[:nWords] {
 		binaryOrder.PutUint64(item, b.set[i])
 		if nn, err := writer.Write(item); err != nil {
 			return int64(i*binary.Size(uint64(0)) + nn), err

bitset_test.go

@@ -707,6 +707,15 @@ func TestShrink(t *testing.T) {
 		t.Error("24 should be set")
 		return
 	}
+
+	b = New(110)
+	b.Set(80)
+	b.Shrink(70)
+	for _, word := range b.set {
+		if (word != 0) {
+			t.Error("word should be 0", word)
+		}
+	}
 }
 
 func TestInsertAtWithSet(t *testing.T) {
@@ -1224,6 +1233,35 @@ func TestMarshalUnmarshalBinary(t *testing.T) {
 		t.Error("Bitsets are not equal:\n\t", a.DumpAsBits(), "\n\t", b.DumpAsBits())
 		return
 	}
+
+	aSetBit := uint(128)
+	a = New(256).Set(aSetBit)
+	aExpectedMarshaledSize := 8 /* length: uint64 */ + 4 * 8 /* set : [4]uint64 */
+	aMarshaled, err := a.MarshalBinary()
+
+	if err != nil || aExpectedMarshaledSize != len(aMarshaled) || aExpectedMarshaledSize != a.BinaryStorageSize() {
+		t.Error("MarshalBinary failed to produce expected (", aExpectedMarshaledSize , ") number of bytes")
+		return
+	}
+
+	shiftAmount := uint(72)
+	// https://github.com/bits-and-blooms/bitset/issues/114
+	for i := uint(0) ; i < shiftAmount; i++ {
+		a.DeleteAt(0)
+	}
+
+	aExpectedMarshaledSize = 8 /* length: uint64 */ + 3 * 8 /* set : [3]uint64 */
+	aMarshaled, err = a.MarshalBinary()
+	if err != nil || aExpectedMarshaledSize != len(aMarshaled) || aExpectedMarshaledSize != a.BinaryStorageSize() {
+		t.Error("MarshalBinary failed to produce expected (", aExpectedMarshaledSize , ") number of bytes")
+		return
+	}
+
+	copyBinary(t, a, b)
+
+	if b.Len() != 256 - shiftAmount || !b.Test(aSetBit - shiftAmount) {
+		t.Error("Shifted bitset is not copied correctly")
+	}
 }
 
 func TestMarshalUnmarshalBinaryByLittleEndian(t *testing.T) {