roaring64: add portable 64-bit serialization API

roaring64/serialization_portable.go

@@ -0,0 +1,100 @@
+package roaring64
+
+import (
+	"bytes"
+	"encoding/binary"
+	"fmt"
+	"io"
+
+	"github.com/RoaringBitmap/roaring/v2"
+)
+
+// maxPortableBucketCount bounds the bucket count a portable stream may declare:
+// bucket keys are uint32, so there are at most 2^32 distinct ones.
+const maxPortableBucketCount = uint64(1) << 32
+
+// WritePortableTo writes this bitmap in the portable 64-bit format, specified at
+// https://github.com/RoaringBitmap/RoaringFormatSpec#extention-for-64-bit-implementations
+// This is the format WriteTo already produces. Call RunOptimize first for better
+// compression.
+func (rb *Bitmap) WritePortableTo(stream io.Writer) (int64, error) {
+	return rb.WriteTo(stream)
+}
+
+// ToPortableBytes returns the bytes WritePortableTo writes.
+func (rb *Bitmap) ToPortableBytes() ([]byte, error) {
+	var buf bytes.Buffer
+	_, err := rb.WritePortableTo(&buf)
+	return buf.Bytes(), err
+}
+
+// GetPortableSerializedSizeInBytes returns the number of bytes WritePortableTo
+// writes, without writing them.
+func (rb *Bitmap) GetPortableSerializedSizeInBytes() uint64 {
+	return rb.GetSerializedSizeInBytes()
+}
+
+// ReadPortableFrom reads a bitmap in the portable 64-bit format. Unlike ReadFrom
+// it checks the bucket count, the bucket key order and the absence of an empty
+// bucket, and leaves rb untouched when the stream is truncated or malformed. The
+// contents of each 32-bit bucket are not checked; call Validate if the source is
+// untrusted.
+func (rb *Bitmap) ReadPortableFrom(stream io.Reader) (p int64, err error) {
+	buf := make([]byte, 8)
+	n, err := io.ReadFull(stream, buf)
+	p += int64(n)
+	if err != nil {
+		return p, err
+	}
+	size := binary.LittleEndian.Uint64(buf)
+	if size > maxPortableBucketCount {
+		return p, fmt.Errorf("error in bitmap.ReadPortableFrom: invalid bucket count %d", size)
+	}
+
+	var hlc roaringArray64
+	// Capped so that a header claiming billions of buckets costs nothing until
+	// the buckets actually show up.
+	capHint := size
+	if capHint > 1024 {
+		capHint = 1024
+	}
+	if capHint > 0 {
+		hlc.keys = make([]uint32, 0, capHint)
+		hlc.containers = make([]*roaring.Bitmap, 0, capHint)
+		hlc.needCopyOnWrite = make([]bool, 0, capHint)
+	}
+
+	keyBuf := buf[:4]
+	previousKey := uint32(0)
+	for i := uint64(0); i < size; i++ {
+		n, err = io.ReadFull(stream, keyBuf)
+		p += int64(n)
+		if err != nil {
+			return p, fmt.Errorf("error in bitmap.ReadPortableFrom: could not read key #%d: %w", i, err)
+		}
+		key := binary.LittleEndian.Uint32(keyBuf)
+		if i > 0 && key <= previousKey {
+			return p, fmt.Errorf("error in bitmap.ReadPortableFrom: bucket keys must be strictly increasing, key #%d is %d after %d", i, key, previousKey)
+		}
+		previousKey = key
+
+		c := roaring.NewBitmap()
+		nc, cerr := c.ReadFrom(stream)
+		p += nc
+		if cerr != nil {
+			return p, fmt.Errorf("error in bitmap.ReadPortableFrom: could not deserialize bitmap for key #%d: %w", i, cerr)
+		}
+		if nc == 0 {
+			return p, fmt.Errorf("error in bitmap.ReadPortableFrom: could not deserialize bitmap for key #%d", i)
+		}
+		// The format never stores an empty bucket; drop it rather than break
+		// the invariant that every container is non-empty.
+		if c.IsEmpty() {
+			continue
+		}
+		hlc.appendContainer(key, c, false)
+	}
+	hlc.copyOnWrite = rb.highlowcontainer.copyOnWrite
+	rb.highlowcontainer = hlc
+	return p, nil
+}

roaring64/serialization_portable_test.go

@@ -0,0 +1,160 @@
+package roaring64
+
+// to run just these tests: go test -run TestPortable
+
+import (
+	"bytes"
+	"encoding/binary"
+	"math"
+	"os"
+	"path/filepath"
+	"runtime"
+	"testing"
+
+	"github.com/RoaringBitmap/roaring/v2"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// portableFixtures are reference streams produced by CRoaring, shared with the
+// Java implementation.
+var portableFixtures = []string{
+	"64mapempty.bin",
+	"64map32bitvals.bin",
+	"64mapspreadvals.bin",
+	"64maphighvals.bin",
+}
+
+// testdataPath resolves a fixture relative to this source file: some CI jobs run
+// the compiled test binary from the repository root rather than from this
+// package's directory.
+func testdataPath(t *testing.T, name string) string {
+	t.Helper()
+	_, self, _, ok := runtime.Caller(0)
+	require.True(t, ok, "could not locate the test source file")
+	return filepath.Join(filepath.Dir(self), "testdata", name)
+}
+
+func TestPortableSerializationMatchesCRoaring(t *testing.T) {
+	for _, name := range portableFixtures {
+		t.Run(name, func(t *testing.T) {
+			reference, err := os.ReadFile(testdataPath(t, name))
+			require.NoError(t, err)
+
+			bm := NewBitmap()
+			n, err := bm.ReadPortableFrom(bytes.NewReader(reference))
+			require.NoError(t, err)
+			assert.Equal(t, int64(len(reference)), n)
+			require.NoError(t, bm.Validate())
+
+			assert.Equal(t, uint64(len(reference)), bm.GetPortableSerializedSizeInBytes())
+			written, err := bm.ToPortableBytes()
+			require.NoError(t, err)
+			assert.Equal(t, reference, written)
+
+			// The portable format is what ReadFrom/WriteTo already speak.
+			legacy := NewBitmap()
+			_, err = legacy.ReadFrom(bytes.NewReader(reference))
+			require.NoError(t, err)
+			assert.True(t, legacy.Equals(bm))
+		})
+	}
+}
+
+func TestPortableDeserializationRejectsInvalidBucketCount(t *testing.T) {
+	buf := make([]byte, 8)
+	binary.LittleEndian.PutUint64(buf, maxPortableBucketCount+1)
+
+	_, err := NewBitmap().ReadPortableFrom(bytes.NewReader(buf))
+	require.Error(t, err)
+}
+
+// portableStream assembles a stream from raw (key, 32-bit bitmap) pairs,
+// bypassing WritePortableTo so that invalid streams can be built.
+func portableStream(t *testing.T, count uint64, keys []uint32, buckets []*roaring.Bitmap) []byte {
+	t.Helper()
+	var out bytes.Buffer
+	header := make([]byte, 8)
+	binary.LittleEndian.PutUint64(header, count)
+	out.Write(header)
+	for i, key := range keys {
+		keyBuf := make([]byte, 4)
+		binary.LittleEndian.PutUint32(keyBuf, key)
+		out.Write(keyBuf)
+		_, err := buckets[i].WriteTo(&out)
+		require.NoError(t, err)
+	}
+	return out.Bytes()
+}
+
+func TestPortableDeserializationRejectsUnsortedBucketKeys(t *testing.T) {
+	for _, tc := range []struct {
+		name string
+		keys []uint32
+	}{
+		{"duplicate", []uint32{1, 1}},
+		{"decreasing", []uint32{2, 1}},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			buckets := []*roaring.Bitmap{roaring.BitmapOf(3), roaring.BitmapOf(4)}
+			data := portableStream(t, 2, tc.keys, buckets)
+
+			_, err := NewBitmap().ReadPortableFrom(bytes.NewReader(data))
+			require.Error(t, err)
+		})
+	}
+}
+
+func TestPortableDeserializationDropsEmptyBucket(t *testing.T) {
+	buckets := []*roaring.Bitmap{roaring.NewBitmap(), roaring.BitmapOf(7)}
+	data := portableStream(t, 2, []uint32{0, 1}, buckets)
+
+	bm := NewBitmap()
+	n, err := bm.ReadPortableFrom(bytes.NewReader(data))
+	require.NoError(t, err)
+	assert.Equal(t, int64(len(data)), n)
+	require.NoError(t, bm.Validate())
+	assert.Equal(t, []uint64{uint64(1)<<32 | 7}, bm.ToArray())
+}
+
+func TestPortableDeserializationIsAtomicWhenTruncated(t *testing.T) {
+	source := BitmapOf(1, 1<<32, math.MaxUint64)
+	complete, err := source.ToPortableBytes()
+	require.NoError(t, err)
+
+	for _, cut := range []int{1, 5, 20} {
+		bm := BitmapOf(42)
+		_, err := bm.ReadPortableFrom(bytes.NewReader(complete[:len(complete)-cut]))
+		require.Error(t, err)
+		assert.True(t, bm.Equals(BitmapOf(42)), "bitmap was modified by a failed read")
+	}
+}
+
+func TestPortableDeserializationKeepsCopyOnWrite(t *testing.T) {
+	data, err := BitmapOf(1, 1<<32).ToPortableBytes()
+	require.NoError(t, err)
+
+	for _, copyOnWrite := range []bool{false, true} {
+		bm := NewBitmap()
+		bm.SetCopyOnWrite(copyOnWrite)
+		_, err := bm.ReadPortableFrom(bytes.NewReader(data))
+		require.NoError(t, err)
+		assert.Equal(t, copyOnWrite, bm.GetCopyOnWrite())
+	}
+}
+
+func TestPortableRoundTrip(t *testing.T) {
+	source := BitmapOf(1, 2, 3, 1<<16, 1<<32, 1<<48, math.MaxUint64)
+	source.AddRange(1<<40, 1<<40+5000)
+	source.RunOptimize()
+
+	data, err := source.ToPortableBytes()
+	require.NoError(t, err)
+	assert.Equal(t, uint64(len(data)), source.GetPortableSerializedSizeInBytes())
+
+	bm := NewBitmap()
+	n, err := bm.ReadPortableFrom(bytes.NewReader(data))
+	require.NoError(t, err)
+	assert.Equal(t, int64(len(data)), n)
+	assert.True(t, bm.Equals(source))
+}