Check for possibly-dirty N values in containers modified in-place

After documenting the semantics, I noticed an arguable hole in them,
which is that you could Freeze() a dirty container, and then Repair()
wouldn't work on it. On further study, I added a roaringparanoia
check for attempts to access the N of dirty containers.

It turns out there's several such. But also, it turns out, there's
circumstances where unionInPlace is relying on the assumption that
N is valid, which it isn't always for dirty containers. Also, there's
at least one case where we rely on the assumption that forcibly
thawing a container, then calling unionInPlace on it, always modifies
that container. But that's not supposed to be true for an empty
container -- an empty container might be better handled by just
returning the container it's being unioned with. So, we drop the
unnecessary thaw (all the *InPlace ops are already thawing if/when
they need to), but we use the return from unionInPlace.

roaring/container_stash.go

@@ -74,16 +74,25 @@ var containerFlagStrings = [...]string{
 	"pristine/mapped",
 	"pristine/frozen",
 	"pristine/frozen/mapped",
+	"dirty",
+	"mapped/dirty",
+	"frozen/dirty",
+	"frozen/mapped/dirty",
+	"pristine/dirty",
+	"pristine/mapped/dirty",
+	"pristine/frozen/dirty",
+	"pristine/frozen/mapped/dirty",
 }
 
 func (f containerFlags) String() string {
-	return containerFlagStrings[f&7]
+	return containerFlagStrings[f&15]
 }
 
 const (
-	flagMapped = containerFlags(1 << iota)
-	flagFrozen
-	flagPristine
+	flagMapped   = containerFlags(1 << iota) // using memory-mapped or otherwise external storage
+	flagFrozen                               // not modifiable
+	flagPristine                             // flagPristine is used for mmapped containers referring to storage
+	flagDirty                                // flagDirty is used for containers which may have invalid N
 )
 
 func (c *Container) String() string {
@@ -232,11 +241,30 @@ func (c *Container) frozen() bool {
 	return (c.flags & flagFrozen) != 0
 }
 
+// SafeN returns N, true if it can, otherwise it returns 0, false. For
+// instance, a container subject to in-place operations can not know its
+// current N, and it's not meaningful or safe to query it until a repair,
+// so you can use this to get N "if it's available".
+func (c *Container) SafeN() (int32, bool) {
+	if c == nil {
+		return 0, true
+	}
+	if (c.flags & flagDirty) != 0 {
+		return 0, false
+	}
+	return c.n, true
+}
+
 // N returns the 1-count of the container.
 func (c *Container) N() int32 {
 	if c == nil {
 		return 0
 	}
+	if roaringParanoia {
+		if c.flags&flagDirty != 0 {
+			panic("trying to call N() on a dirty container")
+		}
+	}
 	return c.n
 }
 
@@ -281,6 +309,21 @@ func (c *Container) setMapped(mapped bool) {
 	}
 }
 
+// setDirty marks a container as "dirty" -- we don't trust container's n.
+// this should never happen except for bitmaps.
+func (c *Container) setDirty(dirty bool) {
+	if roaringParanoia {
+		if c == nil || c.frozen() {
+			panic("setDirty on nil or frozen container")
+		}
+	}
+	if dirty {
+		c.flags |= flagDirty
+	} else {
+		c.flags &^= flagDirty
+	}
+}
+
 // Freeze returns an unmodifiable container identical to c. This might
 // be c, now marked unmodifiable, or might be a new container. If c
 // is currently marked as "mapped", referring to a backing store that's
@@ -291,6 +334,14 @@ func (c *Container) Freeze() *Container {
 	if c == nil {
 		return nil
 	}
+	if c.flags&flagDirty != 0 {
+		if roaringParanoia {
+			panic("freezing dirty container")
+		}
+		// c.Repair won't work if this is already frozen, but in
+		// theory that can't happen?
+		c.Repair()
+	}
 	// don't need to freeze
 	if c.flags&flagFrozen != 0 {
 		return c

roaring/roaring.go

@@ -1366,11 +1366,12 @@ func (b *Bitmap) unionInPlace(others ...*Bitmap) {
 			tContainer := target.Containers.Get(iKey)
 			// if the target's full, short-circuit out.
 			if tContainer != nil {
-				if tContainer.N() == MaxContainerVal+1 {
+				tN, ok := tContainer.SafeN()
+				if ok && tN == MaxContainerVal+1 {
 					bitmapIters.markItersWithKeyAsHandled(i, iKey)
 					continue
 				}
-				expectedN = int64(tContainer.N())
+				expectedN = int64(tN)
 			}
 			// Check i and later iters for any max-range containers, and
 			// find out how many there are.
@@ -1445,8 +1446,7 @@ func (b *Bitmap) unionInPlace(others ...*Bitmap) {
 				jKey, jContainer := iter.iter.Value()
 
 				if iKey == jKey {
-					tContainer = tContainer.Thaw()
-					tContainer.unionInPlace(jContainer)
+					tContainer = tContainer.unionInPlace(jContainer)
 					// "iter" is a local copy from the range
 					// loop, not the actual slice member.
 					itersToUnion[j].handled = true
@@ -3191,15 +3191,24 @@ func (c *Container) optimize() *Container {
 // it is possible that the returned container will not actually be the
 // original container; in-place is a suggestion.
 func (c *Container) unionInPlace(other *Container) *Container {
-	if c == nil {
-		return other.Freeze()
-	}
-	if other == nil {
-		return c
-	}
 	// short-circuit the trivial cases
-	if c.N() == MaxContainerVal+1 || other.N() == MaxContainerVal+1 {
-		return fullContainer
+	cN, cOk := c.SafeN()
+	if cOk {
+		if cN == MaxContainerVal+1 {
+			return fullContainer
+		}
+		if cN == 0 {
+			return other.Clone()
+		}
+	}
+	oN, oOk := other.SafeN()
+	if oOk {
+		if oN == MaxContainerVal+1 {
+			return fullContainer
+		}
+		if oN == 0 {
+			return c
+		}
 	}
 	switch c.typ() {
 	case ContainerBitmap:
@@ -3487,6 +3496,11 @@ func (c *Container) runToBitmap() *Container {
 		}
 		return nil
 	}
+	if roaringParanoia {
+		if c.N() > 65536 {
+			panic(fmt.Sprintf("runToBitmap: container N %d", c.N()))
+		}
+	}
 
 	// return early if empty
 	if c.N() == 0 {
@@ -3864,6 +3878,7 @@ func (c *Container) Repair() {
 	}
 	if c.isBitmap() {
 		c.bitmapRepair()
+		c.setDirty(false)
 	}
 }
 
@@ -4536,6 +4551,7 @@ func unionBitmapRun(a, b *Container) *Container {
 // a will need to be repaired after the fact.
 func unionBitmapRunInPlace(a, b *Container) *Container {
 	a = a.Thaw()
+	a.setDirty(true)
 	bitmap := a.bitmap()
 	statsHit("union/BitmapRun")
 	for _, run := range b.runs() {
@@ -4708,10 +4724,11 @@ func compareArrayArray(a1, a2 []uint16) error {
 // an error describing any difference it finds. This is mostly intended
 // for use in tests that expect equality.
 func (c *Container) BitwiseCompare(c2 *Container) error {
-	if c.N() != c2.N() {
-		return errors.New("containers are different lengths")
+	cn, c2n := c.N(), c2.N()
+	if cn != c2n {
+		return fmt.Errorf("containers are different lengths (%d vs %d)", cn, c2n)
 	}
-	if c.N() == 0 {
+	if cn == 0 {
 		return nil
 	}
 	switch typePair(c.typ(), c2.typ()) {
@@ -4728,7 +4745,7 @@ func (c *Container) BitwiseCompare(c2 *Container) error {
 	default:
 		c3 := xor(c, c2)
 		if c3.N() != 0 {
-			return fmt.Errorf("%d bits differenct between containers", c3.N())
+			return fmt.Errorf("%d bits different between containers", c3.N())
 		}
 	}
 	return nil
@@ -4753,6 +4770,7 @@ func unionArrayBitmap(a, b *Container) *Container {
 func unionBitmapArrayInPlace(a, b *Container) *Container {
 	a = a.Thaw()
 	bitmap := a.bitmap()
+	a.setDirty(true)
 	for _, v := range b.array() {
 		bitmap[v>>6] |= (uint64(1) << (v % 64))
 	}
@@ -4800,6 +4818,7 @@ func unionBitmapBitmapInPlace(a, b *Container) *Container {
 		ab[i+2] |= bb[i+2]
 		ab[i+3] |= bb[i+3]
 	}
+	a.setDirty(true)
 	return a
 }