From 5df2e1c2d298c5632df9e6eab7a4814cc636b98d Mon Sep 17 00:00:00 2001
From: Daniel Lemire <daniel@lemire.me>
Date: Thu, 12 Oct 2023 11:51:27 -0400
Subject: [PATCH] Simplifies readFrom and WriteTo, and improves their
 performance.

---
 bitset.go | 48 ++++++++++++++++++++++--------------------------
 1 file changed, 22 insertions(+), 26 deletions(-)

diff --git a/bitset.go b/bitset.go
index 4eb1733..3d2e05f 100644
--- a/bitset.go
+++ b/bitset.go
@@ -920,22 +920,24 @@ func (b *BitSet) BinaryStorageSize() int {
 //	      f, err := os.Create("myfile")
 //		       w := bufio.NewWriter(f)
 func (b *BitSet) WriteTo(stream io.Writer) (int64, error) {
-	buf := make([]byte, wordBytes)
+
 	length := uint64(b.length)
 
 	// Write length
-	binaryOrder.PutUint64(buf, length)
-	n, err := stream.Write(buf)
+	err := binary.Write(stream, binaryOrder, &length)
 	if err != nil {
-		return int64(n), err
+		// Upon failure, we do not guarantee that we
+		// return the number of bytes written.
+		return int64(0), err
 	}
 
 	nWords := b.wordCount()
-	for i := range b.set[:nWords] {
-		binaryOrder.PutUint64(buf, b.set[i])
-		if nn, err := stream.Write(buf); err != nil {
-			return int64(i*int(wordBytes) + nn + n), err
-		}
+
+	err = binary.Write(stream, binaryOrder, b.set[:nWords])
+	if err != nil {
+		// Upon failure, we do not guarantee that we
+		// return the number of bytes written.
+		return int64(wordBytes), err
 	}
 
 	return int64(b.BinaryStorageSize()), nil
@@ -958,11 +960,8 @@ func (b *BitSet) WriteTo(stream io.Writer) (int64, error) {
 //	f, err := os.Open("myfile")
 //	r := bufio.NewReader(f)
 func (b *BitSet) ReadFrom(stream io.Reader) (int64, error) {
-	buf := make([]byte, wordBytes)
-
-	// Read length first
-	_, err := io.ReadFull(stream, buf[:])
-	length := binaryOrder.Uint64(buf)
+	var length uint64
+	err := binary.Read(stream, binaryOrder, &length)
 	if err != nil {
 		if err == io.EOF {
 			err = io.ErrUnexpectedEOF
@@ -982,19 +981,16 @@ func (b *BitSet) ReadFrom(stream io.Reader) (int64, error) {
 	}
 
 	b.length = newlength
-
-	for i := 0; i < nWords; i++ {
-		if _, err := io.ReadFull(stream, buf); err != nil {
-			if err == io.EOF {
-				err = io.ErrUnexpectedEOF
-			}
-			// We do not want to leave the BitSet partially filled as
-			// it is error prone.
-			b.set = b.set[:0]
-			b.length = 0
-			return 0, err
+	err = binary.Read(stream, binaryOrder, b.set)
+	if err != nil {
+		if err == io.EOF {
+			err = io.ErrUnexpectedEOF
 		}
-		b.set[i] = binaryOrder.Uint64(buf)
+		// We do not want to leave the BitSet partially filled as
+		// it is error prone.
+		b.set = b.set[:0]
+		b.length = 0
+		return 0, err
 	}
 
 	return int64(b.BinaryStorageSize()), nil