From 675263818d90b64b8cc2666a91522ce4a3506551 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 1/2] 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..4720f2a 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 /* 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) {

From fd653373a687cb860efcbaf2fd0743fee18872db Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=96mer=20Faruk=20IRMAK?= <omerfirmak@gmail.com>
Date: Fri, 18 Nov 2022 11:19:53 +0300
Subject: [PATCH 2/2] Fix Shrink not clearing unused bits

when new length requires multiple words
---
 bitset.go      | 5 +++--
 bitset_test.go | 9 +++++++++
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/bitset.go b/bitset.go
index 7f01221..fbf0496 100644
--- a/bitset.go
+++ b/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
 }
diff --git a/bitset_test.go b/bitset_test.go
index 4720f2a..efa202f 100644
--- a/bitset_test.go
+++ b/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) {