From 2e51ec19aa5e4b425f587efb74b1f7861044ccd2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=96mer=20Faruk=20IRMAK?= <omerfirmak@gmail.com>
Date: Thu, 17 Nov 2022 12:46:08 +0300
Subject: [PATCH] Make WriteTo respect BitSet length

Fixes #114
---
 bitset.go      |  6 ++++--
 bitset_test.go | 29 +++++++++++++++++++++++++++++
 2 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/bitset.go b/bitset.go
index 83f6e70..7f01221 100644
--- a/bitset.go
+++ b/bitset.go
@@ -897,7 +897,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 +916,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
diff --git a/bitset_test.go b/bitset_test.go
index b71240d..1dddd63 100644
--- a/bitset_test.go
+++ b/bitset_test.go
@@ -1224,6 +1224,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 /* [4]uint64 : b */
+	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 /* [3]uint64 : b */
+	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) {