Introduce getUnionedWritableContainer to minimize allocations for in-place Or().

this removes double allocations for roaring array containers that have the COW flag set.
previously they would first be cloned, then additional terms would be added, necessitating another slice allocation.

Author: jnewhouse

benchmark_test.go

@@ -3,6 +3,7 @@ package roaring
 import (
 	"bytes"
 	"fmt"
+	"github.com/stretchr/testify/require"
 	"math/rand"
 	"runtime"
 	"testing"
@@ -235,6 +236,75 @@ func BenchmarkUnionRoaring(b *testing.B) {
 	}
 }
 
+// BenchmarkUnionInPlaceCopyOnWrite tests the performance of bitmap.Or()
+// when the bitmap was generated via FromBuffer.
+// In this case all left containers need to be copied in order to be updated.
+// The nested for-loops test a number of different scenarios
+// with respect to the ranges and densities of bitmaps.
+func BenchmarkUnionInPlaceCopyOnWrite(b *testing.B) {
+	//uint32s to maintain 1.12 compatibility, which requires unsigned shifts.
+	startingContainerPower := uint32(4)
+	finalContainerPower := uint32(10)
+	containerIncrement := uint32(3)
+	startingItemsPower := uint32(3)
+	finalItemsPower := uint32(10)
+	itemsIncrement := uint32(7)
+	for leftContainerPower := startingContainerPower; leftContainerPower <= finalContainerPower; leftContainerPower += containerIncrement {
+		for rightContainerPower := startingContainerPower; rightContainerPower <= finalContainerPower; rightContainerPower += containerIncrement {
+			for leftItemsPerContainerPower := startingItemsPower; leftItemsPerContainerPower <= finalItemsPower; leftItemsPerContainerPower += itemsIncrement {
+				for rightItemsPerContainerPower := startingItemsPower; rightItemsPerContainerPower <= finalItemsPower; rightItemsPerContainerPower += itemsIncrement {
+					b.Run(fmt.Sprintf("%d-%d-%d-%d", leftContainerPower, rightContainerPower, leftItemsPerContainerPower, rightItemsPerContainerPower),
+						func(b *testing.B) {
+							leftMax := (1 << 16) << leftContainerPower
+							rightMax := (1 << 16) << rightContainerPower
+							leftItems := 1 << (leftContainerPower + leftItemsPerContainerPower)
+							rightItems := 1 << (rightContainerPower + rightItemsPerContainerPower)
+							left := make([][]byte, 10)
+							right := make([]*Bitmap, 10)
+							for i := 0; i < 10; i++ {
+								right[i] = NewBitmap()
+								left[i] = generateRandomBitmap(b, leftMax, leftItems)
+								_, err := right[i].FromBuffer(generateRandomBitmap(b, rightMax, rightItems))
+								require.NoError(b, err)
+							}
+							// This tests a destructive operation, Or() so have to have a fresh bitmap per test.
+							targetLefts := make([]*Bitmap, b.N)
+							for i := 0; i < b.N; i++ {
+								targetLefts[i] = NewBitmap()
+								_, err := targetLefts[i].FromBuffer(left[i%10])
+								require.NoError(b, err)
+							}
+							runActualBenchmark(b, targetLefts, right)
+						})
+				}
+			}
+		}
+	}
+}
+
+// runActualBenchmark is broken out primarily so you can profile the tests,
+// as otherwise the generation overwhelms the actual test.
+func runActualBenchmark(b *testing.B, targetLefts []*Bitmap, right []*Bitmap) uint64 {
+	b.ResetTimer()
+	b.ReportAllocs()
+	total := uint64(0)
+	for i := 0; i < b.N; i++ {
+		targetLefts[i].Or(right[i%10])
+		total += targetLefts[i].GetCardinality()
+	}
+	return total
+}
+
+func generateRandomBitmap(b *testing.B, max, terms int) []byte {
+	bitmap := NewBitmap()
+	for i := 0; i < terms; i++ {
+		bitmap.Add(uint32(rand.Intn(max)))
+	}
+	result, err := bitmap.ToBytes()
+	require.NoError(b, err)
+	return result
+}
+
 // go test -bench BenchmarkSize -run -
 func BenchmarkSizeBitset(b *testing.B) {
 	b.StopTimer()

roaring.go

@@ -1006,7 +1006,7 @@ main:
 				}
 				s2 = x2.highlowcontainer.getKeyAtIndex(pos2)
 			} else {
-				rb.highlowcontainer.replaceKeyAndContainerAtIndex(pos1, s1, rb.highlowcontainer.getWritableContainerAtIndex(pos1).ior(x2.highlowcontainer.getContainerAtIndex(pos2)), false)
+				rb.highlowcontainer.replaceKeyAndContainerAtIndex(pos1, s1, rb.highlowcontainer.getUnionedWritableContainer(pos1, x2.highlowcontainer.getContainerAtIndex(pos2)), false)
 				pos1++
 				pos2++
 				if (pos1 == length1) || (pos2 == length2) {

roaringarray.go

@@ -328,6 +328,17 @@ func (ra *roaringArray) getFastContainerAtIndex(i int, needsWriteable bool) cont
 	return c
 }
 
+// getUnionedWritableContainer switches behavior for in-place Or
+// depending on whether the container requires a copy on write.
+// If it does using the non-inplace or() method leads to fewer allocations.
+func (ra *roaringArray) getUnionedWritableContainer(pos int, other container) container {
+	if ra.needCopyOnWrite[pos] {
+		return ra.getContainerAtIndex(pos).or(other)
+	}
+	return ra.getContainerAtIndex(pos).ior(other)
+
+}
+
 func (ra *roaringArray) getWritableContainerAtIndex(i int) container {
 	if ra.needCopyOnWrite[i] {
 		ra.containers[i] = ra.containers[i].clone()