From 3c1af4229106019411ea2b785e56ea5c69e31fc7 Mon Sep 17 00:00:00 2001
From: Bob Potter <bobby.potter@gmail.com>
Date: Wed, 1 Aug 2018 16:46:18 -0500
Subject: [PATCH] AndNot: Avoid returning bitmapcontainer with arraycontainer
 cardinality

In some cases this could cause bitmaps to be serialized incorrectly.

Check if a bitmap needs to be converted to an array container after an
inplace-andNot operation.

This also adds a sanity check when writing bitmap containers to make
sure the cardinality is correct.
---
 bitmapcontainer.go            |  5 ++++-
 roaring_test.go               | 25 +++++++++++++++++++++++++
 serialization_generic.go      |  5 +++++
 serialization_littleendian.go |  4 ++++
 serialization_test.go         |  3 +--
 5 files changed, 39 insertions(+), 3 deletions(-)

diff --git a/bitmapcontainer.go b/bitmapcontainer.go
index fcb78d2b..74b7a17c 100644
--- a/bitmapcontainer.go
+++ b/bitmapcontainer.go
@@ -857,12 +857,15 @@ func (bc *bitmapContainer) andNotBitmap(value2 *bitmapContainer) container {
 	return ac
 }
 
-func (bc *bitmapContainer) iandNotBitmapSurely(value2 *bitmapContainer) *bitmapContainer {
+func (bc *bitmapContainer) iandNotBitmapSurely(value2 *bitmapContainer) container {
 	newCardinality := int(popcntMaskSlice(bc.bitmap, value2.bitmap))
 	for k := 0; k < len(bc.bitmap); k++ {
 		bc.bitmap[k] = bc.bitmap[k] &^ value2.bitmap[k]
 	}
 	bc.cardinality = newCardinality
+	if bc.getCardinality() <= arrayDefaultMaxSize {
+		return bc.toArrayContainer()
+	}
 	return bc
 }
 
diff --git a/roaring_test.go b/roaring_test.go
index e81a3bdd..dd031444 100644
--- a/roaring_test.go
+++ b/roaring_test.go
@@ -1,6 +1,7 @@
 package roaring
 
 import (
+	"bytes"
 	"log"
 	"math"
 	"math/rand"
@@ -163,6 +164,30 @@ func TestRoaringBitmapAddOffset(t *testing.T) {
 	}
 }
 
+func TestRoaringInPlaceAndNotBitmapContainer(t *testing.T) {
+	bm := NewBitmap()
+	for i := 0; i < 8192; i++ {
+		bm.Add(uint32(i))
+	}
+	toRemove := NewBitmap()
+	for i := 128; i < 8192; i++ {
+		toRemove.Add(uint32(i))
+	}
+	bm.AndNot(toRemove)
+
+	var b bytes.Buffer
+	_, err := bm.WriteTo(&b)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	bm2 := NewBitmap()
+	bm2.ReadFrom(bytes.NewBuffer(b.Bytes()))
+	if !bm2.Equals(bm) {
+		t.Errorf("expected %s to equal %s", bm2, bm)
+	}
+}
+
 // https://github.com/RoaringBitmap/roaring/issues/64
 func TestFlip64(t *testing.T) {
 	bm := New()
diff --git a/serialization_generic.go b/serialization_generic.go
index 7fcef769..f4805ca3 100644
--- a/serialization_generic.go
+++ b/serialization_generic.go
@@ -4,6 +4,7 @@ package roaring
 
 import (
 	"encoding/binary"
+	"errors"
 	"io"
 )
 
@@ -26,6 +27,10 @@ func (b *arrayContainer) readFrom(stream io.Reader) (int, error) {
 }
 
 func (b *bitmapContainer) writeTo(stream io.Writer) (int, error) {
+	if b.cardinality <= arrayDefaultMaxSize {
+		return 0, errors.New("refusing to write bitmap container with cardinality of array container")
+	}
+
 	// Write set
 	buf := make([]byte, 8*len(b.bitmap))
 	for i, v := range b.bitmap {
diff --git a/serialization_littleendian.go b/serialization_littleendian.go
index c1d3ad30..ba5b7539 100644
--- a/serialization_littleendian.go
+++ b/serialization_littleendian.go
@@ -3,6 +3,7 @@
 package roaring
 
 import (
+	"errors"
 	"io"
 	"reflect"
 	"unsafe"
@@ -14,6 +15,9 @@ func (ac *arrayContainer) writeTo(stream io.Writer) (int, error) {
 }
 
 func (bc *bitmapContainer) writeTo(stream io.Writer) (int, error) {
+	if bc.cardinality <= arrayDefaultMaxSize {
+		return 0, errors.New("refusing to write bitmap container with cardinality of array container")
+	}
 	buf := uint64SliceAsByteSlice(bc.bitmap)
 	return stream.Write(buf)
 }
diff --git a/serialization_test.go b/serialization_test.go
index 32bfa149..97c7536f 100644
--- a/serialization_test.go
+++ b/serialization_test.go
@@ -493,12 +493,11 @@ func TestSerializationRunOnly033(t *testing.T) {
 func TestSerializationBitmapOnly034(t *testing.T) {
 
 	Convey("bitmapContainer writeTo and readFrom should return logically equivalent containers", t, func() {
-
 		seed := int64(42)
 		rand.Seed(seed)
 
 		trials := []trial{
-			{n: 1010, percentFill: .50, ntrial: 10},
+			{n: 8192, percentFill: .99, ntrial: 10},
 		}
 
 		tester := func(tr trial) {