Fix Or/Xor dropping bits when combining bitmaps of different lengths (#51)

* fix: union/xor of different-sized bitmaps drops bits in variadic path

Or and Xor with extra bitmaps used _or_many/_xor_many, which apply a single
word count to every input. The count came from pointersOf, whose length is
the maximum over the extra slice only and ignores other, so a longer other
lost its tail words. More generally a uniform count cannot combine operands
of different lengths: a longer source is truncated and a shorter one is read
past its end.

Fold each source in with the single-operand SIMD primitive over its own
length, matching the scalar reference in simd_generic.go. Fixes the amd64
accelerated and AVX-512 branches and the arm64 branch.

* Keep the _many primitives for equal-sized operands, fold only ragged ones

* test: cover longer ragged bitmap extras

---------

Co-authored-by: Roman <roman.atachiants@gmail.com>

bitmap.go

@@ -236,6 +236,36 @@ func dimensionsOf(n, m int) uint64 {
 
 // pointersOf returns a pointer to an array containing pointers to the
 // first element of each bitmap and the maximum length of all bitmaps
+// simdBinary is one of the accelerated two-operand primitives.
+type simdBinary func(dst, src unsafe.Pointer, n uint64)
+
+// uniformLen returns the word count shared by every source, or zero when the sources
+// differ in length. The _many primitives apply a single count to all of their inputs,
+// so they can only combine sources that are all the same length.
+func uniformLen(other Bitmap, extra []Bitmap) int {
+	n := len(other)
+	for i := range extra {
+		if len(extra[i]) != n {
+			return 0
+		}
+	}
+	return n
+}
+
+// foldEach applies op to every source in turn, each over its own length. This is what
+// ragged inputs need: a longer source would be truncated by a shared count, a shorter
+// one read past its end.
+func foldEach(dst *Bitmap, op simdBinary, other Bitmap, extra []Bitmap) {
+	if len(other) > 0 {
+		op(unsafe.Pointer(&(*dst)[0]), unsafe.Pointer(&other[0]), uint64(len(other)))
+	}
+	for i := range extra {
+		if len(extra[i]) > 0 {
+			op(unsafe.Pointer(&(*dst)[0]), unsafe.Pointer(&extra[i][0]), uint64(len(extra[i])))
+		}
+	}
+}
+
 func pointersOf(other Bitmap, extra []Bitmap) (unsafe.Pointer, int) {
 	out := make([]unsafe.Pointer, len(extra)+1)
 	out[0] = unsafe.Pointer(&other[0])

bitmap_amd64.go

@@ -83,16 +83,26 @@ func (dst *Bitmap) Or(other Bitmap, extra ...Bitmap) {
 		case 0:
 			_or(unsafe.Pointer(&(*dst)[0]), unsafe.Pointer(&other[0]), uint64(len(other)))
 		default:
-			vx, max := pointersOf(other, extra)
-			_or_many(unsafe.Pointer(&(*dst)[0]), vx, dimensionsOf(max, len(extra)+1))
+			// _or_many applies one word count to every input, so it only fits
+			// sources that share a length; ragged ones are folded in one by one.
+			if n := uniformLen(other, extra); n > 0 {
+				vx, _ := pointersOf(other, extra)
+				_or_many(unsafe.Pointer(&(*dst)[0]), vx, dimensionsOf(n, len(extra)+1))
+			} else {
+				foldEach(dst, _or, other, extra)
+			}
 		}
 	case isAVX512:
 		switch len(extra) {
 		case 0:
 			_or_avx512(unsafe.Pointer(&(*dst)[0]), unsafe.Pointer(&other[0]), uint64(len(other)))
 		default:
-			vx, max := pointersOf(other, extra)
-			_or_many_avx512(unsafe.Pointer(&(*dst)[0]), vx, dimensionsOf(max, len(extra)+1))
+			if n := uniformLen(other, extra); n > 0 {
+				vx, _ := pointersOf(other, extra)
+				_or_many_avx512(unsafe.Pointer(&(*dst)[0]), vx, dimensionsOf(n, len(extra)+1))
+			} else {
+				foldEach(dst, _or_avx512, other, extra)
+			}
 		}
 	default:
 		or(*dst, other, extra)
@@ -113,16 +123,26 @@ func (dst *Bitmap) Xor(other Bitmap, extra ...Bitmap) {
 		case 0:
 			_xor(unsafe.Pointer(&(*dst)[0]), unsafe.Pointer(&other[0]), uint64(len(other)))
 		default:
-			vx, max := pointersOf(other, extra)
-			_xor_many(unsafe.Pointer(&(*dst)[0]), vx, dimensionsOf(max, len(extra)+1))
+			// _xor_many applies one word count to every input, so it only fits
+			// sources that share a length; ragged ones are folded in one by one.
+			if n := uniformLen(other, extra); n > 0 {
+				vx, _ := pointersOf(other, extra)
+				_xor_many(unsafe.Pointer(&(*dst)[0]), vx, dimensionsOf(n, len(extra)+1))
+			} else {
+				foldEach(dst, _xor, other, extra)
+			}
 		}
 	case isAVX512:
 		switch len(extra) {
 		case 0:
 			_xor_avx512(unsafe.Pointer(&(*dst)[0]), unsafe.Pointer(&other[0]), uint64(len(other)))
 		default:
-			vx, max := pointersOf(other, extra)
-			_xor_many_avx512(unsafe.Pointer(&(*dst)[0]), vx, dimensionsOf(max, len(extra)+1))
+			if n := uniformLen(other, extra); n > 0 {
+				vx, _ := pointersOf(other, extra)
+				_xor_many_avx512(unsafe.Pointer(&(*dst)[0]), vx, dimensionsOf(n, len(extra)+1))
+			} else {
+				foldEach(dst, _xor_avx512, other, extra)
+			}
 		}
 	default:
 		xor(*dst, other, extra)

bitmap_arm64.go

@@ -67,8 +67,14 @@ func (dst *Bitmap) Or(other Bitmap, extra ...Bitmap) {
 		case 0:
 			_or(unsafe.Pointer(&(*dst)[0]), unsafe.Pointer(&other[0]), uint64(len(other)))
 		default:
-			vx, max := pointersOf(other, extra)
-			_or_many(unsafe.Pointer(&(*dst)[0]), vx, dimensionsOf(max, len(extra)+1))
+			// _or_many applies one word count to every input, so it only fits
+			// sources that share a length; ragged ones are folded in one by one.
+			if n := uniformLen(other, extra); n > 0 {
+				vx, _ := pointersOf(other, extra)
+				_or_many(unsafe.Pointer(&(*dst)[0]), vx, dimensionsOf(n, len(extra)+1))
+			} else {
+				foldEach(dst, _or, other, extra)
+			}
 		}
 	default:
 		or(*dst, other, extra)
@@ -89,8 +95,14 @@ func (dst *Bitmap) Xor(other Bitmap, extra ...Bitmap) {
 		case 0:
 			_xor(unsafe.Pointer(&(*dst)[0]), unsafe.Pointer(&other[0]), uint64(len(other)))
 		default:
-			vx, max := pointersOf(other, extra)
-			_xor_many(unsafe.Pointer(&(*dst)[0]), vx, dimensionsOf(max, len(extra)+1))
+			// _xor_many applies one word count to every input, so it only fits
+			// sources that share a length; ragged ones are folded in one by one.
+			if n := uniformLen(other, extra); n > 0 {
+				vx, _ := pointersOf(other, extra)
+				_xor_many(unsafe.Pointer(&(*dst)[0]), vx, dimensionsOf(n, len(extra)+1))
+			} else {
+				foldEach(dst, _xor, other, extra)
+			}
 		}
 	default:
 		xor(*dst, other, extra)

bitmap_test.go

@@ -130,6 +130,25 @@ func BenchmarkMany(b *testing.B) {
 	run(b, "and4-batch", func(index Bitmap) {
 		index.And(other, other, other, other)
 	})
+
+	run(b, "or4-batch", func(index Bitmap) {
+		index.Or(other, other, other, other)
+	})
+
+	run(b, "xor4-batch", func(index Bitmap) {
+		index.Xor(other, other, other, other)
+	})
+
+	shorter := make(Bitmap, 1000000/64/2)
+	shorter.Set(1000000 / 2)
+
+	run(b, "or4-batch-ragged", func(index Bitmap) {
+		index.Or(other, other, shorter, other)
+	})
+
+	run(b, "xor4-batch-ragged", func(index Bitmap) {
+		index.Xor(other, other, shorter, other)
+	})
 }
 
 func TestSetRemove(t *testing.T) {
@@ -629,3 +648,56 @@ func testTruthTables(t *testing.T) {
 		assert.Equal(t, 0b0110, int(a[0]))
 	}
 }
+
+func TestOrRaggedSizes(t *testing.T) {
+	// Ragged inputs: every set bit must survive regardless of which is longest.
+	var dst Bitmap
+	dst.Set(1)  // word 0
+	dst.Set(70) // word 1
+
+	var other Bitmap
+	other.Set(3)   // word 0
+	other.Set(200) // word 3, beyond every extra bitmap
+
+	var short Bitmap
+	short.Set(5) // word 0 only
+
+	var mid Bitmap
+	mid.Set(130) // word 2
+
+	var long Bitmap
+	long.Set(260) // word 4, beyond other
+
+	dst.Or(other, short, mid, long)
+
+	want := []uint32{1, 70, 3, 200, 5, 130, 260}
+	for _, v := range want {
+		assert.True(t, dst.Contains(v), "missing bit "+strconv.Itoa(int(v)))
+	}
+	assert.Equal(t, len(want), dst.Count())
+}
+
+func TestXorRaggedSizes(t *testing.T) {
+	var dst Bitmap
+	dst.Set(1)
+	dst.Set(200) // shared with other, must cancel out
+
+	var other Bitmap
+	other.Set(200) // cancels dst's bit 200
+	other.Set(201) // survives, beyond every extra bitmap
+
+	var short Bitmap
+	short.Set(5)
+
+	var long Bitmap
+	long.Set(260) // word 4, beyond other
+
+	dst.Xor(other, short, long)
+
+	assert.True(t, dst.Contains(1))
+	assert.False(t, dst.Contains(200), "bit 200 should cancel")
+	assert.True(t, dst.Contains(201), "bit 201 from longer source must survive")
+	assert.True(t, dst.Contains(5))
+	assert.True(t, dst.Contains(260), "bit 260 from longer extra must survive")
+	assert.Equal(t, 4, dst.Count())
+}