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

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,15 +83,13 @@ func (dst *Bitmap) Or(other Bitmap, extra ...Bitmap) {
 		case 0:
 			_or(unsafe.Pointer(&(*dst)[0]), unsafe.Pointer(&other[0]), uint64(len(other)))
 		default:
-			// _or_many uses one word count for every input, so it cannot union
-			// bitmaps of different lengths. Union each source over its own.
-			if len(other) > 0 {
-				_or(unsafe.Pointer(&(*dst)[0]), unsafe.Pointer(&other[0]), uint64(len(other)))
-			}
-			for i := range extra {
-				if len(extra[i]) > 0 {
-					_or(unsafe.Pointer(&(*dst)[0]), unsafe.Pointer(&extra[i][0]), uint64(len(extra[i])))
-				}
+			// _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:
@@ -99,13 +97,11 @@ func (dst *Bitmap) Or(other Bitmap, extra ...Bitmap) {
 		case 0:
 			_or_avx512(unsafe.Pointer(&(*dst)[0]), unsafe.Pointer(&other[0]), uint64(len(other)))
 		default:
-			if len(other) > 0 {
-				_or_avx512(unsafe.Pointer(&(*dst)[0]), unsafe.Pointer(&other[0]), uint64(len(other)))
-			}
-			for i := range extra {
-				if len(extra[i]) > 0 {
-					_or_avx512(unsafe.Pointer(&(*dst)[0]), unsafe.Pointer(&extra[i][0]), uint64(len(extra[i])))
-				}
+			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:
@@ -127,15 +123,13 @@ func (dst *Bitmap) Xor(other Bitmap, extra ...Bitmap) {
 		case 0:
 			_xor(unsafe.Pointer(&(*dst)[0]), unsafe.Pointer(&other[0]), uint64(len(other)))
 		default:
-			// _xor_many uses one word count for every input, so it cannot combine
-			// bitmaps of different lengths. Fold in each source over its own.
-			if len(other) > 0 {
-				_xor(unsafe.Pointer(&(*dst)[0]), unsafe.Pointer(&other[0]), uint64(len(other)))
-			}
-			for i := range extra {
-				if len(extra[i]) > 0 {
-					_xor(unsafe.Pointer(&(*dst)[0]), unsafe.Pointer(&extra[i][0]), uint64(len(extra[i])))
-				}
+			// _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:
@@ -143,13 +137,11 @@ func (dst *Bitmap) Xor(other Bitmap, extra ...Bitmap) {
 		case 0:
 			_xor_avx512(unsafe.Pointer(&(*dst)[0]), unsafe.Pointer(&other[0]), uint64(len(other)))
 		default:
-			if len(other) > 0 {
-				_xor_avx512(unsafe.Pointer(&(*dst)[0]), unsafe.Pointer(&other[0]), uint64(len(other)))
-			}
-			for i := range extra {
-				if len(extra[i]) > 0 {
-					_xor_avx512(unsafe.Pointer(&(*dst)[0]), unsafe.Pointer(&extra[i][0]), uint64(len(extra[i])))
-				}
+			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:

bitmap_arm64.go

@@ -67,15 +67,13 @@ func (dst *Bitmap) Or(other Bitmap, extra ...Bitmap) {
 		case 0:
 			_or(unsafe.Pointer(&(*dst)[0]), unsafe.Pointer(&other[0]), uint64(len(other)))
 		default:
-			// _or_many uses one word count for every input, so it cannot union
-			// bitmaps of different lengths. Union each source over its own.
-			if len(other) > 0 {
-				_or(unsafe.Pointer(&(*dst)[0]), unsafe.Pointer(&other[0]), uint64(len(other)))
-			}
-			for i := range extra {
-				if len(extra[i]) > 0 {
-					_or(unsafe.Pointer(&(*dst)[0]), unsafe.Pointer(&extra[i][0]), uint64(len(extra[i])))
-				}
+			// _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:
@@ -97,15 +95,13 @@ func (dst *Bitmap) Xor(other Bitmap, extra ...Bitmap) {
 		case 0:
 			_xor(unsafe.Pointer(&(*dst)[0]), unsafe.Pointer(&other[0]), uint64(len(other)))
 		default:
-			// _xor_many uses one word count for every input, so it cannot combine
-			// bitmaps of different lengths. Fold in each source over its own.
-			if len(other) > 0 {
-				_xor(unsafe.Pointer(&(*dst)[0]), unsafe.Pointer(&other[0]), uint64(len(other)))
-			}
-			for i := range extra {
-				if len(extra[i]) > 0 {
-					_xor(unsafe.Pointer(&(*dst)[0]), unsafe.Pointer(&extra[i][0]), uint64(len(extra[i])))
-				}
+			// _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:

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) {