Add functions to allow storage of bitarrays

Add functions that allow bitarrays to be stored in offline files
or databases.

Author: evanh

bitarray/bitarray.go

@@ -20,6 +20,11 @@ efficient way.  This is *NOT* a threadsafe package.
 */
 package bitarray
 
+import (
+	"encoding/binary"
+	"io"
+)
+
 // bitArray is a struct that maintains state of a bit array.
 type bitArray struct {
 	blocks  []block
@@ -265,6 +270,70 @@ func (ba *bitArray) copy() BitArray {
 	}
 }
 
+// Write serializes the bitArray and its data and sends it to the writer.
+func Write(w io.Writer, ba *bitArray) error {
+	err := binary.Write(w, binary.LittleEndian, ba.lowest)
+	if err != nil {
+		return err
+	}
+	err = binary.Write(w, binary.LittleEndian, ba.highest)
+	if err != nil {
+		return err
+	}
+
+	var encodedanyset uint8
+	if ba.anyset {
+		encodedanyset = 1
+	} else {
+		encodedanyset = 0
+	}
+	err = binary.Write(w, binary.LittleEndian, encodedanyset)
+	if err != nil {
+		return err
+	}
+
+	err = binary.Write(w, binary.LittleEndian, ba.blocks)
+	return err
+}
+
+// Read takes a reader of a serialized bitArray created by the Write function,
+// and returns a bitArray object.
+func Read(r io.Reader) (*bitArray, error) {
+	ret := &bitArray{}
+
+	err := binary.Read(r, binary.LittleEndian, &ret.lowest)
+	if err != nil {
+		return nil, err
+	}
+
+	err = binary.Read(r, binary.LittleEndian, &ret.highest)
+	if err != nil {
+		return nil, err
+	}
+
+	var encodedanyset uint8
+	err = binary.Read(r, binary.LittleEndian, &encodedanyset)
+	if err != nil {
+		return nil, err
+	}
+
+	// anyset defaults to false so we don't need an else statement
+	if encodedanyset == 1 {
+		ret.anyset = true
+	}
+
+	var nextblock block
+	err = binary.Read(r, binary.LittleEndian, &nextblock)
+	for err == nil {
+		ret.blocks = append(ret.blocks, nextblock)
+		err = binary.Read(r, binary.LittleEndian, &nextblock)
+	}
+	if err != io.EOF {
+		return nil, err
+	}
+	return ret, nil
+}
+
 // newBitArray returns a new dense BitArray at the specified size. This is a
 // separate private constructor so unit tests don't have to constantly cast the
 // BitArray interface to the concrete type.

bitarray/bitarray_test.go

@@ -17,6 +17,7 @@ limitations under the License.
 package bitarray
 
 import (
+	"bytes"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -472,3 +473,30 @@ func BenchmarkBitArrayToNums(b *testing.B) {
 		ba.ToNums()
 	}
 }
+
+func TestBitArrayReadWrite(t *testing.T) {
+	numItems := uint64(1280)
+	input := newBitArray(numItems)
+
+	for i := uint64(0); i < numItems; i++ {
+		if i%3 == 0 {
+			input.SetBit(i)
+		}
+	}
+
+	writebuf := new(bytes.Buffer)
+	err := Write(writebuf, input)
+	assert.Equal(t, err, nil)
+
+	// 1280 bits = 20 blocks = 160 bytes, plus lowest and highest at
+	// 128 bits = 16 bytes plus 1 byte for the anyset param
+	assert.Equal(t, len(writebuf.Bytes()), 177)
+
+	expected := []byte{0, 0, 0, 0, 0, 0, 0, 0, 254}
+	assert.Equal(t, expected, writebuf.Bytes()[:9])
+
+	readbuf := bytes.NewReader(writebuf.Bytes())
+	output, err := Read(readbuf)
+	assert.Equal(t, err, nil)
+	assert.True(t, input.Equals(output))
+}

bitarray/sparse_bitarray.go

@@ -16,7 +16,11 @@ limitations under the License.
 
 package bitarray
 
-import "sort"
+import (
+	"encoding/binary"
+	"io"
+	"sort"
+)
 
 // uintSlice is an alias for a slice of ints.  Len, Swap, and Less
 // are exported to fulfill an interface needed for the search
@@ -296,6 +300,67 @@ func (sba *sparseBitArray) IntersectsBetween(other BitArray, start, stop uint64)
 	return true
 }
 
+// WriteSparse serializes the sparseBitArray and passes it to the writer.
+func WriteSparse(w io.Writer, ba *sparseBitArray) error {
+	blocksLen := uint64(len(ba.blocks))
+	indexLen := uint64(len(ba.indices))
+
+	err := binary.Write(w, binary.LittleEndian, blocksLen)
+	if err != nil {
+		return err
+	}
+
+	err = binary.Write(w, binary.LittleEndian, ba.blocks)
+	if err != nil {
+		return err
+	}
+
+	err = binary.Write(w, binary.LittleEndian, indexLen)
+	if err != nil {
+		return err
+	}
+
+	err = binary.Write(w, binary.LittleEndian, ba.indices)
+	return err
+}
+
+// ReadSparse takes a reader of a serialized sparseBitArray created by the
+// WriteSparse function, and returns a sparseBitArray object.
+func ReadSparse(r io.Reader) (*sparseBitArray, error) {
+	ret := &sparseBitArray{}
+
+	var intsToRead uint64
+	err := binary.Read(r, binary.LittleEndian, &intsToRead)
+	if err != nil {
+		return nil, err
+	}
+
+	var nextblock block
+	for i := intsToRead; i > uint64(0); i-- {
+		err = binary.Read(r, binary.LittleEndian, &nextblock)
+		if err != nil {
+			return nil, err
+		}
+		ret.blocks = append(ret.blocks, nextblock)
+	}
+
+	err = binary.Read(r, binary.LittleEndian, &intsToRead)
+	if err != nil {
+		return nil, err
+	}
+
+	var nextuint uint64
+	for i := intsToRead; i > uint64(0); i-- {
+		err = binary.Read(r, binary.LittleEndian, &nextuint)
+		if err != nil {
+			return nil, err
+		}
+		ret.indices = append(ret.indices, nextuint)
+	}
+
+	return ret, nil
+}
+
 func newSparseBitArray() *sparseBitArray {
 	return &sparseBitArray{}
 }

bitarray/sparse_bitarray_test.go

@@ -17,6 +17,7 @@ limitations under the License.
 package bitarray
 
 import (
+	"bytes"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -308,3 +309,28 @@ func BenchmarkSparseBitArrayToNums(b *testing.B) {
 		sba.ToNums()
 	}
 }
+
+func TestSparseBitArrayReadWrite(t *testing.T) {
+	numItems := uint64(1280)
+	input := newSparseBitArray()
+
+	for i := uint64(0); i < numItems; i++ {
+		if i%3 == 0 {
+			input.SetBit(i)
+		}
+	}
+
+	writebuf := new(bytes.Buffer)
+	err := WriteSparse(writebuf, input)
+	assert.Equal(t, err, nil)
+
+	assert.Equal(t, len(writebuf.Bytes()), 336)
+
+	expected := []byte{20, 0, 0, 0, 0, 0, 0, 0, 73}
+	assert.Equal(t, expected, writebuf.Bytes()[:9])
+
+	readbuf := bytes.NewReader(writebuf.Bytes())
+	output, err := ReadSparse(readbuf)
+	assert.Equal(t, err, nil)
+	assert.True(t, input.Equals(output))
+}