faster bitsets

Signed-off-by: Andres Taylor <[email protected]>

go/vt/vtgate/semantics/bitset/bitset.go

@@ -245,3 +245,45 @@ func Single(bit int) Bitset {
 		return toBitset(words)
 	}
 }
+
+// Merge takes multiple Bitsets and returns a single Bitset that is the
+// bitwise-OR of all of them. It attempts to reduce allocations by creating
+// a single backing slice and merging all Bitsets into it.
+func Merge(bitsets ...Bitset) Bitset {
+	if len(bitsets) == 0 {
+		return ""
+	}
+	if len(bitsets) == 1 {
+		return bitsets[0]
+	}
+
+	// 1. Find the largest length among all input Bitsets.
+	maxLen := 0
+	for _, bs := range bitsets {
+		if len(bs) > maxLen {
+			maxLen = len(bs)
+		}
+	}
+
+	// 2. Allocate a single byte slice of size maxLen.
+	merged := make([]byte, maxLen)
+
+	// 3. Merge each Bitset into the merged buffer in-place.
+	for _, bs := range bitsets {
+		for i := 0; i < len(bs); i++ {
+			merged[i] |= bs[i]
+		}
+	}
+
+	// 4. Trim trailing zeros to avoid unnecessary storage.
+	trimmed := maxLen
+	for trimmed > 0 && merged[trimmed-1] == 0 {
+		trimmed--
+	}
+	if trimmed == 0 {
+		return "" // all bits are zero
+	}
+
+	merged = merged[:trimmed]
+	return toBitset(merged)
+}

go/vt/vtgate/semantics/bitset/bitset_test.go

@@ -17,6 +17,7 @@ limitations under the License.
 package bitset
 
 import (
+	"math/rand/v2"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -219,8 +220,10 @@ func TestOr(t *testing.T) {
 
 	for _, tc := range tt {
 		t.Run(tc.name, func(t *testing.T) {
-			got := tc.bs1.Or(tc.bs2)
-			assert.Equal(t, tc.result, got)
+			gotOR := tc.bs1.Or(tc.bs2)
+			assert.Equal(t, tc.result, gotOR)
+			gotMerge := Merge(tc.bs1, tc.bs2)
+			assert.Equal(t, tc.result, gotMerge)
 		})
 	}
 }
@@ -371,3 +374,49 @@ func TestOverlaps(t *testing.T) {
 		})
 	}
 }
+
+// Setup: create a slice of random bitsets that we can reuse in both benchmarks.
+var bitsets []Bitset
+
+func init() {
+	// For demonstration, generate 100 bitsets of random length up to 128 bytes
+	for i := 0; i < 100; i++ {
+		// random length up to 128
+		bs := generateRandomBitset(rand.IntN(128) + 1)
+		bitsets = append(bitsets, bs)
+	}
+}
+
+func BenchmarkMergeManyWithOr(b *testing.B) {
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		out := Bitset("")
+		for _, bs := range bitsets {
+			out = out.Or(bs)
+		}
+		_ = out
+	}
+}
+
+func BenchmarkMergeManyWithMerge(b *testing.B) {
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		out := Merge(bitsets...)
+		_ = out
+	}
+}
+
+// generateRandomBitset builds a random Bitset of length 'sizeInBytes'
+func generateRandomBitset(sizeInBytes int) Bitset {
+	buf := make([]byte, sizeInBytes)
+	for i := range buf {
+		buf[i] = byte(rand.IntN(256))
+	}
+	// Trim trailing zero if it exists, just to mimic typical usage
+	last := sizeInBytes - 1
+	for last > 0 && buf[last] == 0 {
+		last--
+	}
+	buf = buf[:last+1]
+	return toBitset(buf)
+}