From 35dafbaf23e71129d02b7e3ac97560ad92ffe6f8 Mon Sep 17 00:00:00 2001 From: Perfloop Agent Date: Tue, 4 Aug 2026 01:18:40 +0000 Subject: [PATCH 1/2] test(BSI): cover run-optimized query results at upstream 438e3566 Upstream-Base: 438e356606d4e651d47a1b8a95b5f2fe08f8c7fd Repair-Of: da7ff3dc6e940b3a0f4c701adcf04b9df33b0132 --- AGENTS.md | 36 ++ BitSliceIndexing/bsi.go | 211 ++++++- BitSliceIndexing/bsi_benchmark_test.go | 292 ++++++++++ BitSliceIndexing/bsi_test.go | 145 +++++ README.md | 6 +- benchmark_test.go | 127 +++++ bitmapcontainer.go | 94 ++-- bitmapcontainer_bench_test.go | 26 + bitmapcontainer_test.go | 107 ++++ clz.go | 19 - clz_compat.go | 37 -- clz_test.go | 48 -- container_test.go | 10 - ctz.go | 21 - ctz_compat.go | 72 --- ctz_test.go | 109 ---- go.mod | 2 +- go.sum | 4 +- iter.go | 11 +- popcnt.go | 13 - popcnt_amd64.s | 103 ---- popcnt_asm.go | 68 --- popcnt_avx2_amd64.go | 67 +++ popcnt_avx2_amd64.s | 352 ++++++++++++ popcnt_avx2_amd64_test.go | 129 +++++ popcnt_bench_test.go | 19 +- popcnt_compat.go | 18 - popcnt_generic.go | 4 +- popcnt_neon_arm64.go | 65 +++ popcnt_neon_arm64.s | 329 +++++++++++ popcnt_neon_arm64_test.go | 123 ++++ popcnt_slices.go | 12 +- popcnt_slices_test.go | 66 --- roaring.go | 38 +- roaring64/BSI_BENCHMARKS.md | 55 ++ roaring64/bsi64.go | 626 ++++++++++++++++++++- roaring64/bsi64_batch_equal_test.go | 228 ++++++++ roaring64/bsi64_batch_equal_values_test.go | 149 +++++ roaring64/bsi64_compare_benchmark_test.go | 304 ++++++++++ roaring64/bsi64_compare_bsi_test.go | 166 ++++++ roaring64/bsi64_get_big_values_test.go | 105 ++++ roaring64/roaring64.go | 1 + roaring64/roaringarray64.go | 4 +- roaring_test.go | 221 +++++++- roaringarray.go | 116 ++++ roaringcow_test.go | 15 - runcontainer.go | 5 +- util.go | 27 +- 48 files changed, 4052 insertions(+), 753 deletions(-) create mode 100644 AGENTS.md create mode 100644 BitSliceIndexing/bsi_benchmark_test.go create mode 100644 bitmapcontainer_bench_test.go delete mode 100644 clz.go delete mode 100644 clz_compat.go delete mode 100644 clz_test.go delete mode 100644 ctz.go delete mode 100644 ctz_compat.go delete mode 100644 ctz_test.go delete mode 100644 popcnt.go delete mode 100644 popcnt_amd64.s delete mode 100644 popcnt_asm.go create mode 100644 popcnt_avx2_amd64.go create mode 100644 popcnt_avx2_amd64.s create mode 100644 popcnt_avx2_amd64_test.go delete mode 100644 popcnt_compat.go create mode 100644 popcnt_neon_arm64.go create mode 100644 popcnt_neon_arm64.s create mode 100644 popcnt_neon_arm64_test.go delete mode 100644 popcnt_slices_test.go create mode 100644 roaring64/BSI_BENCHMARKS.md create mode 100644 roaring64/bsi64_batch_equal_test.go create mode 100644 roaring64/bsi64_batch_equal_values_test.go create mode 100644 roaring64/bsi64_compare_benchmark_test.go create mode 100644 roaring64/bsi64_compare_bsi_test.go create mode 100644 roaring64/bsi64_get_big_values_test.go diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 00000000..b57fb688 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,36 @@ +# Agent Guidance for roaring (Go) + +## AI-Generated Bug Reports About Deserialization + +Many AI-generated bug reports claim that deserialization functions (e.g., `ReadFrom`, `FromBuffer`, `FromUnsafeBytes`, `UnmarshalBinary`, and their `roaring64` equivalents) "trigger bugs", "cause crashes", or "panic" when given malformed or untrusted input. + +**These reports are bogus.** + +The library's documented contract is explicit and consistent across the public API and README: + +- The deserialization functions are memory-safe in the ordinary Go sense: they will not read out of bounds of the slices they are given. +- However, if the input does not conform to the format specification (i.e., it is garbage or from an untrusted source), the resulting bitmap may be in an invalid internal state. Using such a bitmap without validation can lead to panics, incorrect results, or excessive memory consumption. +- **If the source is untrusted, you MUST call `Validate()` on the result and only proceed if it returns `nil`. Only after a successful validation is the bitmap considered safe for use.** Equivalently, `MustReadFrom` calls `ReadFrom` followed by `Validate` and panics on an invalid bitmap. + +Relevant documentation locations: + +- roaring.go: doc comments on `ReadFrom`, `MustReadFrom`, `FromBuffer`, `FromUnsafeBytes`, `UnmarshalBinary`, and `Validate`. +- roaring64/roaring64.go: equivalent `ReadFrom`, `UnmarshalBinary`, and `Validate` for 64-bit bitmaps. +- README.md (the serialization example, around lines 250-272): sample code showing the required validate-after-deserialize pattern for untrusted input, with the explicit comment: "if buf is an untrusted source, you should validate the result". + +The format specification is documented at https://github.com/RoaringBitmap/RoaringFormatSpec. The Go, Java, C and C++ implementations are binary compatible. + +A special note on `FromBuffer` and `FromUnsafeBytes`: these are zero-copy entry points (for advanced users only). The resulting bitmap holds references into the caller-provided byte slice and uses copy-on-write. It is the caller's responsibility to ensure that the input slice is not modified and remains valid for the lifetime of the bitmap (and of any bitmap derived from it). Reports that mutate the backing slice after the fact, or that free/reuse it while the bitmap is still alive, are describing documented misuse, not a library bug. + +The fuzzing harnesses (`FuzzSmat`, `FuzzSerializationBuffer`, `FuzzSerializationStream`, and the corpus under `testdata/fuzz/`) and the property tests exist precisely to ensure the deserializers and the validator behave correctly under adversarial input. A report that treats "deserializing attacker-controlled bytes and then using the result without calling `Validate()`" as a bug in the deserializer is a misunderstanding of the stated API contract. + +When triaging such reports, point to the validation requirement in the function documentation and the README example, and close as "not a bug / user error / documented behavior." + +## Building, Testing, and Conventions + +- This is a pure-Go library; there is no cgo. Use the standard Go toolchain. +- Run the test suite with `go test ./...` (the root package and the `roaring64` subpackage both have extensive tests). +- The `Makefile` exposes `make unconvert`, which runs `go tool unconvert -apply ./...` to remove unnecessary type conversions. Run it before proposing changes that touch type conversions. +- The root package implements 32-bit bitmaps; `roaring64/` implements 64-bit bitmaps and should be kept behaviorally consistent with the root package. +- Architecture-specific files exist for performance (e.g., `popcnt_amd64.s`, `setutil_arm64.s`, with `_generic.go` and `_compat.go` fallbacks). Any change to one path must be mirrored in the generic fallback so all build targets stay correct. +- Keep new code consistent with the surrounding style: match existing naming, error handling, and comment density. Public API changes must update doc comments and, where relevant, the README. diff --git a/BitSliceIndexing/bsi.go b/BitSliceIndexing/bsi.go index 2f0a9f89..6b2ddc56 100644 --- a/BitSliceIndexing/bsi.go +++ b/BitSliceIndexing/bsi.go @@ -4,6 +4,7 @@ import ( "fmt" "math/bits" "runtime" + "sort" "sync" "sync/atomic" @@ -271,7 +272,6 @@ type task struct { op Operation valueOrStart int64 end int64 - values map[int64]struct{} bits *roaring.Bitmap } @@ -743,36 +743,209 @@ func (b *BSI) MarshalBinary() ([][]byte, error) { return data, nil } -// BatchEqual returns a bitmap containing the column IDs where the values are contained within the list of values provided. +// BatchEqual returns a bitmap containing the column IDs where the values are contained +// within the list of values provided. The trie path shares work across values and runs +// on the calling goroutine; on scattered queries that would fan the trie out it falls +// back to a linear existence-bitmap scan, which parallelism splits across goroutines. func (b *BSI) BatchEqual(parallelism int, values []int64) *roaring.Bitmap { + if b.eBM.IsEmpty() || len(values) == 0 { + return roaring.NewBitmap() + } + + bitCount := b.BitCount() - valMap := make(map[int64]struct{}, len(values)) - for i := 0; i < len(values); i++ { - valMap[values[i]] = struct{}{} + // Deduplicate, and drop values that cannot be represented in bitCount + // planes: GetValue can never observe such a value, so it matches no column. + seen := make(map[uint64]struct{}, len(values)) + vals := make([]uint64, 0, len(values)) + for _, v := range values { + u := uint64(v) + if bitCount < 64 && (v < 0 || u >= uint64(1)<= 128 && b.shouldUseParallelScan(vals, bitCount) { + result := b.parallelBatchEqualScan(parallelism, vals) + if b.runOptimized { + result.RunOptimize() + } + return result + } + + result := b.matchTrie(vals, bitCount-1, b.eBM, false) + if b.runOptimized { + result.RunOptimize() + } + return result } -func batchEqual(e *task, batch []uint32, resultsChan chan *roaring.Bitmap, - wg *sync.WaitGroup) { +// shouldUseParallelScan reports whether BatchEqual should skip the match trie in +// favor of a linear existence-bitmap scan. estimateBranchCount caps the trie's +// branch fan-out; past the crossover the trie degenerates into an intermediate- +// bitmap blowup, and only then is the scan's goroutine cost worth paying. The +// scan also needs a large existence bitmap for the parallelism to earn its keep. +func (b *BSI) shouldUseParallelScan(vals []uint64, bitCount int) bool { + return estimateBranchCount(vals, bitCount-1, 64) >= 64 && b.eBM.GetCardinality() >= 100000 +} - defer wg.Done() +// estimateBranchCount bounds the branches the match trie would take on vals over +// planes [0, p], stopping once the count reaches limit. Perfectly contiguous +// ranges collapse to zero. It reads sorted query values only, so the estimate +// costs nothing against the data. +func estimateBranchCount(vals []uint64, p int, limit int) int { + if len(vals) <= 1 || limit <= 0 { + return 0 + } + if vals[len(vals)-1]-vals[0] == uint64(len(vals)-1) { + return 0 + } + if p >= 64 { + p = 63 + } + if p < 0 || (p < 63 && uint64(len(vals)) == uint64(1)< card { + n = int(card) + } + + x := card / uint64(n) + remainder := card - (x * uint64(n)) + iter := b.eBM.ManyIterator() + resultsChan := make(chan *roaring.Bitmap, n) + + var wg sync.WaitGroup + for i := 0; i < n; i++ { + size := x + if i == n-1 { + size += remainder + } + batch := make([]uint32, size) + iter.NextMany(batch) + wg.Add(1) + go func(cols []uint32) { + defer wg.Done() + out := roaring.NewBitmap() + for _, col := range cols { + if v, ok := b.GetValue(uint64(col)); ok { + if _, hit := want[uint64(v)]; hit { + out.Add(col) + } + } } + resultsChan <- out + }(batch) + } + wg.Wait() + close(resultsChan) + + ba := make([]*roaring.Bitmap, 0, n) + for bm := range resultsChan { + ba = append(ba, bm) + } + return roaring.ParOr(0, ba...) +} + +// matchTrie returns the columns in prefix whose bits on planes [0, p] equal the low +// p+1 bits of any value in vals. vals must be unique, sorted, and share identical bits +// above plane p; prefix holds the columns matching those shared upper bits, so results +// stay rooted in the existence bitmap the recursion starts from. owned reports whether +// prefix is a private intermediate the callee may consume in place. Each shared value +// prefix is intersected exactly once and empty intermediates prune the recursion, so +// total work is bounded by the size of the values' bit trie and the data actually +// present, not len(vals) × BitCount(). +func (b *BSI) matchTrie(vals []uint64, p int, prefix *roaring.Bitmap, owned bool) *roaring.Bitmap { + if prefix.IsEmpty() { + if owned { + return prefix } + return roaring.NewBitmap() } - resultsChan <- results + // Either all planes are consumed (exactly one value remains and prefix matched + // every one of its bits), or vals covers every bit pattern of the remaining + // planes, which collapses dense value ranges to a single operation. In both + // cases every column in prefix matches. + if p < 0 || (p < 63 && uint64(len(vals)) == uint64(1)<= 2^7) + // must both return an empty bitmap and must equal a GetValue-derived ground truth. + bsi := NewBSI(100, 0) + assert.Equal(t, 7, bsi.BitCount()) + + // Set some values inside range + bsi.SetValue(10, 42) + bsi.SetValue(20, 99) + + // Ground truth function + getGroundTruth := func(query []int64) *roaring.Bitmap { + expected := roaring.NewBitmap() + valMap := make(map[int64]bool) + for _, q := range query { + valMap[q] = true + } + iter := bsi.GetExistenceBitmap().Iterator() + for iter.HasNext() { + col := iter.Next() + val, ok := bsi.GetValue(uint64(col)) + if ok && valMap[val] { + expected.Add(col) + } + } + return expected + } + + for _, q := range []int64{-5, 200} { + res := bsi.BatchEqual(0, []int64{q}) + assert.True(t, res.IsEmpty()) + + expected := getGroundTruth([]int64{q}) + assert.True(t, expected.IsEmpty()) + assert.True(t, res.Equals(expected)) + } +} + +func TestBatchEqualResultIsolation(t *testing.T) { + bsi := NewDefaultBSI() + bsi.SetValue(10, 42) + bsi.SetValue(20, 100) + + // Get batch equal result + res := bsi.BatchEqual(0, []int64{42}) + assert.True(t, res.Contains(10)) + + // Mutate the returned bitmap + res.Add(999) + res.Remove(10) + + // Assert that the source BSI's internal state (existence bitmap and bit planes) is completely unaffected + assert.False(t, bsi.GetExistenceBitmap().Contains(999)) + assert.True(t, bsi.GetExistenceBitmap().Contains(10)) + + val, ok := bsi.GetValue(10) + assert.True(t, ok) + assert.Equal(t, int64(42), val) + + val, ok = bsi.GetValue(999) + assert.False(t, ok) +} + +func TestBatchEqualConsistentWithGetValue(t *testing.T) { + rg := rand.New(rand.NewSource(42)) + for run := 0; run < 15; run++ { + // Create a randomized BSI + bsi := NewDefaultBSI() + numCols := rg.Intn(1000) + 10 + for col := 0; col < numCols; col++ { + if rg.Float64() < 0.8 { + val := rg.Int63n(500) - 250 // Mix of positive, zero, and negative values + bsi.SetValue(uint64(col), val) + } + } + + // Generate query values (small, medium, and large list sizes to test the hybrid threshold) + querySizes := []int{rg.Intn(10) + 1, rg.Intn(50) + 50, rg.Intn(200) + 100} + for _, querySize := range querySizes { + query := make([]int64, querySize) + for i := range query { + query[i] = rg.Int63n(600) - 300 + } + + // Ground truth + expected := roaring.NewBitmap() + valMap := make(map[int64]bool) + for _, q := range query { + valMap[q] = true + } + iter := bsi.GetExistenceBitmap().Iterator() + for iter.HasNext() { + col := iter.Next() + val, ok := bsi.GetValue(uint64(col)) + if ok && valMap[val] { + expected.Add(col) + } + } + + // Test different parallelism settings + for _, parallelism := range []int{0, 1, 2, 4} { + actual := bsi.BatchEqual(parallelism, query) + if !actual.Equals(expected) { + t.Fatalf("Mismatch in run %d querySize %d parallelism %d. Query: %v. Expected: %v, Got: %v", run, querySize, parallelism, query, expected.ToArray(), actual.ToArray()) + } + } + } + } +} + +// TestBatchEqualExistenceAuthority pins BatchEqual results to the existence +// bitmap. UnmarshalBinary accepts plane data that is not a subset of eBM (the +// checked-in testdata/age fixture is such data), and every read path treats +// eBM as authoritative, so columns present in a plane but absent from eBM must +// never appear in results. +func TestBatchEqualExistenceAuthority(t *testing.T) { + // Synthetic state: column 2 has bits in plane 0 but is absent from eBM. + ebm := roaring.BitmapOf(1) + plane := roaring.BitmapOf(1, 2) + ebmData, err := ebm.MarshalBinary() + if err != nil { + t.Fatal(err) + } + planeData, err := plane.MarshalBinary() + if err != nil { + t.Fatal(err) + } + bsi := NewDefaultBSI() + if err := bsi.UnmarshalBinary([][]byte{ebmData, planeData}); err != nil { + t.Fatal(err) + } + res := bsi.BatchEqual(0, []int64{1}) + assert.True(t, res.Contains(1)) + assert.False(t, res.Contains(2), "column 2 is not in eBM and must not match") + + // The age fixture ships with plane cardinalities above the eBM cardinality; + // results must still be a subset of eBM. + large := setupLargeBSI(t) + if large == nil { + t.Skip("skipping, large BSI setup failed") + } + for _, vals := range [][]int64{{16}, {55, 57}, {0, 1, 2, 3}} { + res := large.BatchEqual(0, vals) + outside := roaring.AndNot(res, large.GetExistenceBitmap()) + assert.True(t, outside.IsEmpty(), "BatchEqual(%v) returned %d columns outside eBM", vals, outside.GetCardinality()) + } +} + +// Benchmarks across query-list shapes: work sharing behaves differently for +// small lists, dense contiguous ranges (which collapse to a few plane +// operations), and scattered values (no range collapse). +func BenchmarkBatchEqualM128(b *testing.B) { benchmarkBatchEqualM(b, 128, 1) } +func BenchmarkBatchEqualM128Scattered(b *testing.B) { benchmarkBatchEqualM(b, 128, 2) } +func BenchmarkBatchEqualM200(b *testing.B) { benchmarkBatchEqualM(b, 200, 1) } + +func benchmarkBatchEqualM(b *testing.B, m int, stride int64) { + bsi := setupLargeBSI(b) + if bsi == nil { + b.Skip("skipping, large BSI setup failed") + } + vals := make([]int64, m) + for i := range vals { + vals[i] = int64(i)*stride + stride - 1 + } + b.ResetTimer() + for i := 0; i < b.N; i++ { + res := bsi.BatchEqual(0, vals) + _ = res + } +} diff --git a/BitSliceIndexing/bsi_test.go b/BitSliceIndexing/bsi_test.go index b28d0044..dee4c83a 100644 --- a/BitSliceIndexing/bsi_test.go +++ b/BitSliceIndexing/bsi_test.go @@ -127,6 +127,19 @@ func setupRandom() *BSI { return bsi } +const runOptimizedQueryResultCardinality = 10_000 + +func setupRunOptimizedQueryBSI(t testing.TB, cardinality int, valueForColumn func(int) int64) *BSI { + t.Helper() + + bsi := NewDefaultBSI() + for columnID := 0; columnID < cardinality; columnID++ { + bsi.SetValue(uint64(columnID), valueForColumn(columnID)) + } + bsi.RunOptimize() + return bsi +} + func TestEQ(t *testing.T) { bsi := setup() eq := bsi.CompareValue(0, EQ, 50, 0, nil) @@ -365,6 +378,91 @@ func TestTransposeWithCounts(t *testing.T) { assert.Equal(t, int64(2), a) } +func TestRunOptimizedBitmapQueryResults(t *testing.T) { + expected := roaring.NewBitmap() + expected.AddRange(0, runOptimizedQueryResultCardinality) + + constantValues := setupRunOptimizedQueryBSI(t, runOptimizedQueryResultCardinality, func(int) int64 { return 1 }) + sequentialValues := setupRunOptimizedQueryBSI(t, runOptimizedQueryResultCardinality, func(columnID int) int64 { return int64(columnID) }) + results := []struct { + name string + result *roaring.Bitmap + }{ + {"BatchEqual", constantValues.BatchEqual(4, []int64{1})}, + {"CompareValue", constantValues.CompareValue(4, EQ, 1, 0, nil)}, + {"Transpose", sequentialValues.IntersectAndTranspose(4, nil)}, + } + + for _, test := range results { + t.Run(test.name, func(t *testing.T) { + assert.True(t, test.result.Equals(expected)) + assert.True(t, test.result.HasRunCompression()) + }) + } +} + +func TestRunOptimizedTransposeWithCountsResult(t *testing.T) { + input := setupRunOptimizedQueryBSI(t, runOptimizedQueryResultCardinality, func(columnID int) int64 { return int64(columnID) }) + result := input.TransposeWithCounts(4, nil) + + assert.Equal(t, uint64(runOptimizedQueryResultCardinality), result.GetCardinality()) + assert.True(t, result.HasRunCompression()) + assert.True(t, result.GetExistenceBitmap().HasRunCompression()) + require.Len(t, result.bA, 1) + assert.True(t, result.bA[0].HasRunCompression()) + for _, columnID := range []uint64{0, 1, runOptimizedQueryResultCardinality - 1} { + value, exists := result.GetValue(columnID) + assert.True(t, exists) + assert.Equal(t, int64(1), value) + } +} + +func TestRunOptimizedTransposeWithCountsEmptyResultIsMutable(t *testing.T) { + input := NewDefaultBSI() + input.SetValue(0, 1) + input.RunOptimize() + + result := input.TransposeWithCounts(4, roaring.NewBitmap()) + assert.Zero(t, result.GetCardinality()) + assert.False(t, result.HasRunCompression()) + + result.SetValue(0, 1) + value, exists := result.GetValue(0) + assert.True(t, exists) + assert.Equal(t, int64(1), value) +} + +func TestRunOptimizedBSISettersCanGrowBitSlices(t *testing.T) { + const expandedValue int64 = 1 << 10 + + t.Run("SetValue", func(t *testing.T) { + bsi := NewDefaultBSI() + bsi.SetValue(0, 1) + bsi.RunOptimize() + bsi.SetValue(1, expandedValue) + + value, exists := bsi.GetValue(1) + assert.True(t, exists) + assert.Equal(t, expandedValue, value) + assert.Equal(t, 11, bsi.BitCount()) + }) + + t.Run("SetMany", func(t *testing.T) { + bsi := NewDefaultBSI() + bsi.SetValue(0, 1) + bsi.RunOptimize() + foundSet := roaring.BitmapOf(1, 2) + bsi.SetMany(foundSet, expandedValue) + + for _, columnID := range []uint64{1, 2} { + value, exists := bsi.GetValue(columnID) + assert.True(t, exists) + assert.Equal(t, expandedValue, value) + } + assert.Equal(t, 11, bsi.BitCount()) + }) +} + func TestRangeAllNegative(t *testing.T) { bsi := setupAllNegative() assert.Equal(t, uint64(100), bsi.GetCardinality()) @@ -502,3 +600,50 @@ func TestTransposeWithCountsNil(t *testing.T) { assert.True(t, ok) assert.Equal(t, int64(2), a) } + +// TestBatchEqualLargeQueryValues drives the scattered-query scan path: a large +// existence bitmap and a 128+ value query push BatchEqual past the crossover, and +// the result is pinned to the GetValue ground truth across parallelism levels, so +// the linear scan and its partitioning stay authoritative (subset of eBM) and exact. +func TestBatchEqualLargeQueryValues(t *testing.T) { + rg := rand.New(rand.NewSource(12345)) + for run := 0; run < 10; run++ { + // Large values (>= 2^20) and a big column set push past the scan crossover. + bsi := NewDefaultBSI() + numCols := rg.Intn(50000) + 120000 + for col := 0; col < numCols; col++ { + if rg.Float64() < 0.8 { + val := rg.Int63n(100000) + 1048500 + bsi.SetValue(uint64(col), val) + } + } + + querySize := rg.Intn(100) + 128 + query := make([]int64, querySize) + for i := range query { + query[i] = rg.Int63n(100100) + 1048500 + } + + // Ground truth: GetValue per existing column. + expected := roaring.NewBitmap() + valMap := make(map[int64]bool) + for _, q := range query { + valMap[q] = true + } + iter := bsi.GetExistenceBitmap().Iterator() + for iter.HasNext() { + col := iter.Next() + val, ok := bsi.GetValue(uint64(col)) + if ok && valMap[val] { + expected.Add(col) + } + } + + for _, parallelism := range []int{0, 1, 2, 4} { + actual := bsi.BatchEqual(parallelism, query) + if !actual.Equals(expected) { + t.Fatalf("mismatch in run %d parallelism %d: expected %v, got %v", run, parallelism, expected.ToArray(), actual.ToArray()) + } + } + } +} diff --git a/README.md b/README.md index eb99b826..7957a463 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # roaring -[![GoDoc](https://godoc.org/github.com/RoaringBitmap/roaring?status.svg)](https://godoc.org/github.com/RoaringBitmap/roaring) [![Go Report Card](https://goreportcard.com/badge/RoaringBitmap/roaring)](https://goreportcard.com/report/github.com/RoaringBitmap/roaring) +[![GoDoc](https://godoc.org/github.com/RoaringBitmap/roaring?status.svg)](https://godoc.org/github.com/RoaringBitmap/roaring) ![Go-CI](https://github.com/RoaringBitmap/roaring/workflows/Go-CI/badge.svg) ![Go-ARM-CI](https://github.com/RoaringBitmap/roaring/workflows/Go-ARM-CI/badge.svg) @@ -421,10 +421,6 @@ The two versions were written independently. https://groups.google.com/g/roaring-bitmaps -## Stars - - -[![Star History Chart](https://api.star-history.com/svg?repos=RoaringBitmap/roaring&type=Date)](https://www.star-history.com/#RoaringBitmap/roaring&Date) ### Further reading diff --git a/benchmark_test.go b/benchmark_test.go index 62153322..f733e1c0 100644 --- a/benchmark_test.go +++ b/benchmark_test.go @@ -1418,3 +1418,130 @@ func BenchmarkFastOrRunContainers(b *testing.B) { } } } + +// BenchmarkBitmapOrBulkMerge exercises the in-place Bitmap.Or bulk merge path +// with a variety of key shapes. The fixtures live in roaring_test.go. +func BenchmarkBitmapOrBulkMerge(b *testing.B) { + fixtures := []struct { + name string + fixture bitmapOrBulkMergeFixture + }{ + {"fresh-interleaved-64", bitmapOrBulkMergeInterleavedFixture(64, false)}, + {"fresh-interleaved-65", bitmapOrBulkMergeInterleavedFixture(65, false)}, + {"fresh-interleaved-1024", bitmapOrBulkMergeInterleavedFixture(1024, false)}, + {"fresh-interleaved-4096", bitmapOrBulkMergeInterleavedFixture(4096, false)}, + {"fresh-append-only-4096", bitmapOrBulkMergeAppendFixture(4096)}, + {"fresh-overlap-4096", bitmapOrBulkMergeOverlapFixture(4096)}, + {"fresh-copy-on-write-interleaved-4096", bitmapOrBulkMergeInterleavedFixture(4096, true)}, + {"fresh-single-interior-4096", bitmapOrBulkMergeSingleInteriorFixture(4096)}, + } + for _, benchmark := range fixtures { + b.Run(benchmark.name, func(b *testing.B) { + b.ReportAllocs() + var cardinality uint64 + b.ResetTimer() + for i := 0; i < b.N; i++ { + fixtureIndex := i & 1 + receiver := benchmark.fixture.lefts[fixtureIndex].Clone() + receiver.Or(benchmark.fixture.rights[fixtureIndex]) + cardinality += receiver.GetCardinality() + } + b.StopTimer() + if cardinality != benchmark.fixture.cardinality*uint64(b.N) { + b.Fatalf("unexpected total cardinality: got %d, want %d", cardinality, benchmark.fixture.cardinality*uint64(b.N)) + } + }) + } +} + +// bitmapOrBulkMergeTailAdjacentFixture builds a receiver holding keys +// [0, 4093] plus 4095, unioned with the single source-only key 4094. The lone +// insert lands just before the receiver's tail. +func bitmapOrBulkMergeTailAdjacentFixture() bitmapOrBulkMergeFixture { + const containers = 4096 + + leftKeys := make([]uint16, 0, containers-1) + for key := 0; key < containers-2; key++ { + leftKeys = append(leftKeys, uint16(key)) + } + leftKeys = append(leftKeys, containers-1) + + return newBitmapOrBulkMergeFixture(leftKeys, []uint16{containers - 2}, false) +} + +func BenchmarkBitmapOrBulkMergeTailAdjacent(b *testing.B) { + fixture := bitmapOrBulkMergeTailAdjacentFixture() + + b.ReportAllocs() + var cardinality uint64 + b.ResetTimer() + for i := 0; i < b.N; i++ { + fixtureIndex := i & 1 + receiver := fixture.lefts[fixtureIndex].Clone() + receiver.Or(fixture.rights[fixtureIndex]) + cardinality += receiver.GetCardinality() + } + b.StopTimer() + if cardinality != fixture.cardinality*uint64(b.N) { + b.Fatalf("unexpected total cardinality: got %d, want %d", cardinality, fixture.cardinality*uint64(b.N)) + } +} + +// bitmapXorBulkMergeFixture builds two variants of a left/right bitmap pair for +// the in-place Bitmap.Xor bulk merge path. shift controls how many low bits the +// right side is offset from the left: shift==0 makes every aligned container +// cancel to empty (stressing the removal path), while a non-zero shift keeps +// them disjoint so aligned pairs survive. +type bitmapXorBulkMergeFixture struct { + lefts [2]*Bitmap + rights [2]*Bitmap +} + +func newBitmapXorBulkMergeFixture(leftKeys, rightKeys []uint16, shift uint16) bitmapXorBulkMergeFixture { + fixture := bitmapXorBulkMergeFixture{} + for variant := range fixture.lefts { + left := NewBitmap() + right := NewBitmap() + base := uint16(variant * 100) + for _, key := range leftKeys { + left.Add(uint32(key)<<16 | uint32(base)) + } + for _, key := range rightKeys { + right.Add(uint32(key)<<16 | uint32(base+shift)) + } + fixture.lefts[variant] = left + fixture.rights[variant] = right + } + return fixture +} + +func BenchmarkBitmapXorBulkMerge(b *testing.B) { + keysAll := bitmapOrBulkMergeKeys(0, 4096, 1) + fixtures := []struct { + name string + fixture bitmapXorBulkMergeFixture + }{ + // disjoint even/odd keys: every source key is an interior insert. + {"interleaved-insert-4096", newBitmapXorBulkMergeFixture( + bitmapOrBulkMergeKeys(0, 4096, 2), bitmapOrBulkMergeKeys(1, 4096, 2), 1)}, + // same keys, same low bits: every aligned container cancels to empty. + {"overlap-cancel-4096", newBitmapXorBulkMergeFixture(keysAll, keysAll, 0)}, + // same keys, disjoint low bits: every aligned pair survives (no shift). + {"overlap-survive-4096", newBitmapXorBulkMergeFixture(keysAll, keysAll, 1)}, + // receiver holds all keys, single interior source-only insert. + {"single-interior-4096", newBitmapXorBulkMergeFixture( + append(append([]uint16{}, bitmapOrBulkMergeKeys(0, 2048, 1)...), + bitmapOrBulkMergeKeys(2049, 2047, 1)...), []uint16{2048}, 1)}, + } + for _, benchmark := range fixtures { + b.Run(benchmark.name, func(b *testing.B) { + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + fixtureIndex := i & 1 + receiver := benchmark.fixture.lefts[fixtureIndex].Clone() + receiver.Xor(benchmark.fixture.rights[fixtureIndex]) + } + }) + } +} diff --git a/bitmapcontainer.go b/bitmapcontainer.go index 99729634..34c68784 100644 --- a/bitmapcontainer.go +++ b/bitmapcontainer.go @@ -51,7 +51,7 @@ func (bc *bitmapContainer) minimum() uint16 { for i := 0; i < len(bc.bitmap); i++ { w := bc.bitmap[i] if w != 0 { - r := countTrailingZeros(w) + r := bits.TrailingZeros64(w) return uint16(r + i*64) } } @@ -69,39 +69,11 @@ func (bc *bitmapContainer) safeMinimum() (uint16, error) { return val, nil } -// i should be non-zero -func clz(i uint64) int { - n := 1 - x := uint32(i >> 32) - if x == 0 { - n += 32 - x = uint32(i) - } - if x>>16 == 0 { - n += 16 - x = x << 16 - } - if x>>24 == 0 { - n += 8 - x = x << 8 - } - if x>>28 == 0 { - n += 4 - x = x << 4 - } - if x>>30 == 0 { - n += 2 - x = x << 2 - } - return n - int(x>>31) -} - func (bc *bitmapContainer) maximum() uint16 { for i := len(bc.bitmap); i > 0; i-- { w := bc.bitmap[i-1] if w != 0 { - r := clz(w) - return uint16((i-1)*64 + 63 - r) + return uint16((i-1)*64 + 63 - bits.LeadingZeros64(w)) } } return uint16(0) @@ -216,7 +188,7 @@ func (bcmi *bitmapContainerManyIterator) nextMany(hs uint32, buf []uint32) int { continue } t := bitset & -bitset - buf[n] = uint32(((base * 64) + int(popcount(t-1)))) | hs + buf[n] = uint32(((base * 64) + bits.OnesCount64(t-1))) | hs n = n + 1 bitset ^= t } @@ -244,7 +216,7 @@ func (bcmi *bitmapContainerManyIterator) nextMany64(hs uint64, buf []uint64) int continue } t := bitset & -bitset - buf[n] = uint64(((base * 64) + int(popcount(t-1)))) | hs + buf[n] = uint64(((base * 64) + bits.OnesCount64(t-1))) | hs n = n + 1 bitset ^= t } @@ -323,16 +295,16 @@ func bitmapEquals(a, b []uint64) bool { } func (bc *bitmapContainer) fillLeastSignificant16bits(x []uint32, i int, mask uint32) int { - // TODO: should be written as optimized assembly + // On amd64 this loop compiles to TZCNT/BLSR; the remaining headroom is + // vectorized decode (cf. CRoaring bitset_extract_setbits_avx2/avx512). pos := i base := mask for k := 0; k < len(bc.bitmap); k++ { bitset := bc.bitmap[k] for bitset != 0 { - t := bitset & -bitset - x[pos] = base + uint32(popcount(t-1)) + x[pos] = base + uint32(bits.TrailingZeros64(bitset)) pos++ - bitset ^= t + bitset &= bitset - 1 } base += 64 } @@ -726,13 +698,13 @@ func (bc *bitmapContainer) rank(x uint16) int { if leftover == 0 { return int(popcntSlice(bc.bitmap[:(uint(x)+1)/64])) } - return int(popcntSlice(bc.bitmap[:(uint(x)+1)/64]) + popcount(bc.bitmap[(uint(x)+1)/64]<<(64-leftover))) + return int(popcntSlice(bc.bitmap[:(uint(x)+1)/64])) + bits.OnesCount64(bc.bitmap[(uint(x)+1)/64]<<(64-leftover)) } func (bc *bitmapContainer) selectInt(x uint16) int { remaining := x for k := 0; k < len(bc.bitmap); k++ { - w := popcount(bc.bitmap[k]) + w := bits.OnesCount64(bc.bitmap[k]) if uint16(w) > remaining { return k*64 + selectBitPosition(bc.bitmap[k], int(remaining)) } @@ -858,11 +830,11 @@ func (bc *bitmapContainer) getCardinalityInRange(start, end uint) int { endword := (end - 1) / 64 const allones = ^uint64(0) if firstword == endword { - return int(popcount(bc.bitmap[firstword] & ((allones << (start % 64)) & (allones >> ((64 - end) & 63))))) + return bits.OnesCount64(bc.bitmap[firstword] & ((allones << (start % 64)) & (allones >> ((64 - end) & 63)))) } - answer := popcount(bc.bitmap[firstword] & (allones << (start % 64))) + answer := uint64(bits.OnesCount64(bc.bitmap[firstword] & (allones << (start % 64)))) answer += popcntSlice(bc.bitmap[firstword+1 : endword]) - answer += popcount(bc.bitmap[endword] & (allones >> ((64 - end) & 63))) + answer += uint64(bits.OnesCount64(bc.bitmap[endword] & (allones >> ((64 - end) & 63)))) return int(answer) } @@ -1000,7 +972,7 @@ func (bc *bitmapContainer) iandNotArray(ac *arrayContainer) container { // are set in the mask and in the current word. mask &= bc.bitmap[wordIdx] bc.bitmap[wordIdx] &= ^mask - bc.cardinality -= int(popcount(mask)) + bc.cardinality -= bits.OnesCount64(mask) } wordIdx = v / 64 @@ -1012,7 +984,7 @@ func (bc *bitmapContainer) iandNotArray(ac *arrayContainer) container { // Flush the last word. mask &= bc.bitmap[wordIdx] bc.bitmap[wordIdx] &= ^mask - bc.cardinality -= int(popcount(mask)) + bc.cardinality -= bits.OnesCount64(mask) if bc.getCardinality() <= arrayDefaultMaxSize { return bc.toArrayContainer() @@ -1152,7 +1124,7 @@ func (bc *bitmapContainer) fillArray(container []uint16) { bitset := bc.bitmap[k] for bitset != 0 { t := bitset & -bitset - container[pos] = uint16((base + int(popcount(t-1)))) + container[pos] = uint16((base + bits.OnesCount64(t-1))) pos = pos + 1 bitset ^= t } @@ -1172,12 +1144,12 @@ func (bc *bitmapContainer) NextSetBit(i uint) int { w := bc.bitmap[x] w = w >> (i % 64) if w != 0 { - return int(i) + countTrailingZeros(w) + return int(i) + bits.TrailingZeros64(w) } x++ for ; x < length; x++ { if bc.bitmap[x] != 0 { - return int(x*64) + countTrailingZeros(bc.bitmap[x]) + return int(x*64) + bits.TrailingZeros64(bc.bitmap[x]) } } return -1 @@ -1195,12 +1167,12 @@ func (bc *bitmapContainer) NextUnsetBit(i uint) int { w = w >> (i % 64) w = ^w if w != 0 { - return int(i) + countTrailingZeros(w) + return int(i) + bits.TrailingZeros64(w) } x++ for ; x < length; x++ { if bc.bitmap[x] != 0xFFFFFFFFFFFFFFFF { - return int(x*64) + countTrailingZeros(^bc.bitmap[x]) + return int(x*64) + bits.TrailingZeros64(^bc.bitmap[x]) } } return int(length * 64) @@ -1231,7 +1203,7 @@ func (bc *bitmapContainer) uPrevSetBit(i uint) int { w = w << (63 - b) if w != 0 { - return int(i) - countLeadingZeros(w) + return int(i) - bits.LeadingZeros64(w) } orig := x x-- @@ -1240,7 +1212,7 @@ func (bc *bitmapContainer) uPrevSetBit(i uint) int { } for ; x < orig; x-- { if bc.bitmap[x] != 0 { - return int((x*64)+63) - countLeadingZeros(bc.bitmap[x]) + return int((x*64)+63) - bits.LeadingZeros64(bc.bitmap[x]) } } return -1 @@ -1259,11 +1231,11 @@ func (bc *bitmapContainer) numberOfRuns() int { for i := 0; i < len(bc.bitmap)-1; i++ { word := nextWord nextWord = bc.bitmap[i+1] - numRuns += popcount((^word)&(word<<1)) + ((word >> 63) &^ nextWord) + numRuns += uint64(bits.OnesCount64((^word)&(word<<1))) + ((word >> 63) &^ nextWord) } word := nextWord - numRuns += popcount((^word) & (word << 1)) + numRuns += uint64(bits.OnesCount64((^word) & (word << 1))) if (word & 0x8000000000000000) != 0 { numRuns++ } @@ -1407,14 +1379,14 @@ func (bc *bitmapContainer) nextAbsentValue(target uint16) int { // if statement - we skip the if we have all ones [1,1,1,1...1] if ^w != 0 { - if countTrailingZeros(w) > 0 { + if bits.TrailingZeros64(w) > 0 { // we have something like [X,Y,Z, 0,0,0]. This means the target bit is zero return int(target) } // other wise something like [X,Y,0,1,1,1..1], where x and y can be either 1 or 0. - trailing := countTrailingOnes(w) + trailing := bits.TrailingZeros64(^w) return int(target) + trailing } @@ -1424,7 +1396,7 @@ func (bc *bitmapContainer) nextAbsentValue(target uint16) int { return int(x * 64) } if ^bc.bitmap[x] != 0 { - trailing := countTrailingOnes(bc.bitmap[x]) + trailing := bits.TrailingZeros64(^bc.bitmap[x]) return int(x*64) + trailing } @@ -1474,25 +1446,25 @@ func (bc *bitmapContainer) previousAbsentValue(target uint16) int { // if statement - we skip if we have all ones [1,1,1,1...1] as no value is absent if ^shifted != 0 { - if countTrailingZeros(shifted) > 0 { + if bits.TrailingZeros64(shifted) > 0 { // we have something like shifted=[X,Y,Z,..., 0,0,0]. This means the target bit is zero return int(target) } // The rotate will rotate the target bit into the leading position. - // We know the target bit is not zero because of the countTrailingZero check above + // We know the target bit is not zero because of the TrailingZeros64 check above // We then shift the target bit out of the way. // Assume a structure like an original structure like [X,Y,Z,..., Target, A, B,C...] // shifted will be [X,Y,Z...Target] // shiftedRotated will be [A,B,C....] - // If countLeadingZeros > 0 then A is zero, if not at least A is 1 return + // If LeadingZeros64 > 0 then A is zero, if not at least A is 1 return // Else count the number of ones's until a 0 shiftedRotated := bits.RotateLeft64(w, int(64-uint(target%64))-1) << 1 - leadingZeros := countLeadingZeros(shiftedRotated) + leadingZeros := bits.LeadingZeros64(shiftedRotated) if leadingZeros > 0 { return int(target) - 1 } - leadingOnes := countLeadingOnes(shiftedRotated) + leadingOnes := bits.LeadingZeros64(^shiftedRotated) if leadingOnes > 0 { return int(target) - leadingOnes - 1 } @@ -1504,7 +1476,7 @@ func (bc *bitmapContainer) previousAbsentValue(target uint16) int { return int(x * 64) } if ^bc.bitmap[x] != 0 { - trailing := countTrailingOnes(bc.bitmap[x]) + trailing := bits.TrailingZeros64(^bc.bitmap[x]) return int(x*64) + trailing } diff --git a/bitmapcontainer_bench_test.go b/bitmapcontainer_bench_test.go new file mode 100644 index 00000000..067bcd51 --- /dev/null +++ b/bitmapcontainer_bench_test.go @@ -0,0 +1,26 @@ +package roaring + +import ( + "math/rand" + "testing" +) + +var sink uint32 + +func BenchmarkBitmapContainerFillLeastSignificant16bits(b *testing.B) { + r := rand.New(rand.NewSource(42)) + bc := newBitmapContainer() + for i := 0; i < 32768; i++ { + val := uint16(r.Intn(65536)) + bc.iadd(val) + } + + x := make([]uint32, 65536) + mask := uint32(123) << 16 + + b.ResetTimer() + for i := 0; i < b.N; i++ { + pos := bc.fillLeastSignificant16bits(x, 0, mask) + sink += x[pos-1] + } +} diff --git a/bitmapcontainer_test.go b/bitmapcontainer_test.go index 51be3dd2..b74fb5c6 100644 --- a/bitmapcontainer_test.go +++ b/bitmapcontainer_test.go @@ -532,3 +532,110 @@ func TestBitmapcontainerOrArrayCardinality(t *testing.T) { assert.Equal(t, 1024, result) }) } + +func TestBitmapContainerFillLeastSignificant16bitsProperties(t *testing.T) { + runTest := func(t *testing.T, vals []uint16, mask uint32) { + bc := newBitmapContainer() + for _, val := range vals { + bc.iadd(val) + } + + cardinality := len(vals) + assert.Equal(t, cardinality, bc.getCardinality()) + + for _, startIdx := range []int{0, 13} { + x := make([]uint32, startIdx+cardinality+10) + // Fill x with a sentinel value to detect out-of-bound writes + const sentinel = 0xDEADC0DE + for j := range x { + x[j] = sentinel + } + + pos := bc.fillLeastSignificant16bits(x, startIdx, mask) + + // Assert return value matches container cardinality contract + assert.Equal(t, startIdx+cardinality, pos) + + // Assert prefix before startIdx is untouched + for j := 0; j < startIdx; j++ { + assert.Equal(t, uint32(sentinel), x[j]) + } + + // Assert suffix after pos is untouched + for j := pos; j < len(x); j++ { + assert.Equal(t, uint32(sentinel), x[j]) + } + + // Assert output contents and order + for j, val := range vals { + expected := mask + uint32(val) + assert.Equal(t, expected, x[startIdx+j], "Mismatch at index %d for val %d", startIdx+j, val) + } + } + } + + t.Run("words >= 2 boundary and bits 63", func(t *testing.T) { + // covers: words >= 2, bit 63 within words, word boundaries 0/63/64/127 + vals := []uint16{ + 0, // boundary 0 + 63, // word 0 bit 63 + 64, // boundary 64 + 127, // word 1 bit 63 + 128, // word 2 boundary 0 (words >= 2) + 191, // word 2 bit 63 + 255, // word 3 bit 63 + 1024, + 1024 + 63, + } + runTest(t, vals, 0xFFFF0000) + runTest(t, vals, 0x12340000) + }) + + t.Run("full container 65536 bits at max mask 0xFFFF0000", func(t *testing.T) { + vals := make([]uint16, 65536) + for i := 0; i < 65536; i++ { + vals[i] = uint16(i) + } + runTest(t, vals, 0xFFFF0000) + }) + + t.Run("dense regime p0.95", func(t *testing.T) { + r := rand.New(rand.NewSource(12345)) + vals := []uint16{} + for i := 0; i < 65536; i++ { + if r.Float64() < 0.95 { + vals = append(vals, uint16(i)) + } + } + runTest(t, vals, 0xABCDE000) + }) + + t.Run("sparse regime", func(t *testing.T) { + r := rand.New(rand.NewSource(54321)) + vals := []uint16{} + for i := 0; i < 65536; i++ { + if r.Float64() < 0.01 { + vals = append(vals, uint16(i)) + } + } + runTest(t, vals, 0x10000000) + }) + + t.Run("random-word regimes", func(t *testing.T) { + r := rand.New(rand.NewSource(999)) + vals := []uint16{} + for word := 0; word < 1024; word++ { + // 30% chance to populate this 64-bit word + if r.Float64() < 0.3 { + // Random word content + w := r.Uint64() + for bit := 0; bit < 64; bit++ { + if (w & (1 << bit)) != 0 { + vals = append(vals, uint16(word*64+bit)) + } + } + } + } + runTest(t, vals, 0x55550000) + }) +} diff --git a/clz.go b/clz.go deleted file mode 100644 index ff49ac89..00000000 --- a/clz.go +++ /dev/null @@ -1,19 +0,0 @@ -//go:build go1.9 -// +build go1.9 - -// "go1.9", from Go version 1.9 onward -// See https://golang.org/pkg/go/build/#hdr-Build_Constraints - -package roaring - -import "math/bits" - -// countLeadingOnes returns the number of leading zeros bits in x; the result is 64 for x == 0. -func countLeadingZeros(x uint64) int { - return bits.LeadingZeros64(x) -} - -// countLeadingOnes returns the number of leading ones bits in x; the result is 0 for x == 0. -func countLeadingOnes(x uint64) int { - return bits.LeadingZeros64(^x) -} diff --git a/clz_compat.go b/clz_compat.go deleted file mode 100644 index 7ee16b4a..00000000 --- a/clz_compat.go +++ /dev/null @@ -1,37 +0,0 @@ -//go:build !go1.9 -// +build !go1.9 - -package roaring - -// LeadingZeroBits returns the number of consecutive most significant zero -// bits of x. -func countLeadingZeros(i uint64) int { - if i == 0 { - return 64 - } - n := 1 - x := uint32(i >> 32) - if x == 0 { - n += 32 - x = uint32(i) - } - if (x >> 16) == 0 { - n += 16 - x <<= 16 - } - if (x >> 24) == 0 { - n += 8 - x <<= 8 - } - if x>>28 == 0 { - n += 4 - x <<= 4 - } - if x>>30 == 0 { - n += 2 - x <<= 2 - - } - n -= int(x >> 31) - return n -} diff --git a/clz_test.go b/clz_test.go deleted file mode 100644 index d9d94002..00000000 --- a/clz_test.go +++ /dev/null @@ -1,48 +0,0 @@ -package roaring - -import ( - "github.com/stretchr/testify/assert" - "testing" -) - -func numberOfLeadingZeros(i uint64) int { - if i == 0 { - return 64 - } - n := 1 - x := uint32(i >> 32) - if x == 0 { - n += 32 - x = uint32(i) - } - if (x >> 16) == 0 { - n += 16 - x <<= 16 - } - if (x >> 24) == 0 { - n += 8 - x <<= 8 - } - if x>>28 == 0 { - n += 4 - x <<= 4 - } - if x>>30 == 0 { - n += 2 - x <<= 2 - - } - n -= int(x >> 31) - return n -} - -func TestCountLeadingZeros072(t *testing.T) { - assert.Equal(t, 64, numberOfLeadingZeros(0)) - assert.Equal(t, 60, numberOfLeadingZeros(8)) - assert.Equal(t, 64-17-1, numberOfLeadingZeros(1<<17)) - assert.Equal(t, 0, numberOfLeadingZeros(0xFFFFFFFFFFFFFFFF)) - assert.Equal(t, 64, countLeadingZeros(0)) - assert.Equal(t, 60, countLeadingZeros(8)) - assert.Equal(t, 64-17-1, countLeadingZeros(1<<17)) - assert.Equal(t, 0, countLeadingZeros(0xFFFFFFFFFFFFFFFF)) -} diff --git a/container_test.go b/container_test.go index 9b89fcb9..7d943d53 100644 --- a/container_test.go +++ b/container_test.go @@ -213,16 +213,6 @@ func TestContainerReverseIterator(t *testing.T) { } func TestRoaringContainer(t *testing.T) { - t.Run("countTrailingZeros", func(t *testing.T) { - x := uint64(0) - o := countTrailingZeros(x) - assert.Equal(t, 64, o) - - x = 1 << 3 - o = countTrailingZeros(x) - assert.Equal(t, 3, o) - }) - t.Run("ArrayShortIterator", func(t *testing.T) { content := []uint16{1, 3, 5, 7, 9} c := makeContainer(content) diff --git a/ctz.go b/ctz.go deleted file mode 100644 index b09bfbd1..00000000 --- a/ctz.go +++ /dev/null @@ -1,21 +0,0 @@ -//go:build go1.9 -// +build go1.9 - -// "go1.9", from Go version 1.9 onward -// See https://golang.org/pkg/go/build/#hdr-Build_Constraints - -package roaring - -import "math/bits" - -// countTrailingZeros returns the number of trailing zero bits in x; the result is 64 for x == 0. -func countTrailingZeros(x uint64) int { - return bits.TrailingZeros64(x) -} - -// countTrailingOnes returns the number of trailing one bits in x -// The result is 64 for x == 9,223,372,036,854,775,807. -// The result is 0 for x == 0. -func countTrailingOnes(x uint64) int { - return bits.TrailingZeros64(^x) -} diff --git a/ctz_compat.go b/ctz_compat.go deleted file mode 100644 index d01df825..00000000 --- a/ctz_compat.go +++ /dev/null @@ -1,72 +0,0 @@ -//go:build !go1.9 -// +build !go1.9 - -package roaring - -// Reuse of portions of go/src/math/big standard lib code -// under this license: -/* -Copyright (c) 2009 The Go Authors. All rights reserved. - -Redistribution and use in source and binary forms, with or without -modification, are permitted provided that the following conditions are -met: - - * Redistributions of source code must retain the above copyright -notice, this list of conditions and the following disclaimer. - * Redistributions in binary form must reproduce the above -copyright notice, this list of conditions and the following disclaimer -in the documentation and/or other materials provided with the -distribution. - * Neither the name of Google Inc. nor the names of its -contributors may be used to endorse or promote products derived from -this software without specific prior written permission. - -THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS -"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT -LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR -A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT -OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, -SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT -LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, -DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY -THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT -(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -*/ - -const deBruijn32 = 0x077CB531 - -var deBruijn32Lookup = []byte{ - 0, 1, 28, 2, 29, 14, 24, 3, 30, 22, 20, 15, 25, 17, 4, 8, - 31, 27, 13, 23, 21, 19, 16, 7, 26, 12, 18, 6, 11, 5, 10, 9, -} - -const deBruijn64 = 0x03f79d71b4ca8b09 - -var deBruijn64Lookup = []byte{ - 0, 1, 56, 2, 57, 49, 28, 3, 61, 58, 42, 50, 38, 29, 17, 4, - 62, 47, 59, 36, 45, 43, 51, 22, 53, 39, 33, 30, 24, 18, 12, 5, - 63, 55, 48, 27, 60, 41, 37, 16, 46, 35, 44, 21, 52, 32, 23, 11, - 54, 26, 40, 15, 34, 20, 31, 10, 25, 14, 19, 9, 13, 8, 7, 6, -} - -// trailingZeroBits returns the number of consecutive least significant zero -// bits of x. -func countTrailingZeros(x uint64) int { - // x & -x leaves only the right-most bit set in the word. Let k be the - // index of that bit. Since only a single bit is set, the value is two - // to the power of k. Multiplying by a power of two is equivalent to - // left shifting, in this case by k bits. The de Bruijn constant is - // such that all six bit, consecutive substrings are distinct. - // Therefore, if we have a left shifted version of this constant we can - // find by how many bits it was shifted by looking at which six bit - // substring ended up at the top of the word. - // (Knuth, volume 4, section 7.3.1) - if x == 0 { - // We have to special case 0; the fomula - // below doesn't work for 0. - return 64 - } - return int(deBruijn64Lookup[((x&-x)*(deBruijn64))>>58]) -} diff --git a/ctz_test.go b/ctz_test.go deleted file mode 100644 index 16d94320..00000000 --- a/ctz_test.go +++ /dev/null @@ -1,109 +0,0 @@ -package roaring - -import ( - "encoding/binary" - "github.com/stretchr/testify/assert" - "math/rand" - "testing" -) - -func TestCountTrailingZeros072(t *testing.T) { - assert.Equal(t, 64, numberOfTrailingZeros(0)) - assert.Equal(t, 3, numberOfTrailingZeros(8)) - assert.Equal(t, 0, numberOfTrailingZeros(7)) - assert.Equal(t, 17, numberOfTrailingZeros(1<<17)) - assert.Equal(t, 17, numberOfTrailingZeros(7<<17)) - assert.Equal(t, 33, numberOfTrailingZeros(255<<33)) - - assert.Equal(t, 64, countTrailingZeros(0)) - assert.Equal(t, 3, countTrailingZeros(8)) - assert.Equal(t, 0, countTrailingZeros(7)) - assert.Equal(t, 17, countTrailingZeros(1<<17)) - assert.Equal(t, 17, countTrailingZeros(7<<17)) - assert.Equal(t, 33, countTrailingZeros(255<<33)) -} - -func getRandomUint64Set(n int) []uint64 { - seed := int64(42) - rand.Seed(seed) - - var buf [8]byte - var o []uint64 - for i := 0; i < n; i++ { - rand.Read(buf[:]) - o = append(o, binary.LittleEndian.Uint64(buf[:])) - } - return o -} - -func getAllOneBitUint64Set() []uint64 { - var o []uint64 - for i := uint(0); i < 64; i++ { - o = append(o, 1<>63)) -} diff --git a/go.mod b/go.mod index 5caef981..34ab80ef 100644 --- a/go.mod +++ b/go.mod @@ -5,7 +5,7 @@ go 1.24.0 toolchain go1.24.4 require ( - github.com/bits-and-blooms/bitset v1.24.2 + github.com/bits-and-blooms/bitset v1.24.4 github.com/google/uuid v1.6.0 github.com/mschoch/smat v0.2.0 github.com/stretchr/testify v1.11.1 diff --git a/go.sum b/go.sum index c045343e..4666bde7 100644 --- a/go.sum +++ b/go.sum @@ -1,5 +1,5 @@ -github.com/bits-and-blooms/bitset v1.24.2 h1:M7/NzVbsytmtfHbumG+K2bremQPMJuqv1JD3vOaFxp0= -github.com/bits-and-blooms/bitset v1.24.2/go.mod h1:7hO7Gc7Pp1vODcmWvKMRA9BNmbv6a/7QIWpPxHddWR8= +github.com/bits-and-blooms/bitset v1.24.4 h1:95H15Og1clikBrKr/DuzMXkQzECs1M6hhoGXLwLQOZE= +github.com/bits-and-blooms/bitset v1.24.4/go.mod h1:7hO7Gc7Pp1vODcmWvKMRA9BNmbv6a/7QIWpPxHddWR8= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI= diff --git a/iter.go b/iter.go index 03f17a6a..4b5dafe4 100644 --- a/iter.go +++ b/iter.go @@ -1,6 +1,9 @@ package roaring -import "iter" +import ( + "iter" + "math/bits" +) // Values returns an iterator that yields the elements of the bitmap in // increasing order. Starting with Go 1.23, users can use a for loop to iterate @@ -99,10 +102,10 @@ func (b *Bitmap) Ranges() iter.Seq2[uint32, uint64] { } for w != 0 { - lo := uint(countTrailingZeros(w)) + lo := uint(bits.TrailingZeros64(w)) bitStart := pos*64 + lo - ones := uint(countTrailingOnes(w >> lo)) + ones := uint(bits.TrailingZeros64(^(w >> lo))) if lo+ones < 64 { if !emit(hs+uint64(bitStart), hs+uint64(bitStart+ones)) { return @@ -115,7 +118,7 @@ func (b *Bitmap) Ranges() iter.Seq2[uint32, uint64] { } var bitEnd uint if pos < length { - trailing := uint(countTrailingOnes(bm[pos])) + trailing := uint(bits.TrailingZeros64(^bm[pos])) bitEnd = pos*64 + trailing w = bm[pos] & ^((uint64(1) << trailing) - 1) } else { diff --git a/popcnt.go b/popcnt.go deleted file mode 100644 index b4980aad..00000000 --- a/popcnt.go +++ /dev/null @@ -1,13 +0,0 @@ -//go:build go1.9 -// +build go1.9 - -// "go1.9", from Go version 1.9 onward -// See https://golang.org/pkg/go/build/#hdr-Build_Constraints - -package roaring - -import "math/bits" - -func popcount(x uint64) uint64 { - return uint64(bits.OnesCount64(x)) -} diff --git a/popcnt_amd64.s b/popcnt_amd64.s deleted file mode 100644 index 1f13fa2e..00000000 --- a/popcnt_amd64.s +++ /dev/null @@ -1,103 +0,0 @@ -// +build amd64,!appengine,!go1.9 - -TEXT ·hasAsm(SB),4,$0-1 -MOVQ $1, AX -CPUID -SHRQ $23, CX -ANDQ $1, CX -MOVB CX, ret+0(FP) -RET - -#define POPCNTQ_DX_DX BYTE $0xf3; BYTE $0x48; BYTE $0x0f; BYTE $0xb8; BYTE $0xd2 - -TEXT ·popcntSliceAsm(SB),4,$0-32 -XORQ AX, AX -MOVQ s+0(FP), SI -MOVQ s_len+8(FP), CX -TESTQ CX, CX -JZ popcntSliceEnd -popcntSliceLoop: -BYTE $0xf3; BYTE $0x48; BYTE $0x0f; BYTE $0xb8; BYTE $0x16 // POPCNTQ (SI), DX -ADDQ DX, AX -ADDQ $8, SI -LOOP popcntSliceLoop -popcntSliceEnd: -MOVQ AX, ret+24(FP) -RET - -TEXT ·popcntMaskSliceAsm(SB),4,$0-56 -XORQ AX, AX -MOVQ s+0(FP), SI -MOVQ s_len+8(FP), CX -TESTQ CX, CX -JZ popcntMaskSliceEnd -MOVQ m+24(FP), DI -popcntMaskSliceLoop: -MOVQ (DI), DX -NOTQ DX -ANDQ (SI), DX -POPCNTQ_DX_DX -ADDQ DX, AX -ADDQ $8, SI -ADDQ $8, DI -LOOP popcntMaskSliceLoop -popcntMaskSliceEnd: -MOVQ AX, ret+48(FP) -RET - -TEXT ·popcntAndSliceAsm(SB),4,$0-56 -XORQ AX, AX -MOVQ s+0(FP), SI -MOVQ s_len+8(FP), CX -TESTQ CX, CX -JZ popcntAndSliceEnd -MOVQ m+24(FP), DI -popcntAndSliceLoop: -MOVQ (DI), DX -ANDQ (SI), DX -POPCNTQ_DX_DX -ADDQ DX, AX -ADDQ $8, SI -ADDQ $8, DI -LOOP popcntAndSliceLoop -popcntAndSliceEnd: -MOVQ AX, ret+48(FP) -RET - -TEXT ·popcntOrSliceAsm(SB),4,$0-56 -XORQ AX, AX -MOVQ s+0(FP), SI -MOVQ s_len+8(FP), CX -TESTQ CX, CX -JZ popcntOrSliceEnd -MOVQ m+24(FP), DI -popcntOrSliceLoop: -MOVQ (DI), DX -ORQ (SI), DX -POPCNTQ_DX_DX -ADDQ DX, AX -ADDQ $8, SI -ADDQ $8, DI -LOOP popcntOrSliceLoop -popcntOrSliceEnd: -MOVQ AX, ret+48(FP) -RET - -TEXT ·popcntXorSliceAsm(SB),4,$0-56 -XORQ AX, AX -MOVQ s+0(FP), SI -MOVQ s_len+8(FP), CX -TESTQ CX, CX -JZ popcntXorSliceEnd -MOVQ m+24(FP), DI -popcntXorSliceLoop: -MOVQ (DI), DX -XORQ (SI), DX -POPCNTQ_DX_DX -ADDQ DX, AX -ADDQ $8, SI -ADDQ $8, DI -LOOP popcntXorSliceLoop -popcntXorSliceEnd: -MOVQ AX, ret+48(FP) -RET diff --git a/popcnt_asm.go b/popcnt_asm.go deleted file mode 100644 index ba2dac91..00000000 --- a/popcnt_asm.go +++ /dev/null @@ -1,68 +0,0 @@ -//go:build amd64 && !appengine && !go1.9 -// +build amd64,!appengine,!go1.9 - -package roaring - -// *** the following functions are defined in popcnt_amd64.s - -//go:noescape - -func hasAsm() bool - -// useAsm is a flag used to select the GO or ASM implementation of the popcnt function -var useAsm = hasAsm() - -//go:noescape - -func popcntSliceAsm(s []uint64) uint64 - -//go:noescape - -func popcntMaskSliceAsm(s, m []uint64) uint64 - -//go:noescape - -func popcntAndSliceAsm(s, m []uint64) uint64 - -//go:noescape - -func popcntOrSliceAsm(s, m []uint64) uint64 - -//go:noescape - -func popcntXorSliceAsm(s, m []uint64) uint64 - -func popcntSlice(s []uint64) uint64 { - if useAsm { - return popcntSliceAsm(s) - } - return popcntSliceGo(s) -} - -func popcntMaskSlice(s, m []uint64) uint64 { - if useAsm { - return popcntMaskSliceAsm(s, m) - } - return popcntMaskSliceGo(s, m) -} - -func popcntAndSlice(s, m []uint64) uint64 { - if useAsm { - return popcntAndSliceAsm(s, m) - } - return popcntAndSliceGo(s, m) -} - -func popcntOrSlice(s, m []uint64) uint64 { - if useAsm { - return popcntOrSliceAsm(s, m) - } - return popcntOrSliceGo(s, m) -} - -func popcntXorSlice(s, m []uint64) uint64 { - if useAsm { - return popcntXorSliceAsm(s, m) - } - return popcntXorSliceGo(s, m) -} diff --git a/popcnt_avx2_amd64.go b/popcnt_avx2_amd64.go new file mode 100644 index 00000000..25b6065c --- /dev/null +++ b/popcnt_avx2_amd64.go @@ -0,0 +1,67 @@ +//go:build amd64 && !appengine +// +build amd64,!appengine + +package roaring + +// The functions below are implemented in popcnt_avx2_amd64.s using AVX2. +// They are only used when the CPU supports AVX2 (see useAVX2); otherwise the +// pure-Go fallbacks in popcnt_slices.go are used. This keeps behavior identical +// on every target: appengine and non-amd64 builds compile popcnt_generic.go +// instead, and amd64 CPUs without AVX2 take the scalar path at runtime. + +//go:noescape +func _hasAVX2() bool + +//go:noescape +func _popcntSliceAVX2(s []uint64) uint64 + +//go:noescape +func _popcntMaskSliceAVX2(s, m []uint64) uint64 + +//go:noescape +func _popcntAndSliceAVX2(s, m []uint64) uint64 + +//go:noescape +func _popcntOrSliceAVX2(s, m []uint64) uint64 + +//go:noescape +func _popcntXorSliceAVX2(s, m []uint64) uint64 + +// useAVX2 selects the AVX2 assembly implementations when the running CPU +// supports AVX2. It is evaluated once at package initialization. +var useAVX2 = _hasAVX2() + +func popcntSlice(s []uint64) uint64 { + if useAVX2 { + return _popcntSliceAVX2(s) + } + return popcntSliceGo(s) +} + +func popcntMaskSlice(s, m []uint64) uint64 { + if useAVX2 { + return _popcntMaskSliceAVX2(s, m) + } + return popcntMaskSliceGo(s, m) +} + +func popcntAndSlice(s, m []uint64) uint64 { + if useAVX2 { + return _popcntAndSliceAVX2(s, m) + } + return popcntAndSliceGo(s, m) +} + +func popcntOrSlice(s, m []uint64) uint64 { + if useAVX2 { + return _popcntOrSliceAVX2(s, m) + } + return popcntOrSliceGo(s, m) +} + +func popcntXorSlice(s, m []uint64) uint64 { + if useAVX2 { + return _popcntXorSliceAVX2(s, m) + } + return popcntXorSliceGo(s, m) +} diff --git a/popcnt_avx2_amd64.s b/popcnt_avx2_amd64.s new file mode 100644 index 00000000..1c61a237 --- /dev/null +++ b/popcnt_avx2_amd64.s @@ -0,0 +1,352 @@ +// +build amd64,!appengine + +//go:build amd64 && !appengine + +#include "textflag.h" + +// AVX2 population-count routines for amd64. They count the set bits across a +// []uint64 (the backing storage of a bitmap container), optionally combining +// each pair of words with a boolean op first: And, Or, Xor, and Mask (s &^ m). +// +// Algorithm (Mula/Lemire VPSHUFB nibble lookup) +// --------------------------------------------- +// AVX2 has no single "popcount a whole vector" instruction, so each byte's +// popcount is taken from a 16-entry lookup table indexed by a 4-bit nibble: +// a byte is split into its low and high nibble, each nibble is looked up (one +// VPSHUFB performs all 32 lookups in a 256-bit register at once), the two +// results are added to give a per-byte popcount, and VPSADBW then sums each +// group of 8 byte-counts into a 64-bit lane total that is accumulated. After +// the loop the four lane totals are summed (HSUM) into a scalar register. +// Each iteration handles 256 bits (4 uint64); a scalar POPCNTQ tail handles +// the trailing len%4 words, so any slice length is counted correctly. +// +// Go assembler conventions used below +// ----------------------------------- +// - Operands are written source(s) first, destination LAST. So +// "VPAND Ymask, Ydata, Ylo" means Ylo = Ydata AND Ymask. +// - Yn are the 256-bit AVX registers; Xn aliases the low 128 bits of Yn. +// - Arguments/results are read from the frame pointer (FP). A Go slice is a +// 3-word header {ptr,len,cap}: s_base+0(FP), s_len+8(FP); a second slice +// argument starts at +24(FP). The uint64 result slot follows the args +// (e.g. ret+24(FP) for one slice arg, ret+48(FP) for two). +// - Every routine is a leaf (makes no calls): NOSPLIT with a $0 local frame. +// - Loads/stores use VMOVDQU (unaligned): container slices are only 8-byte +// aligned, not 32. VZEROUPPER precedes every RET to avoid the AVX<->SSE +// transition penalty in any non-VEX SSE code that runs afterwards. + +// lutmask is a 64-byte read-only blob holding two constants used by every +// routine: +// bytes 0..31 - the nibble popcount table, i.e. table[i] = number of set +// bits in the 4-bit value i. VPSHUFB indexes within each +// 128-bit lane independently, so the 16-entry table is stored +// twice (once per lane). Read low-byte-first, the first qword +// 0x0302020102010100 is the bytes {0,1,1,2,1,2,2,3} for +// nibbles 0..7, and 0x0403030203020201 is {1,2,2,3,2,3,3,4} +// for nibbles 8..15. +// bytes 32..63 - 0x0F in every byte: a mask that isolates the low nibble of +// each byte. +// RODATA|NOPTR marks it read-only and pointer-free (so the GC ignores it). +DATA lutmask<>+0(SB)/8, $0x0302020102010100 +DATA lutmask<>+8(SB)/8, $0x0403030203020201 +DATA lutmask<>+16(SB)/8, $0x0302020102010100 +DATA lutmask<>+24(SB)/8, $0x0403030203020201 +DATA lutmask<>+32(SB)/8, $0x0f0f0f0f0f0f0f0f +DATA lutmask<>+40(SB)/8, $0x0f0f0f0f0f0f0f0f +DATA lutmask<>+48(SB)/8, $0x0f0f0f0f0f0f0f0f +DATA lutmask<>+56(SB)/8, $0x0f0f0f0f0f0f0f0f +GLOBL lutmask<>(SB), RODATA|NOPTR, $64 + +// Register aliases. Ylut/Ymask/Yzero are constants set up once per call (see +// SETUP); Yacc is the running accumulator of lane totals; Ydata/Yb hold the +// current input vector(s); Ylo/Yhi/Yc1/Yc2 are scratch used by COUNTBLOCK. +#define Ylut Y0 +#define Ymask Y1 +#define Yzero Y2 +#define Yacc Y3 +#define Ydata Y4 +#define Yb Y5 +#define Ylo Y6 +#define Yhi Y7 +#define Yc1 Y8 +#define Yc2 Y9 + +// COUNTBLOCK folds the popcount of the 32 bytes currently in Ydata into the +// accumulator Yacc. Line by line: +// VPAND Ymask,Ydata,Ylo : Ylo = low nibble of every byte +// VPSRLW $4,Ydata,Yhi : shift each 16-bit lane right by 4... +// VPAND Ymask,Yhi,Yhi : ...then mask, leaving the high nibble of each byte +// VPSHUFB Ylo,Ylut,Yc1 : Yc1[b] = popcount(low nibble of byte b) +// VPSHUFB Yhi,Ylut,Yc2 : Yc2[b] = popcount(high nibble of byte b) +// VPADDB Yc2,Yc1,Yc1 : Yc1[b] = popcount(byte b) (0..8 each) +// VPSADBW Yzero,Yc1,Yc1 : sum each group of 8 bytes -> 4 lane totals (0..512) +// VPADDQ Yc1,Yacc,Yacc : add the 4 lane totals into the accumulator +// Per-byte counts max at 8 and lane totals at 512, so accumulating across the +// whole loop never overflows the 64-bit lanes. +#define COUNTBLOCK \ + VPAND Ymask, Ydata, Ylo \ + VPSRLW $4, Ydata, Yhi \ + VPAND Ymask, Yhi, Yhi \ + VPSHUFB Ylo, Ylut, Yc1 \ + VPSHUFB Yhi, Ylut, Yc2 \ + VPADDB Yc2, Yc1, Yc1 \ + VPSADBW Yzero, Yc1, Yc1 \ + VPADDQ Yc1, Yacc, Yacc + +// SETUP loads the lookup table and nibble mask and zeroes Yzero (the VPSADBW +// addend) and Yacc (the accumulator). Run once at the top of each routine. +#define SETUP \ + VMOVDQU lutmask<>+0(SB), Ylut \ + VMOVDQU lutmask<>+32(SB), Ymask \ + VPXOR Yzero, Yzero, Yzero \ + VPXOR Yacc, Yacc, Yacc + +// HSUM reduces Yacc's four 64-bit lane totals to a single sum in AX. X3 is the +// low 128 bits of Yacc (Y3); VEXTRACTI128 pulls the high 128 bits into X5, the +// two halves are added (giving two qwords in X3), and those two qwords are then +// added into AX. +#define HSUM \ + VEXTRACTI128 $1, Yacc, X5 \ + VPADDQ X5, X3, X3 \ + VPEXTRQ $1, X3, DX \ + MOVQ X3, R9 \ + ADDQ R9, AX \ + ADDQ DX, AX + +// func _popcntSliceAVX2(s []uint64) uint64 +// Returns the total number of set bits in s. This is the canonical routine; +// the And/Or/Xor/Mask variants below share its structure and differ only by +// the boolean op applied before counting. +TEXT ·_popcntSliceAVX2(SB), NOSPLIT, $0-32 + MOVQ s_base+0(FP), SI // SI = &s[0] + MOVQ s_len+8(FP), CX // CX = len(s), in 64-bit words + XORQ AX, AX // AX = running result + SETUP // load table/mask; zero Yzero and Yacc + MOVQ CX, R8 + SHRQ $2, R8 // R8 = len/4 = number of full 256-bit blocks + TESTQ R8, R8 + JZ slicetail // fewer than 4 words: skip the vector loop +sliceloop: + VMOVDQU (SI), Ydata // load 4 words (32 bytes) + COUNTBLOCK // Yacc += popcount(those 32 bytes) + ADDQ $32, SI // advance to the next block + DECQ R8 + JNZ sliceloop + HSUM // AX += sum of Yacc's lane totals +slicetail: + ANDQ $3, CX // CX = len % 4 = leftover words (0..3) + TESTQ CX, CX + JZ slicedone +slicetailloop: + MOVQ (SI), DX + POPCNTQ DX, DX // scalar popcount of one word + ADDQ DX, AX + ADDQ $8, SI + DECQ CX + JNZ slicetailloop +slicedone: + VZEROUPPER // clear upper YMM state before returning + MOVQ AX, ret+24(FP) // return AX + RET + +// func _popcntAndSliceAVX2(s, m []uint64) uint64 +// Returns the sum of popcount(s[i] & m[i]). Mirrors _popcntSliceAVX2 but loads +// a vector from each of s and m and ANDs them before counting. s and m are +// assumed to have equal length. +TEXT ·_popcntAndSliceAVX2(SB), NOSPLIT, $0-56 + MOVQ s_base+0(FP), SI // SI = &s[0] + MOVQ m_base+24(FP), DI // DI = &m[0] + MOVQ s_len+8(FP), CX // CX = len + XORQ AX, AX + SETUP + MOVQ CX, R8 + SHRQ $2, R8 + TESTQ R8, R8 + JZ andtail +andloop: + VMOVDQU (SI), Ydata + VMOVDQU (DI), Yb + VPAND Yb, Ydata, Ydata // Ydata = s & m + COUNTBLOCK + ADDQ $32, SI + ADDQ $32, DI + DECQ R8 + JNZ andloop + HSUM +andtail: + ANDQ $3, CX + TESTQ CX, CX + JZ anddone +andtailloop: + MOVQ (SI), DX + ANDQ (DI), DX // s & m, one word + POPCNTQ DX, DX + ADDQ DX, AX + ADDQ $8, SI + ADDQ $8, DI + DECQ CX + JNZ andtailloop +anddone: + VZEROUPPER + MOVQ AX, ret+48(FP) // +48: result follows two 24-byte slice headers + RET + +// func _popcntOrSliceAVX2(s, m []uint64) uint64 +// Returns the sum of popcount(s[i] | m[i]); see _popcntAndSliceAVX2 for the +// shared structure. +TEXT ·_popcntOrSliceAVX2(SB), NOSPLIT, $0-56 + MOVQ s_base+0(FP), SI + MOVQ m_base+24(FP), DI + MOVQ s_len+8(FP), CX + XORQ AX, AX + SETUP + MOVQ CX, R8 + SHRQ $2, R8 + TESTQ R8, R8 + JZ ortail +orloop: + VMOVDQU (SI), Ydata + VMOVDQU (DI), Yb + VPOR Yb, Ydata, Ydata // Ydata = s | m + COUNTBLOCK + ADDQ $32, SI + ADDQ $32, DI + DECQ R8 + JNZ orloop + HSUM +ortail: + ANDQ $3, CX + TESTQ CX, CX + JZ ordone +ortailloop: + MOVQ (SI), DX + ORQ (DI), DX // s | m, one word + POPCNTQ DX, DX + ADDQ DX, AX + ADDQ $8, SI + ADDQ $8, DI + DECQ CX + JNZ ortailloop +ordone: + VZEROUPPER + MOVQ AX, ret+48(FP) + RET + +// func _popcntXorSliceAVX2(s, m []uint64) uint64 +// Returns the sum of popcount(s[i] ^ m[i]); see _popcntAndSliceAVX2 for the +// shared structure. +TEXT ·_popcntXorSliceAVX2(SB), NOSPLIT, $0-56 + MOVQ s_base+0(FP), SI + MOVQ m_base+24(FP), DI + MOVQ s_len+8(FP), CX + XORQ AX, AX + SETUP + MOVQ CX, R8 + SHRQ $2, R8 + TESTQ R8, R8 + JZ xortail +xorloop: + VMOVDQU (SI), Ydata + VMOVDQU (DI), Yb + VPXOR Yb, Ydata, Ydata // Ydata = s ^ m + COUNTBLOCK + ADDQ $32, SI + ADDQ $32, DI + DECQ R8 + JNZ xorloop + HSUM +xortail: + ANDQ $3, CX + TESTQ CX, CX + JZ xordone +xortailloop: + MOVQ (SI), DX + XORQ (DI), DX // s ^ m, one word + POPCNTQ DX, DX + ADDQ DX, AX + ADDQ $8, SI + ADDQ $8, DI + DECQ CX + JNZ xortailloop +xordone: + VZEROUPPER + MOVQ AX, ret+48(FP) + RET + +// func _popcntMaskSliceAVX2(s, m []uint64) uint64 +// Returns the sum of popcount(s[i] &^ m[i]) == popcount(s & ~m). Same structure +// as _popcntAndSliceAVX2; the combine is VPANDN, which computes (NOT first) AND +// second, i.e. VPANDN Ydata, Yb, Ydata -> Ydata = (NOT Yb) AND Ydata = s &^ m. +TEXT ·_popcntMaskSliceAVX2(SB), NOSPLIT, $0-56 + MOVQ s_base+0(FP), SI + MOVQ m_base+24(FP), DI + MOVQ s_len+8(FP), CX + XORQ AX, AX + SETUP + MOVQ CX, R8 + SHRQ $2, R8 + TESTQ R8, R8 + JZ masktail +maskloop: + VMOVDQU (SI), Ydata + VMOVDQU (DI), Yb + VPANDN Ydata, Yb, Ydata // Ydata = s &^ m (= (NOT m) AND s) + COUNTBLOCK + ADDQ $32, SI + ADDQ $32, DI + DECQ R8 + JNZ maskloop + HSUM +masktail: + ANDQ $3, CX + TESTQ CX, CX + JZ maskdone +masktailloop: + MOVQ (DI), R10 + NOTQ R10 // ~m + MOVQ (SI), DX + ANDQ R10, DX // s &^ m = s & ~m, one word + POPCNTQ DX, DX + ADDQ DX, AX + ADDQ $8, SI + ADDQ $8, DI + DECQ CX + JNZ masktailloop +maskdone: + VZEROUPPER + MOVQ AX, ret+48(FP) + RET + +// func _hasAVX2() bool +// Reports whether the CPU supports AVX2 and the OS has enabled the wide (YMM) +// register state. All three checks must pass; otherwise the Go wrappers fall +// back to the scalar implementation. Note CPUID clobbers AX/BX/CX/DX. +TEXT ·_hasAVX2(SB), NOSPLIT, $0-1 + // CPUID leaf 1: require OSXSAVE (ECX bit 27) and AVX (ECX bit 28). Both must + // be set, so mask and compare against the combined bit pattern. + MOVL $1, AX + XORL CX, CX + CPUID + ANDL $0x18000000, CX + CMPL CX, $0x18000000 + JNE noavx2 + // XGETBV(0): the OS must have enabled saving of SSE and AVX/YMM state, i.e. + // XCR0 bits 1 and 2. Without this the YMM registers would be corrupted + // across a context switch even though the CPU supports the instructions. + XORL CX, CX + XGETBV + ANDL $0x6, AX + CMPL AX, $0x6 + JNE noavx2 + // CPUID leaf 7, sub-leaf 0: require AVX2 itself (EBX bit 5). The sub-leaf is + // selected via ECX, which must be 0. + MOVL $7, AX + XORL CX, CX + CPUID + ANDL $0x20, BX + CMPL BX, $0x20 + JNE noavx2 + MOVB $1, ret+0(FP) + RET +noavx2: + MOVB $0, ret+0(FP) + RET diff --git a/popcnt_avx2_amd64_test.go b/popcnt_avx2_amd64_test.go new file mode 100644 index 00000000..c0e91b0f --- /dev/null +++ b/popcnt_avx2_amd64_test.go @@ -0,0 +1,129 @@ +//go:build amd64 && !appengine +// +build amd64,!appengine + +package roaring + +import ( + "math/rand" + "testing" + + "github.com/stretchr/testify/assert" +) + +// edge lengths exercise the AVX2 main loop (multiples of 4) and the scalar +// POPCNTQ tail (len % 4 != 0), including the empty and sub-block cases. +var avx2TestLengths = []int{0, 1, 2, 3, 4, 5, 7, 8, 15, 16, 17, 31, 63, 64, 65, 1023, 1024, 1025} + +func randomUint64Slice(r *rand.Rand, n int) []uint64 { + s := make([]uint64, n) + for i := range s { + // mix fully-random words with sparse and dense ones to vary popcounts. + switch i % 4 { + case 0: + s[i] = r.Uint64() + case 1: + s[i] = 0 + case 2: + s[i] = ^uint64(0) + default: + s[i] = r.Uint64() & r.Uint64() + } + } + return s +} + +func benchPopcntPair(b *testing.B, avx2 bool, fn func(s, m []uint64) uint64) { + if avx2 && !useAVX2 { + b.Skip("AVX2 not available") + } + r := rand.New(rand.NewSource(1)) + s := randomUint64Slice(r, 1024) + m := randomUint64Slice(r, 1024) + var sink uint64 + b.ResetTimer() + for i := 0; i < b.N; i++ { + sink += fn(s, m) + } + _ = sink +} + +func BenchmarkPopcntAndSlice1024AVX2(b *testing.B) { + benchPopcntPair(b, true, _popcntAndSliceAVX2) +} + +func BenchmarkPopcntAndSlice1024Go(b *testing.B) { + benchPopcntPair(b, false, popcntAndSliceGo) +} + +func BenchmarkPopcntSlice1024AVX2(b *testing.B) { + if !useAVX2 { + b.Skip("AVX2 not available") + } + r := rand.New(rand.NewSource(1)) + s := randomUint64Slice(r, 1024) + var sink uint64 + b.ResetTimer() + for i := 0; i < b.N; i++ { + sink += _popcntSliceAVX2(s) + } + _ = sink +} + +func BenchmarkPopcntSlice1024Go(b *testing.B) { + r := rand.New(rand.NewSource(1)) + s := randomUint64Slice(r, 1024) + var sink uint64 + b.ResetTimer() + for i := 0; i < b.N; i++ { + sink += popcntSliceGo(s) + } + _ = sink +} + +func TestAVX2PopcntDispatch(t *testing.T) { + // Verify the runtime dispatch wrappers agree with the Go reference both + // when AVX2 is selected and when the scalar fallback is forced. + saved := useAVX2 + defer func() { useAVX2 = saved }() + + r := rand.New(rand.NewSource(7)) + for _, on := range []bool{false, true} { + if on && !saved { + continue // CPU has no AVX2; only the fallback exists + } + useAVX2 = on + for _, n := range avx2TestLengths { + s := randomUint64Slice(r, n) + m := randomUint64Slice(r, n) + assert.Equalf(t, popcntSliceGo(s), popcntSlice(s), "popcntSlice avx2=%v len=%d", on, n) + assert.Equalf(t, popcntAndSliceGo(s, m), popcntAndSlice(s, m), "popcntAndSlice avx2=%v len=%d", on, n) + assert.Equalf(t, popcntOrSliceGo(s, m), popcntOrSlice(s, m), "popcntOrSlice avx2=%v len=%d", on, n) + assert.Equalf(t, popcntXorSliceGo(s, m), popcntXorSlice(s, m), "popcntXorSlice avx2=%v len=%d", on, n) + assert.Equalf(t, popcntMaskSliceGo(s, m), popcntMaskSlice(s, m), "popcntMaskSlice avx2=%v len=%d", on, n) + } + } +} + +func TestAVX2PopcntDifferential(t *testing.T) { + if !useAVX2 { + t.Skip("AVX2 not available on this CPU") + } + r := rand.New(rand.NewSource(42)) + for _, n := range avx2TestLengths { + for iter := 0; iter < 64; iter++ { + s := randomUint64Slice(r, n) + m := randomUint64Slice(r, n) + + assert.Equalf(t, popcntSliceGo(s), _popcntSliceAVX2(s), + "popcntSlice len=%d", n) + assert.Equalf(t, popcntAndSliceGo(s, m), _popcntAndSliceAVX2(s, m), + "popcntAndSlice len=%d", n) + assert.Equalf(t, popcntOrSliceGo(s, m), _popcntOrSliceAVX2(s, m), + "popcntOrSlice len=%d", n) + assert.Equalf(t, popcntXorSliceGo(s, m), _popcntXorSliceAVX2(s, m), + "popcntXorSlice len=%d", n) + assert.Equalf(t, popcntMaskSliceGo(s, m), _popcntMaskSliceAVX2(s, m), + "popcntMaskSlice len=%d", n) + } + } +} diff --git a/popcnt_bench_test.go b/popcnt_bench_test.go index 25a99577..7f3e031e 100644 --- a/popcnt_bench_test.go +++ b/popcnt_bench_test.go @@ -1,6 +1,23 @@ package roaring -import "testing" +import ( + "encoding/binary" + "math/rand" + "testing" +) + +func getRandomUint64Set(n int) []uint64 { + seed := int64(42) + rand.Seed(seed) + + var buf [8]byte + var o []uint64 + for i := 0; i < n; i++ { + rand.Read(buf[:]) + o = append(o, binary.LittleEndian.Uint64(buf[:])) + } + return o +} func BenchmarkPopcount(b *testing.B) { b.StopTimer() diff --git a/popcnt_compat.go b/popcnt_compat.go deleted file mode 100644 index 5933e52f..00000000 --- a/popcnt_compat.go +++ /dev/null @@ -1,18 +0,0 @@ -//go:build !go1.9 -// +build !go1.9 - -package roaring - -// bit population count, take from -// https://code.google.com/p/go/issues/detail?id=4988#c11 -// credit: https://code.google.com/u/arnehormann/ -// credit: https://play.golang.org/p/U7SogJ7psJ -// credit: http://graphics.stanford.edu/~seander/bithacks.html#CountBitsSetParallel -func popcount(x uint64) uint64 { - x -= (x >> 1) & 0x5555555555555555 - x = (x>>2)&0x3333333333333333 + x&0x3333333333333333 - x += x >> 4 - x &= 0x0f0f0f0f0f0f0f0f - x *= 0x0101010101010101 - return x >> 56 -} diff --git a/popcnt_generic.go b/popcnt_generic.go index 4ae6d5af..63235fe4 100644 --- a/popcnt_generic.go +++ b/popcnt_generic.go @@ -1,5 +1,5 @@ -//go:build !amd64 || appengine || go1.9 -// +build !amd64 appengine go1.9 +//go:build (!amd64 && !arm64) || appengine +// +build !amd64,!arm64 appengine package roaring diff --git a/popcnt_neon_arm64.go b/popcnt_neon_arm64.go new file mode 100644 index 00000000..c348688f --- /dev/null +++ b/popcnt_neon_arm64.go @@ -0,0 +1,65 @@ +//go:build arm64 && !appengine +// +build arm64,!appengine + +package roaring + +// The functions below are implemented in popcnt_neon_arm64.s using NEON +// (Advanced SIMD). NEON is mandatory in the ARMv8-A baseline that every arm64 +// CPU implements, so — unlike the amd64 AVX2 path, which is gated on a runtime +// _hasAVX2 check — these routines are always used on arm64. The pure-Go +// fallbacks in popcnt_slices.go remain in use on other architectures and on +// appengine builds, which compile popcnt_generic.go instead. + +//go:noescape +func _popcntSliceNEON(s []uint64) uint64 + +//go:noescape +func _popcntMaskSliceNEON(s, m []uint64) uint64 + +//go:noescape +func _popcntAndSliceNEON(s, m []uint64) uint64 + +//go:noescape +func _popcntOrSliceNEON(s, m []uint64) uint64 + +//go:noescape +func _popcntXorSliceNEON(s, m []uint64) uint64 + +// useNEON is always true on arm64; it exists so tests can force the scalar +// fallback path and to mirror the amd64 dispatch structure. +var useNEON = true + +func popcntSlice(s []uint64) uint64 { + if useNEON { + return _popcntSliceNEON(s) + } + return popcntSliceGo(s) +} + +func popcntMaskSlice(s, m []uint64) uint64 { + if useNEON { + return _popcntMaskSliceNEON(s, m) + } + return popcntMaskSliceGo(s, m) +} + +func popcntAndSlice(s, m []uint64) uint64 { + if useNEON { + return _popcntAndSliceNEON(s, m) + } + return popcntAndSliceGo(s, m) +} + +func popcntOrSlice(s, m []uint64) uint64 { + if useNEON { + return _popcntOrSliceNEON(s, m) + } + return popcntOrSliceGo(s, m) +} + +func popcntXorSlice(s, m []uint64) uint64 { + if useNEON { + return _popcntXorSliceNEON(s, m) + } + return popcntXorSliceGo(s, m) +} diff --git a/popcnt_neon_arm64.s b/popcnt_neon_arm64.s new file mode 100644 index 00000000..9a4d651e --- /dev/null +++ b/popcnt_neon_arm64.s @@ -0,0 +1,329 @@ +// +build arm64,!appengine + +//go:build arm64 && !appengine + +#include "textflag.h" + +// NEON (Advanced SIMD) population-count routines for arm64. They count the set +// bits across a []uint64 (the backing storage of a bitmap container), optionally +// combining each pair of words with a boolean op first: And, Or, Xor, and Mask +// (s &^ m). NEON is mandatory in the ARMv8-A baseline that every arm64 CPU +// implements, so unlike the amd64 AVX2 code there is no runtime feature check: +// these routines are always used on arm64 (see popcnt_neon_arm64.go). +// +// Algorithm (VCNT byte popcount + widening accumulation) +// ------------------------------------------------------ +// arm64 has a dedicated per-byte popcount instruction, VCNT, which replaces each +// byte of a 128-bit register with the popcount (0..8) of the input byte. Turning +// those per-byte counts into a running total means widening and accumulating, +// and the loop is shaped to keep the arithmetic units busy: +// - each iteration loads four 16-byte vectors (64 bytes = 8 words) and VCNTs +// them independently, then sums the four with byte-wise VADD. Four counts of +// at most 8 sum to at most 32, so no byte lane overflows. +// - Go's arm64 assembler exposes no pairwise-add-long (UADALP), so the summed +// bytes are folded into 16-bit lanes with add-wide: VUADDW takes the low 8 +// bytes into partial accumulator V16 and VUADDW2 the high 8 into V18. Two +// separate accumulators keep those adds off each other's dependency chain. +// - a 16-bit lane would eventually overflow, so every INNERMAX iterations the +// partials are drained (widened again) into a 4x32-bit accumulator (V17) +// that cannot realistically overflow, and the partials are re-zeroed. +// - at the end VUADDLV sums the four 32-bit lanes into a scalar. +// A scalar-width NEON tail (VCNT + VUADDLV on one 64-bit word at a time) mops up +// the trailing len%8 words, so any slice length is counted correctly. +// +// Go assembler conventions used below +// ----------------------------------- +// - Operands are written source(s) first, destination LAST. So +// "VAND V4.B16, V0.B16, V0.B16" means V0 = V0 AND V4. +// - Vn.B16/H8/H4/S4/D1 name the arrangement (element size x count) an +// instruction operates on: B16 = 16 bytes, H8/H4 = 8/4 halfwords, S4 = 4 +// words, D1 = 1 doubleword. The same physical register is viewed either way. +// - VLD1.P post-increments the pointer register by the number of bytes loaded. +// - Arguments/results are read from the frame pointer (FP). A Go slice is a +// 3-word header {ptr,len,cap}: s_base+0(FP), s_len+8(FP); a second slice +// argument starts at +24(FP). The uint64 result slot follows the args +// (ret+24(FP) for one slice arg, ret+48(FP) for two). +// - Every routine is a leaf (makes no calls): NOSPLIT with a $0 local frame. + +// INNERMAX bounds how many 64-byte iterations fold into the 16-bit partial +// accumulators before they are drained into the wider one. Each iteration adds +// at most 32 (four byte-popcounts of at most 8) to a 16-bit lane, and +// 1024*32 = 32768 stays well under the 65535 lane limit. +#define INNERMAX $1024 + +// FOLD4 assumes 64 bytes of input (post-combine) sit in V0..V3 and folds their +// popcount into the partial accumulators V16/V18. VADD sums the four VCNT +// results byte-wise (each lane 0..32); VUADDW/VUADDW2 then widen the low/high +// halves into the two 16-bit accumulators. +#define FOLD4 \ + VCNT V0.B16, V0.B16 \ + VCNT V1.B16, V1.B16 \ + VCNT V2.B16, V2.B16 \ + VCNT V3.B16, V3.B16 \ + VADD V1.B16, V0.B16, V0.B16 \ + VADD V3.B16, V2.B16, V2.B16 \ + VADD V2.B16, V0.B16, V0.B16 \ + VUADDW V0.B8, V16.H8, V16.H8 \ + VUADDW2 V0.B16, V18.H8, V18.H8 + +// ZEROPART re-zeroes the two 16-bit partial accumulators at the start of each +// INNERMAX batch. +#define ZEROPART \ + VEOR V16.B16, V16.B16, V16.B16 \ + VEOR V18.B16, V18.B16, V18.B16 + +// DRAIN widens the 16-bit partials V16/V18 into the 32-bit accumulator V17 +// (VUADDW low four halfwords, VUADDW2 high four, for each) and re-zeroes them. +#define DRAIN \ + VUADDW V16.H4, V17.S4, V17.S4 \ + VUADDW2 V16.H8, V17.S4, V17.S4 \ + VUADDW V18.H4, V17.S4, V17.S4 \ + VUADDW2 V18.H8, V17.S4, V17.S4 \ + ZEROPART + +// REDUCE sums the four 32-bit lanes of V17 into a scalar and adds it to R2 (the +// running result). VUADDLV over .S4 yields a 64-bit sum; VMOV lifts it to a GPR. +#define REDUCE \ + VUADDLV V17.S4, V0 \ + VMOV V0.D[0], R4 \ + ADD R4, R2, R2 + +// TAILWORD popcounts the single 64-bit word already loaded into V0's low lane +// and adds it to R2: VCNT counts each of the 8 bytes, VUADDLV sums them. +#define TAILWORD \ + VCNT V0.B8, V0.B8 \ + VUADDLV V0.B8, V0 \ + VMOV V0.S[0], R4 \ + ADD R4, R2, R2 + +// func _popcntSliceNEON(s []uint64) uint64 +// Returns the total number of set bits in s. This is the canonical routine; the +// And/Or/Xor/Mask variants below share its structure and differ only by the +// boolean op applied to the two inputs before counting. +TEXT ·_popcntSliceNEON(SB), NOSPLIT, $0-32 + MOVD s_base+0(FP), R0 // R0 = &s[0] + MOVD s_len+8(FP), R1 // R1 = len(s), in 64-bit words + MOVD $0, R2 // R2 = running result + VEOR V17.B16, V17.B16, V17.B16 // zero the 32-bit accumulator + LSR $3, R1, R3 // R3 = len/8 = number of 64-byte blocks + CBZ R3, sltail // fewer than 8 words: skip the vector loop +slblock: + MOVD INNERMAX, R4 // R4 = min(remaining blocks, INNERMAX) + CMP R4, R3 + BHS slinner + MOVD R3, R4 +slinner: + SUB R4, R3, R3 // R3 -= this batch's block count + ZEROPART +slloop: + VLD1.P 64(R0), [V0.B16, V1.B16, V2.B16, V3.B16] // load 8 words (64 bytes) + FOLD4 // partials += popcount(those 64 bytes) + SUBS $1, R4, R4 + BNE slloop + DRAIN // fold partials into V17, re-zero them + CBNZ R3, slblock // more blocks remain + REDUCE // R2 += sum of V17's lanes +sltail: + AND $7, R1, R1 // leftover words (0..7) + CBZ R1, sldone +sltailloop: + VLD1.P 8(R0), [V0.D1] // load one word, advance R0 by 8 + TAILWORD + SUBS $1, R1, R1 + BNE sltailloop +sldone: + MOVD R2, ret+24(FP) + RET + +// func _popcntAndSliceNEON(s, m []uint64) uint64 +// Returns the sum of popcount(s[i] & m[i]). Mirrors _popcntSliceNEON but loads +// four vectors from each of s and m and ANDs them before counting. s and m are +// assumed to have equal length. +TEXT ·_popcntAndSliceNEON(SB), NOSPLIT, $0-56 + MOVD s_base+0(FP), R0 // R0 = &s[0] + MOVD m_base+24(FP), R1 // R1 = &m[0] + MOVD s_len+8(FP), R5 // R5 = len + MOVD $0, R2 + VEOR V17.B16, V17.B16, V17.B16 + LSR $3, R5, R3 + CBZ R3, andtail +andblock: + MOVD INNERMAX, R4 + CMP R4, R3 + BHS andinner + MOVD R3, R4 +andinner: + SUB R4, R3, R3 + ZEROPART +andloop: + VLD1.P 64(R0), [V0.B16, V1.B16, V2.B16, V3.B16] + VLD1.P 64(R1), [V4.B16, V5.B16, V6.B16, V7.B16] + VAND V4.B16, V0.B16, V0.B16 // V0 = s & m + VAND V5.B16, V1.B16, V1.B16 + VAND V6.B16, V2.B16, V2.B16 + VAND V7.B16, V3.B16, V3.B16 + FOLD4 + SUBS $1, R4, R4 + BNE andloop + DRAIN + CBNZ R3, andblock + REDUCE +andtail: + AND $7, R5, R5 + CBZ R5, anddone +andtailloop: + VLD1.P 8(R0), [V0.D1] + VLD1.P 8(R1), [V1.D1] + VAND V1.B8, V0.B8, V0.B8 // s & m, one word + TAILWORD + SUBS $1, R5, R5 + BNE andtailloop +anddone: + MOVD R2, ret+48(FP) // +48: result follows two 24-byte slice headers + RET + +// func _popcntOrSliceNEON(s, m []uint64) uint64 +// Returns the sum of popcount(s[i] | m[i]); see _popcntAndSliceNEON for the +// shared structure. +TEXT ·_popcntOrSliceNEON(SB), NOSPLIT, $0-56 + MOVD s_base+0(FP), R0 + MOVD m_base+24(FP), R1 + MOVD s_len+8(FP), R5 + MOVD $0, R2 + VEOR V17.B16, V17.B16, V17.B16 + LSR $3, R5, R3 + CBZ R3, ortail +orblock: + MOVD INNERMAX, R4 + CMP R4, R3 + BHS orinner + MOVD R3, R4 +orinner: + SUB R4, R3, R3 + ZEROPART +orloop: + VLD1.P 64(R0), [V0.B16, V1.B16, V2.B16, V3.B16] + VLD1.P 64(R1), [V4.B16, V5.B16, V6.B16, V7.B16] + VORR V4.B16, V0.B16, V0.B16 // V0 = s | m + VORR V5.B16, V1.B16, V1.B16 + VORR V6.B16, V2.B16, V2.B16 + VORR V7.B16, V3.B16, V3.B16 + FOLD4 + SUBS $1, R4, R4 + BNE orloop + DRAIN + CBNZ R3, orblock + REDUCE +ortail: + AND $7, R5, R5 + CBZ R5, ordone +ortailloop: + VLD1.P 8(R0), [V0.D1] + VLD1.P 8(R1), [V1.D1] + VORR V1.B8, V0.B8, V0.B8 // s | m, one word + TAILWORD + SUBS $1, R5, R5 + BNE ortailloop +ordone: + MOVD R2, ret+48(FP) + RET + +// func _popcntXorSliceNEON(s, m []uint64) uint64 +// Returns the sum of popcount(s[i] ^ m[i]); see _popcntAndSliceNEON for the +// shared structure. +TEXT ·_popcntXorSliceNEON(SB), NOSPLIT, $0-56 + MOVD s_base+0(FP), R0 + MOVD m_base+24(FP), R1 + MOVD s_len+8(FP), R5 + MOVD $0, R2 + VEOR V17.B16, V17.B16, V17.B16 + LSR $3, R5, R3 + CBZ R3, xortail +xorblock: + MOVD INNERMAX, R4 + CMP R4, R3 + BHS xorinner + MOVD R3, R4 +xorinner: + SUB R4, R3, R3 + ZEROPART +xorloop: + VLD1.P 64(R0), [V0.B16, V1.B16, V2.B16, V3.B16] + VLD1.P 64(R1), [V4.B16, V5.B16, V6.B16, V7.B16] + VEOR V4.B16, V0.B16, V0.B16 // V0 = s ^ m + VEOR V5.B16, V1.B16, V1.B16 + VEOR V6.B16, V2.B16, V2.B16 + VEOR V7.B16, V3.B16, V3.B16 + FOLD4 + SUBS $1, R4, R4 + BNE xorloop + DRAIN + CBNZ R3, xorblock + REDUCE +xortail: + AND $7, R5, R5 + CBZ R5, xordone +xortailloop: + VLD1.P 8(R0), [V0.D1] + VLD1.P 8(R1), [V1.D1] + VEOR V1.B8, V0.B8, V0.B8 // s ^ m, one word + TAILWORD + SUBS $1, R5, R5 + BNE xortailloop +xordone: + MOVD R2, ret+48(FP) + RET + +// func _popcntMaskSliceNEON(s, m []uint64) uint64 +// Returns the sum of popcount(s[i] &^ m[i]) == popcount(s & ~m). Same structure +// as _popcntAndSliceNEON; arm64's NEON has no "and-not" here, so ~m is formed by +// XORing m with the all-ones register V15 (materialized once by VMOVI) before +// the AND. +TEXT ·_popcntMaskSliceNEON(SB), NOSPLIT, $0-56 + MOVD s_base+0(FP), R0 + MOVD m_base+24(FP), R1 + MOVD s_len+8(FP), R5 + MOVD $0, R2 + VMOVI $255, V15.B16 // V15 = all ones, used to invert m + VEOR V17.B16, V17.B16, V17.B16 + LSR $3, R5, R3 + CBZ R3, masktail +maskblock: + MOVD INNERMAX, R4 + CMP R4, R3 + BHS maskinner + MOVD R3, R4 +maskinner: + SUB R4, R3, R3 + ZEROPART +maskloop: + VLD1.P 64(R0), [V0.B16, V1.B16, V2.B16, V3.B16] + VLD1.P 64(R1), [V4.B16, V5.B16, V6.B16, V7.B16] + VEOR V15.B16, V4.B16, V4.B16 // V4 = ~m + VEOR V15.B16, V5.B16, V5.B16 + VEOR V15.B16, V6.B16, V6.B16 + VEOR V15.B16, V7.B16, V7.B16 + VAND V4.B16, V0.B16, V0.B16 // V0 = s & ~m = s &^ m + VAND V5.B16, V1.B16, V1.B16 + VAND V6.B16, V2.B16, V2.B16 + VAND V7.B16, V3.B16, V3.B16 + FOLD4 + SUBS $1, R4, R4 + BNE maskloop + DRAIN + CBNZ R3, maskblock + REDUCE +masktail: + AND $7, R5, R5 + CBZ R5, maskdone +masktailloop: + VLD1.P 8(R0), [V0.D1] + VLD1.P 8(R1), [V1.D1] + VEOR V15.B8, V1.B8, V1.B8 // ~m, one word + VAND V1.B8, V0.B8, V0.B8 // s &^ m, one word + TAILWORD + SUBS $1, R5, R5 + BNE masktailloop +maskdone: + MOVD R2, ret+48(FP) + RET diff --git a/popcnt_neon_arm64_test.go b/popcnt_neon_arm64_test.go new file mode 100644 index 00000000..fd087e37 --- /dev/null +++ b/popcnt_neon_arm64_test.go @@ -0,0 +1,123 @@ +//go:build arm64 && !appengine +// +build arm64,!appengine + +package roaring + +import ( + "math/rand" + "testing" + + "github.com/stretchr/testify/assert" +) + +// edge lengths exercise the NEON main loop (multiples of 2) and the scalar-width +// tail (len % 2 != 0), including the empty and sub-block cases. +var neonTestLengths = []int{0, 1, 2, 3, 4, 5, 7, 8, 15, 16, 17, 31, 63, 64, 65, 1023, 1024, 1025} + +func randomUint64Slice(r *rand.Rand, n int) []uint64 { + s := make([]uint64, n) + for i := range s { + // mix fully-random words with sparse and dense ones to vary popcounts. + switch i % 4 { + case 0: + s[i] = r.Uint64() + case 1: + s[i] = 0 + case 2: + s[i] = ^uint64(0) + default: + s[i] = r.Uint64() & r.Uint64() + } + } + return s +} + +func benchPopcntPair(b *testing.B, neon bool, fn func(s, m []uint64) uint64) { + if neon && !useNEON { + b.Skip("NEON not available") + } + r := rand.New(rand.NewSource(1)) + s := randomUint64Slice(r, 1024) + m := randomUint64Slice(r, 1024) + var sink uint64 + b.ResetTimer() + for i := 0; i < b.N; i++ { + sink += fn(s, m) + } + _ = sink +} + +func BenchmarkPopcntAndSlice1024NEON(b *testing.B) { + benchPopcntPair(b, true, _popcntAndSliceNEON) +} + +func BenchmarkPopcntAndSlice1024Go(b *testing.B) { + benchPopcntPair(b, false, popcntAndSliceGo) +} + +func BenchmarkPopcntSlice1024NEON(b *testing.B) { + if !useNEON { + b.Skip("NEON not available") + } + r := rand.New(rand.NewSource(1)) + s := randomUint64Slice(r, 1024) + var sink uint64 + b.ResetTimer() + for i := 0; i < b.N; i++ { + sink += _popcntSliceNEON(s) + } + _ = sink +} + +func BenchmarkPopcntSlice1024Go(b *testing.B) { + r := rand.New(rand.NewSource(1)) + s := randomUint64Slice(r, 1024) + var sink uint64 + b.ResetTimer() + for i := 0; i < b.N; i++ { + sink += popcntSliceGo(s) + } + _ = sink +} + +func TestNEONPopcntDispatch(t *testing.T) { + // Verify the runtime dispatch wrappers agree with the Go reference both when + // NEON is selected and when the scalar fallback is forced. + saved := useNEON + defer func() { useNEON = saved }() + + r := rand.New(rand.NewSource(7)) + for _, on := range []bool{false, true} { + useNEON = on + for _, n := range neonTestLengths { + s := randomUint64Slice(r, n) + m := randomUint64Slice(r, n) + assert.Equalf(t, popcntSliceGo(s), popcntSlice(s), "popcntSlice neon=%v len=%d", on, n) + assert.Equalf(t, popcntAndSliceGo(s, m), popcntAndSlice(s, m), "popcntAndSlice neon=%v len=%d", on, n) + assert.Equalf(t, popcntOrSliceGo(s, m), popcntOrSlice(s, m), "popcntOrSlice neon=%v len=%d", on, n) + assert.Equalf(t, popcntXorSliceGo(s, m), popcntXorSlice(s, m), "popcntXorSlice neon=%v len=%d", on, n) + assert.Equalf(t, popcntMaskSliceGo(s, m), popcntMaskSlice(s, m), "popcntMaskSlice neon=%v len=%d", on, n) + } + } +} + +func TestNEONPopcntDifferential(t *testing.T) { + r := rand.New(rand.NewSource(42)) + for _, n := range neonTestLengths { + for iter := 0; iter < 64; iter++ { + s := randomUint64Slice(r, n) + m := randomUint64Slice(r, n) + + assert.Equalf(t, popcntSliceGo(s), _popcntSliceNEON(s), + "popcntSlice len=%d", n) + assert.Equalf(t, popcntAndSliceGo(s, m), _popcntAndSliceNEON(s, m), + "popcntAndSlice len=%d", n) + assert.Equalf(t, popcntOrSliceGo(s, m), _popcntOrSliceNEON(s, m), + "popcntOrSlice len=%d", n) + assert.Equalf(t, popcntXorSliceGo(s, m), _popcntXorSliceNEON(s, m), + "popcntXorSlice len=%d", n) + assert.Equalf(t, popcntMaskSliceGo(s, m), _popcntMaskSliceNEON(s, m), + "popcntMaskSlice len=%d", n) + } + } +} diff --git a/popcnt_slices.go b/popcnt_slices.go index d27c5f38..4afdd713 100644 --- a/popcnt_slices.go +++ b/popcnt_slices.go @@ -1,9 +1,11 @@ package roaring +import "math/bits" + func popcntSliceGo(s []uint64) uint64 { cnt := uint64(0) for _, x := range s { - cnt += popcount(x) + cnt += uint64(bits.OnesCount64(x)) } return cnt } @@ -11,7 +13,7 @@ func popcntSliceGo(s []uint64) uint64 { func popcntMaskSliceGo(s, m []uint64) uint64 { cnt := uint64(0) for i := range s { - cnt += popcount(s[i] &^ m[i]) + cnt += uint64(bits.OnesCount64(s[i] &^ m[i])) } return cnt } @@ -19,7 +21,7 @@ func popcntMaskSliceGo(s, m []uint64) uint64 { func popcntAndSliceGo(s, m []uint64) uint64 { cnt := uint64(0) for i := range s { - cnt += popcount(s[i] & m[i]) + cnt += uint64(bits.OnesCount64(s[i] & m[i])) } return cnt } @@ -27,7 +29,7 @@ func popcntAndSliceGo(s, m []uint64) uint64 { func popcntOrSliceGo(s, m []uint64) uint64 { cnt := uint64(0) for i := range s { - cnt += popcount(s[i] | m[i]) + cnt += uint64(bits.OnesCount64(s[i] | m[i])) } return cnt } @@ -35,7 +37,7 @@ func popcntOrSliceGo(s, m []uint64) uint64 { func popcntXorSliceGo(s, m []uint64) uint64 { cnt := uint64(0) for i := range s { - cnt += popcount(s[i] ^ m[i]) + cnt += uint64(bits.OnesCount64(s[i] ^ m[i])) } return cnt } diff --git a/popcnt_slices_test.go b/popcnt_slices_test.go deleted file mode 100644 index 53d8f5ff..00000000 --- a/popcnt_slices_test.go +++ /dev/null @@ -1,66 +0,0 @@ -//go:build amd64 && !appengine && !go1.9 -// +build amd64,!appengine,!go1.9 - -// This file tests the popcnt functions - -package roaring - -import ( - "testing" - - "github.com/stretchr/testify/assert" -) - -func TestPopcntSlice(t *testing.T) { - s := []uint64{2, 3, 5, 7, 11, 13, 17, 19, 23, 29} - resGo := popcntSliceGo(s) - resAsm := popcntSliceAsm(s) - res := popcntSlice(s) - - assert.Equal(t, resGo, resAsm) - assert.Equal(t, resGo, res) -} - -func TestPopcntMaskSlice(t *testing.T) { - s := []uint64{2, 3, 5, 7, 11, 13, 17, 19, 23, 29} - m := []uint64{31, 37, 41, 43, 47, 53, 59, 61, 67, 71} - resGo := popcntMaskSliceGo(s, m) - resAsm := popcntMaskSliceAsm(s, m) - res := popcntMaskSlice(s, m) - - assert.Equal(t, resGo, resAsm) - assert.Equal(t, resGo, res) -} - -func TestPopcntAndSlice(t *testing.T) { - s := []uint64{2, 3, 5, 7, 11, 13, 17, 19, 23, 29} - m := []uint64{31, 37, 41, 43, 47, 53, 59, 61, 67, 71} - resGo := popcntAndSliceGo(s, m) - resAsm := popcntAndSliceAsm(s, m) - res := popcntAndSlice(s, m) - - assert.Equal(t, resGo, resAsm) - assert.Equal(t, resGo, res) -} - -func TestPopcntOrSlice(t *testing.T) { - s := []uint64{2, 3, 5, 7, 11, 13, 17, 19, 23, 29} - m := []uint64{31, 37, 41, 43, 47, 53, 59, 61, 67, 71} - resGo := popcntOrSliceGo(s, m) - resAsm := popcntOrSliceAsm(s, m) - res := popcntOrSlice(s, m) - - assert.Equal(t, resGo, resAsm) - assert.Equal(t, resGo, res) -} - -func TestPopcntXorSlice(t *testing.T) { - s := []uint64{2, 3, 5, 7, 11, 13, 17, 19, 23, 29} - m := []uint64{31, 37, 41, 43, 47, 53, 59, 61, 67, 71} - resGo := popcntXorSliceGo(s, m) - resAsm := popcntXorSliceAsm(s, m) - res := popcntXorSlice(s, m) - - assert.Equal(t, resGo, resAsm) - assert.Equal(t, resGo, res) -} diff --git a/roaring.go b/roaring.go index 50ae58e9..129034e3 100644 --- a/roaring.go +++ b/roaring.go @@ -10,6 +10,7 @@ import ( "encoding/base64" "fmt" "io" + "math/bits" "strconv" "github.com/RoaringBitmap/roaring/v2/internal" @@ -166,7 +167,7 @@ func (rb *Bitmap) FromDense(bitmap []uint64, doCopy bool) { for _, w := range words { for w != 0 { t := w & -w - c.content[pos] = uint16(base + int(popcount(t-1))) + c.content[pos] = uint16(base + bits.OnesCount64(t-1)) pos++ w ^= t } @@ -1601,20 +1602,26 @@ func (rb *Bitmap) Xor(x2 *Bitmap) { break } } else if s1 > s2 { - rb.highlowcontainer.insertNewKeyValueAt(pos1, x2.highlowcontainer.getKeyAtIndex(pos2), x2.highlowcontainer.getContainerAtIndex(pos2).clone()) - length1++ - pos1++ - pos2++ + // A source-only key must be inserted at pos1. Inserting in + // place shifts the aligned suffix once per inserted key, which + // is quadratic when many keys are interleaved. Finish the merge + // into fresh slices in a single linear pass instead. + rb.highlowcontainer.mergeBulk(&x2.highlowcontainer, pos1, pos1, pos2, true) + return } else { c := rb.highlowcontainer.getWritableContainerAtIndex(pos1).ixor(x2.highlowcontainer.getContainerAtIndex(pos2)) if !c.isEmpty() { rb.highlowcontainer.setContainerAtIndex(pos1, c) pos1++ + pos2++ } else { - rb.highlowcontainer.removeAtIndex(pos1) - length1-- + // The aligned containers cancelled out. Removing pos1 in + // place would shift the suffix once per removed key (also + // quadratic), so finish the merge into fresh slices, + // dropping this now-empty container. + rb.highlowcontainer.mergeBulk(&x2.highlowcontainer, pos1, pos1+1, pos2+1, true) + return } - pos2++ } } else { break @@ -1644,14 +1651,13 @@ main: } s1 = rb.highlowcontainer.getKeyAtIndex(pos1) } else if s1 > s2 { - rb.highlowcontainer.insertNewKeyValueAt(pos1, s2, x2.highlowcontainer.getContainerAtIndex(pos2).clone()) - pos1++ - length1++ - pos2++ - if pos2 == length2 { - break main - } - s2 = x2.highlowcontainer.getKeyAtIndex(pos2) + // The receiver has run ahead of the source: a source-only key + // must be inserted at pos1. Inserting in place shifts the + // aligned suffix once per inserted key, which is quadratic when + // many source-only keys are interleaved. Finish the merge into + // fresh slices in a single linear pass instead. + rb.highlowcontainer.mergeBulk(&x2.highlowcontainer, pos1, pos1, pos2, false) + return } else { newcont := rb.highlowcontainer.getUnionedWritableContainer(pos1, x2.highlowcontainer.getContainerAtIndex(pos2)) rb.highlowcontainer.replaceKeyAndContainerAtIndex(pos1, s1, newcont, false) diff --git a/roaring64/BSI_BENCHMARKS.md b/roaring64/BSI_BENCHMARKS.md new file mode 100644 index 00000000..091e1242 --- /dev/null +++ b/roaring64/BSI_BENCHMARKS.md @@ -0,0 +1,55 @@ +# BSI64 Benchmarks + +These notes capture local benchmark results for the BSI64 `BatchEqual` and +comparison paths. They are intended as reproducible PR evidence, not as +contractual performance guarantees. + +Environment: + +- CPU: 12th Gen Intel(R) Core(TM) i7-1255U +- OS/arch: linux/amd64 +- Package: `github.com/RoaringBitmap/roaring/v2/roaring64` + +Commands: + +```sh +go test ./roaring64 -count=1 +go test ./roaring64 -run '^$' -bench 'BenchmarkBSI64BatchEqual' -benchmem -count 3 +go test ./roaring64 -run '^$' -bench 'BenchmarkBSI64Compare(Big)?Value|BenchmarkBSI64BatchEqual(Big)?LargeAgeFixture' -benchmem -count 1 +go test ./roaring64 -run '^$' -bench 'BenchmarkBSI64CompareBSISameRow' -benchmem -count=5 +go test ./roaring64 -run '^$' -bench 'BenchmarkBSI64GetBigValue' -benchmem -count=3 +go test ./roaring64 -run '^$' -bench 'BenchmarkBSI64BatchEqual.*LargeFixture' -benchmem -benchtime=2s -count=5 +``` + +Representative results: + +| Benchmark | Before | After | Notes | +| --- | ---: | ---: | --- | +| `BenchmarkBSI64BatchEqualLargeAgeFixture` | ~13-14s/op, ~12.4GB/op | ~145-205ms/op, ~25.5MB/op | Avoids row-by-row `GetBigValue` for int64-width values. | +| `BenchmarkBSI64BatchEqualM128Scattered` | ~1.25s/op, ~458MB/op | ~11-17ms/op, ~12.5MB/op | Detects complete bit-cube value patterns. | +| `BenchmarkBSI64CompareValueEQLargeAgeFixture` | ~4.44s/op, ~461MB/op | ~100-118ms/op, ~19.7MB/op | `EQ` delegates to optimized `BatchEqual`. | +| `BenchmarkBSI64CompareValueRangeLargeAgeFixture` | ~7.49s/op, ~501MB/op | ~204-224ms/op, ~122.6MB/op | Uses bitmap-native signed int64 comparison. | +| `BenchmarkBSI64CompareValueGELargeAgeFixture` | ~3.45s/op, ~500MB/op | ~168-184ms/op, ~82.3MB/op | Uses bitmap-native signed int64 comparison. | +| `BenchmarkBSI64CompareBSISameRowBitwise` | ~127-168ms/op, ~69.7MB/op | ~568-795us/op, ~619KB/op | Compares two BSI values per column ID through bitplane algebra instead of row-by-row `GetBigValue`. | +| `BenchmarkBSI64GetBigValuesLargeFixture` | ~69-92ms/op, ~35.6MB/op, ~1.3M allocs/op for a row-by-row `GetBigValue` loop | ~23-34ms/op, ~8.2MB/op, ~200k allocs/op | Extracts aligned BSI values for a column batch by walking bit-slices once. | +| `BenchmarkBSI64BatchEqualValuesLargeFixture` | ~5.4-7.1ms/op for `BatchEqual` plus `GetBigValues`; ~10.9-13.0ms/op for `BatchEqual` plus row-by-row `GetValue` | ~1.6-2.3ms/op, ~2.0MB/op, ~432 allocs/op | Emits matched column IDs and int64 values directly from trie leaves, avoiding a second value lookup pass. | + +Compatibility: + +- Public method signatures are unchanged. +- `CompareBigValue` and `BatchEqualBig` internally delegate to the optimized + int64 paths only when the BSI and query values fit in signed 64-bit space. +- True wider-than-64-bit values continue to use the existing generic paths. +- `BatchEqualBig` now keys values by sign and magnitude so positive and negative + values with the same magnitude do not collide. +- `GetBigValues` returns values aligned to the requested column IDs, with nil + entries for missing values, while preserving `GetBigValue` semantics. +- `BatchEqualValues` returns matched column IDs and int64 values for `BatchEqual` + shapes, optionally restricted by a found set. Result order is intentionally + unspecified. + +Follow-up: + +- This change is scoped to `roaring64`. The 32-bit `BitSliceIndexing` package + already has separate `BatchEqual` coverage, and `CompareValue` parity can be + addressed in a follow-up PR with its own benchmarks and signed-value tests. diff --git a/roaring64/bsi64.go b/roaring64/bsi64.go index e9b1eca4..4e918ddb 100644 --- a/roaring64/bsi64.go +++ b/roaring64/bsi64.go @@ -5,6 +5,7 @@ import ( "io" "math/big" "runtime" + "sort" "sync" ) @@ -24,6 +25,12 @@ type BSI struct { runOptimized bool } +// BSIValuePair is a column ID and its BSI value. +type BSIValuePair struct { + ColumnID uint64 + Value int64 +} + // NewBSI constructs a new BSI. Note that it is your responsibility to ensure that // the min/max values are set correctly. Queries CompareValue, MinMax, etc. will not // work correctly if the min/max values are not set correctly. @@ -127,6 +134,7 @@ func (b *BSI) SetBigValue(columnID uint64, value *big.Int) { b.eBM.Add(columnID) } +// SetBigMany sets value for all columns in foundSet. func (b *BSI) SetBigMany(foundSet *Bitmap, value *big.Int) { // If max/min values are set to zero then automatically determine bit array size if b.MaxValue == 0 && b.MinValue == 0 { @@ -206,6 +214,131 @@ func (b *BSI) GetBigValue(columnID uint64) (value *big.Int, exists bool) { return val, exists } +// GetBigValues gets values for the column IDs. Returned values are aligned with +// columnIDs, and a nil entry means the corresponding column ID has no value. +func (b *BSI) GetBigValues(columnIDs []uint64) []*big.Int { + values := make([]*big.Int, len(columnIDs)) + if len(columnIDs) == 0 { + return values + } + if len(columnIDs) == 1 { + if value, ok := b.GetBigValue(columnIDs[0]); ok { + values[0] = value + } + return values + } + request := newBSIGetBigValuesRequest(columnIDs) + if !b.isBig() { + return b.getBigValuesInt64(request, values) + } + return b.getBigValuesGeneric(request, values) +} + +type bsiGetBigValuesRequest struct { + foundSet *Bitmap + positions map[uint64]int + duplicatePositions map[uint64][]int +} + +func newBSIGetBigValuesRequest(columnIDs []uint64) bsiGetBigValuesRequest { + foundSet := NewBitmap() + positions := make(map[uint64]int, len(columnIDs)) + var duplicatePositions map[uint64][]int + for position, columnID := range columnIDs { + if _, ok := positions[columnID]; ok { + if duplicatePositions == nil { + duplicatePositions = make(map[uint64][]int) + } + duplicatePositions[columnID] = append(duplicatePositions[columnID], position) + continue + } + positions[columnID] = position + foundSet.Add(columnID) + } + return bsiGetBigValuesRequest{ + foundSet: foundSet, + positions: positions, + duplicatePositions: duplicatePositions, + } +} + +func (b *BSI) getBigValuesInt64(request bsiGetBigValuesRequest, values []*big.Int) []*big.Int { + existing := And(&b.eBM, request.foundSet) + if existing.IsEmpty() { + return values + } + + rawValues := make([]uint64, len(values)) + signBit := b.BitCount() + for bit := 0; bit <= signBit; bit++ { + bitSet := And(&b.bA[bit], existing) + iter := bitSet.Iterator() + for iter.HasNext() { + columnID := iter.Next() + rawValues[request.positions[columnID]] |= uint64(1) << uint(bit) + } + } + + width := uint(signBit + 1) + signMask := uint64(1) << uint(signBit) + iter := existing.Iterator() + for iter.HasNext() { + columnID := iter.Next() + position := request.positions[columnID] + rawValue := rawValues[position] + if rawValue&signMask != 0 && width < 64 { + rawValue |= ^uint64(0) << width + } + values[position] = big.NewInt(int64(rawValue)) + } + fillDuplicateBigValues(values, request) + return values +} + +func (b *BSI) getBigValuesGeneric(request bsiGetBigValuesRequest, values []*big.Int) []*big.Int { + existing := And(&b.eBM, request.foundSet) + if existing.IsEmpty() { + return values + } + + iter := existing.Iterator() + for iter.HasNext() { + values[request.positions[iter.Next()]] = big.NewInt(0) + } + for bit := b.BitCount(); bit >= 0; bit-- { + bitSet := And(&b.bA[bit], existing) + iter := bitSet.Iterator() + for iter.HasNext() { + columnID := iter.Next() + position := request.positions[columnID] + values[position].SetBit(values[position], bit, 1) + } + } + + signBit := b.BitCount() + negativeSet := And(&b.bA[signBit], existing) + iter = negativeSet.Iterator() + for iter.HasNext() { + position := request.positions[iter.Next()] + values[position] = negativeTwosComplementToInt(values[position]) + } + + fillDuplicateBigValues(values, request) + return values +} + +func fillDuplicateBigValues(values []*big.Int, request bsiGetBigValuesRequest) { + for columnID, extraPositions := range request.duplicatePositions { + value := values[request.positions[columnID]] + if value == nil { + continue + } + for _, position := range extraPositions { + values[position] = new(big.Int).Set(value) + } + } +} + func negativeTwosComplementToInt(val *big.Int) *big.Int { inverted := new(big.Int).Not(val) mask := new(big.Int).Lsh(big.NewInt(1), uint(val.BitLen())) @@ -348,9 +481,173 @@ type task struct { func (b *BSI) CompareValue(parallelism int, op Operation, valueOrStart, end int64, foundSet *Bitmap) *Bitmap { + if result, ok := b.compareInt64Value(parallelism, op, valueOrStart, end, foundSet); ok { + return result + } return b.CompareBigValue(parallelism, op, big.NewInt(valueOrStart), big.NewInt(end), foundSet) } +// CompareBSI compares values from two BSIs by column ID and returns the column +// IDs where b[columnID] op other[columnID] is true. Only column IDs present in +// both existence bitmaps are considered. When foundSet is not nil, it further +// restricts the comparison universe. +func (b *BSI) CompareBSI(op Operation, other *BSI, foundSet *Bitmap) *Bitmap { + if b == nil || other == nil || b.eBM.IsEmpty() || other.eBM.IsEmpty() { + return NewBitmap() + } + universe := b.eBM.Clone() + universe.And(&other.eBM) + if foundSet != nil { + universe.And(foundSet) + } + if universe.IsEmpty() { + return universe + } + + commonSign := b.BitCount() + if other.BitCount() > commonSign { + commonSign = other.BitCount() + } + less, equal := b.compareBSILessAndEqual(other, commonSign, universe) + + switch op { + case LT: + return less + case LE: + less.Or(equal) + return less + case EQ: + return equal + case GE: + universe.AndNot(less) + return universe + case GT: + less.Or(equal) + universe.AndNot(less) + return universe + default: + panic(fmt.Sprintf("Operation [%v] not supported for BSI comparison", op)) + } +} + +func (b *BSI) compareBSILessAndEqual(other *BSI, commonSign int, universe *Bitmap) (*Bitmap, *Bitmap) { + less := NewBitmap() + equalPrefix := universe.Clone() + for i := commonSign; i >= 0; i-- { + leftOnes := b.compareBSIPlaneChild(equalPrefix, i, commonSign, true, false) + rightOnes := other.compareBSIPlaneChild(equalPrefix, i, commonSign, true, false) + + rightOnly := rightOnes.Clone() + rightOnly.AndNot(leftOnes) + less.Or(rightOnly) + + leftOnly := leftOnes + leftOnly.AndNot(rightOnes) + rightOnly.Or(leftOnly) + equalPrefix.AndNot(rightOnly) + if equalPrefix.IsEmpty() { + break + } + } + return less, equalPrefix +} + +func (b *BSI) compareBSIPlaneChild(prefix *Bitmap, planeIndex, commonSign int, set, owned bool) *Bitmap { + sourcePlane := planeIndex + if sourcePlane > b.BitCount() { + sourcePlane = b.BitCount() + } + rawSet := set + if planeIndex == commonSign { + rawSet = !rawSet + } + return bsi64PlaneChild(prefix, &b.bA[sourcePlane], rawSet, owned) +} + +func (b *BSI) compareInt64Value(parallelism int, op Operation, valueOrStart, end int64, foundSet *Bitmap) (*Bitmap, bool) { + bitCount := b.BitCount() + if bitCount > 63 || !bsi64ValueFitsBitCount(valueOrStart, bitCount) { + return nil, false + } + if op == EQ { + result := b.BatchEqual(parallelism, []int64{valueOrStart}) + if foundSet != nil { + result.And(foundSet) + } + return result, true + } + if op == RANGE && !bsi64ValueFitsBitCount(end, bitCount) { + return nil, false + } + + universe := b.eBM.Clone() + if foundSet != nil { + universe.And(foundSet) + } + if universe.IsEmpty() { + return universe, true + } + + start := transformBSI64SignedEncoding(encodeBSI64Value(valueOrStart, bitCount), bitCount) + less, equal := b.compareInt64LessAndEqual(start, universe) + + switch op { + case LT: + return less, true + case LE: + less.Or(equal) + return less, true + case GE: + universe.AndNot(less) + return universe, true + case GT: + less.Or(equal) + universe.AndNot(less) + return universe, true + case RANGE: + if valueOrStart > end { + return NewBitmap(), true + } + universe.AndNot(less) + finish := transformBSI64SignedEncoding(encodeBSI64Value(end, bitCount), bitCount) + rangeLess, rangeEqual := b.compareInt64LessAndEqual(finish, universe) + rangeLess.Or(rangeEqual) + return rangeLess, true + default: + return nil, false + } +} + +func transformBSI64SignedEncoding(encoded uint64, bitCount int) uint64 { + return encoded ^ (uint64(1) << uint(bitCount)) +} + +func (b *BSI) compareInt64LessAndEqual(target uint64, universe *Bitmap) (*Bitmap, *Bitmap) { + less := NewBitmap() + equalPrefix := universe.Clone() + for i := b.BitCount(); i >= 0; i-- { + targetBitSet := target&(uint64(1)<= 64 { + // Fall back to the arbitrary-precision path when the BSI has more than + // int64's finite bit width. This preserves correctness for big-value BSIs. + bigValues := make([]*big.Int, len(values)) + for i, v := range values { + bigValues[i] = big.NewInt(v) + } + return b.BatchEqualBig(parallelism, bigValues) + } + + vals := b.batchEqualInt64Values(values, bitCount) + if len(vals) == 0 { + return NewBitmap() + } + + if result, ok := b.matchInt64Cube(vals, bitCount); ok { + if b.runOptimized { + result.RunOptimize() + } + return result + } + result := b.matchInt64Trie(vals, bitCount, &b.eBM, false) + if b.runOptimized { + result.RunOptimize() + } + return result +} + +// BatchEqualValues returns column IDs and values where the BSI value is +// contained in values. When foundSet is not nil, only column IDs in foundSet are +// considered. Result order is not guaranteed. +func (b *BSI) BatchEqualValues(parallelism int, values []int64, foundSet *Bitmap) []BSIValuePair { + if b.eBM.IsEmpty() || len(values) == 0 { + return nil + } + + bitCount := b.BitCount() + if bitCount >= 64 { + matched := b.BatchEqual(parallelism, values) + if foundSet != nil { + matched.And(foundSet) + } + return b.bsiValuePairsFromBitmap(matched) + } + + vals := b.batchEqualInt64Values(values, bitCount) + if len(vals) == 0 { + return nil + } + + var universe *Bitmap + owned := false + if foundSet == nil { + universe = &b.eBM + } else { + universe = And(&b.eBM, foundSet) + owned = true + } + if universe.IsEmpty() { + return nil + } + + pairs := make([]BSIValuePair, 0) + b.matchInt64TrieValues(vals, bitCount, universe, owned, 0, &pairs) + return pairs +} + +func (b *BSI) batchEqualInt64Values(values []int64, bitCount int) []uint64 { + seen := make(map[uint64]struct{}, len(values)) + vals := make([]uint64, 0, len(values)) + for _, v := range values { + if !bsi64ValueFitsBitCount(v, bitCount) { + continue + } + encoded := encodeBSI64Value(v, bitCount) + if _, ok := seen[encoded]; ok { + continue + } + seen[encoded] = struct{}{} + vals = append(vals, encoded) + } + sort.Slice(vals, func(i, j int) bool { return vals[i] < vals[j] }) + return vals +} + +func bsi64ValueFitsBitCount(value int64, bitCount int) bool { + if bitCount >= 63 { + return true } - return b.BatchEqualBig(parallelism, bigValues) + min := -(int64(1) << uint(bitCount)) + max := (int64(1) << uint(bitCount)) - 1 + return value >= min && value <= max +} + +func encodeBSI64Value(value int64, bitCount int) uint64 { + if bitCount >= 63 { + return uint64(value) + } + mask := (uint64(1) << uint(bitCount+1)) - 1 + return uint64(value) & mask +} + +func decodeBSI64Value(encoded uint64, bitCount int) int64 { + if bitCount >= 63 { + return int64(encoded) + } + width := uint(bitCount + 1) + signMask := uint64(1) << uint(bitCount) + if encoded&signMask != 0 && width < 64 { + encoded |= ^uint64(0) << width + } + return int64(encoded) +} + +func (b *BSI) matchInt64Cube(vals []uint64, bitCount int) (*Bitmap, bool) { + if bitCount >= 63 { + return nil, false + } + widthMask := (uint64(1) << uint(bitCount+1)) - 1 + fixedOnes := vals[0] & widthMask + fixedZeros := ^vals[0] & widthMask + for _, v := range vals[1:] { + fixedOnes &= v + fixedZeros &= ^v & widthMask + } + + variableMask := ^(fixedOnes | fixedZeros) & widthMask + combinations := uint64(1) << uint(countBSI64Bits(variableMask)) + if uint64(len(vals)) != combinations { + return nil, false + } + for _, v := range vals { + if v&fixedOnes != fixedOnes || (^v)&fixedZeros != fixedZeros { + return nil, false + } + } + + result := b.eBM.Clone() + for i := 0; i <= bitCount; i++ { + bit := uint64(1) << uint(i) + if variableMask&bit != 0 { + continue + } + if fixedOnes&bit != 0 { + result.And(&b.bA[i]) + } else { + result.AndNot(&b.bA[i]) + } + if result.IsEmpty() { + break + } + } + return result, true +} + +func countBSI64Bits(value uint64) int { + count := 0 + for value != 0 { + value &= value - 1 + count++ + } + return count +} + +func (b *BSI) matchInt64Trie(vals []uint64, p int, prefix *Bitmap, owned bool) *Bitmap { + if prefix.IsEmpty() { + if owned { + return prefix + } + return NewBitmap() + } + if p < 0 || (p < 63 && uint64(len(vals)) == uint64(1)<= len(values) || values[i] == nil { + continue + } + pairs = append(pairs, BSIValuePair{ + ColumnID: columnID, + Value: values[i].Int64(), + }) + } + return pairs +} + +func bsi64PlaneChild(prefix, plane *Bitmap, set, owned bool) *Bitmap { + if owned { + if set { + prefix.And(plane) + } else { + prefix.AndNot(plane) + } + return prefix + } + if set { + return And(prefix, plane) + } + return AndNot(prefix, plane) } // BatchEqualBig returns a bitmap containing the column IDs where the values are contained within the list of values provided. func (b *BSI) BatchEqualBig(parallelism int, values []*big.Int) *Bitmap { + if b.eBM.IsEmpty() || len(values) == 0 { + return NewBitmap() + } + + if intValues, ok := b.batchEqualBigValuesAsInt64(values); ok { + return b.BatchEqual(parallelism, intValues) + } valMap := make(map[string]struct{}, len(values)) for i := 0; i < len(values); i++ { - valMap[string(values[i].Bytes())] = struct{}{} + if values[i] == nil { + continue + } + valMap[batchEqualBigKey(values[i])] = struct{}{} + } + if len(valMap) == 0 { + return NewBitmap() } comp := &task{bsi: b, values: valMap} return parallelExecutor(parallelism, comp, batchEqual, &b.eBM) } +func batchEqualBigKey(value *big.Int) string { + bytes := value.Bytes() + key := make([]byte, len(bytes)+1) + key[0] = byte(value.Sign() + 1) + copy(key[1:], bytes) + return string(key) +} + +func (b *BSI) batchEqualBigValuesAsInt64(values []*big.Int) ([]int64, bool) { + if b.BitCount() > 63 { + return nil, false + } + intValues := make([]int64, 0, len(values)) + for _, value := range values { + if value == nil { + continue + } + if !value.IsInt64() { + return nil, false + } + intValues = append(intValues, value.Int64()) + } + if len(intValues) == 0 { + return nil, false + } + return intValues, true +} + func batchEqual(e *task, batch []uint64, resultsChan chan *Bitmap, wg *sync.WaitGroup) { @@ -1002,7 +1614,7 @@ func batchEqual(e *task, batch []uint64, resultsChan chan *Bitmap, for i := 0; i < len(batch); i++ { cID := batch[i] if value, ok := e.bsi.GetBigValue(cID); ok { - if _, yes := e.values[string(value.Bytes())]; yes { + if _, yes := e.values[batchEqualBigKey(value)]; yes { results.Add(cID) } } diff --git a/roaring64/bsi64_batch_equal_test.go b/roaring64/bsi64_batch_equal_test.go new file mode 100644 index 00000000..73716daa --- /dev/null +++ b/roaring64/bsi64_batch_equal_test.go @@ -0,0 +1,228 @@ +package roaring64 + +import ( + "math/big" + "math/rand" + "testing" + + "github.com/stretchr/testify/assert" +) + +func expectedBSI64BatchEqual(bsi *BSI, query []int64) *Bitmap { + expected := NewBitmap() + want := make(map[int64]struct{}, len(query)) + for _, q := range query { + want[q] = struct{}{} + } + iter := bsi.GetExistenceBitmap().Iterator() + for iter.HasNext() { + col := iter.Next() + val, ok := bsi.GetValue(col) + if ok { + if _, hit := want[val]; hit { + expected.Add(col) + } + } + } + return expected +} + +func TestBSI64BatchEqualEdgeCases(t *testing.T) { + bsi := NewDefaultBSI() + res := bsi.BatchEqual(0, nil) + assert.True(t, res.IsEmpty()) + + res = bsi.BatchEqual(0, []int64{}) + assert.True(t, res.IsEmpty()) + + bsi.SetValue(10, 42) + bsi.SetValue(20, 100) + bsi.SetValue(30, 42) + bsi.SetValue(40, -5) + bsi.SetValue(50, 5) + + res = bsi.BatchEqual(0, []int64{42}) + assert.Equal(t, uint64(2), res.GetCardinality()) + assert.True(t, res.Contains(10)) + assert.True(t, res.Contains(30)) + + res = bsi.BatchEqual(0, []int64{42, 100, 42, 999}) + assert.Equal(t, uint64(3), res.GetCardinality()) + assert.True(t, res.Contains(10)) + assert.True(t, res.Contains(20)) + assert.True(t, res.Contains(30)) + + res = bsi.BatchEqual(0, []int64{-5}) + assert.Equal(t, uint64(1), res.GetCardinality()) + assert.True(t, res.Contains(40)) + assert.False(t, res.Contains(50), "negative and positive values with the same magnitude must not collide") + + res = bsi.BatchEqual(0, []int64{5}) + assert.Equal(t, uint64(1), res.GetCardinality()) + assert.True(t, res.Contains(50)) + assert.False(t, res.Contains(40), "positive and negative values with the same magnitude must not collide") + + bsi62 := NewBSI(1<<62, 0) + bsi62.SetValue(10, 5) + res = bsi62.BatchEqual(0, []int64{5}) + assert.Equal(t, uint64(1), res.GetCardinality()) + assert.True(t, res.Contains(10)) +} + +func TestBSI64BatchEqualSubBitWidthMatchesGetValue(t *testing.T) { + bsi := NewBSI(100, 0) + assert.Equal(t, 7, bsi.BitCount()) + + bsi.SetValue(10, 42) + bsi.SetValue(20, 99) + + for _, query := range [][]int64{{-5}, {200}, {-5, 42, 200}} { + expected := expectedBSI64BatchEqual(bsi, query) + actual := bsi.BatchEqual(0, query) + assert.True(t, actual.Equals(expected), "query %v expected %v got %v", query, expected.ToArray(), actual.ToArray()) + } +} + +func TestBSI64BatchEqualResultIsolation(t *testing.T) { + bsi := NewDefaultBSI() + bsi.SetValue(10, 42) + bsi.SetValue(20, 100) + + res := bsi.BatchEqual(0, []int64{42}) + assert.True(t, res.Contains(10)) + + res.Add(999) + res.Remove(10) + + assert.False(t, bsi.GetExistenceBitmap().Contains(999)) + assert.True(t, bsi.GetExistenceBitmap().Contains(10)) + + val, ok := bsi.GetValue(10) + assert.True(t, ok) + assert.Equal(t, int64(42), val) + + _, ok = bsi.GetValue(999) + assert.False(t, ok) +} + +func TestBSI64BatchEqualConsistentWithGetValue(t *testing.T) { + rg := rand.New(rand.NewSource(42)) + for run := 0; run < 15; run++ { + bsi := NewDefaultBSI() + numCols := rg.Intn(1000) + 10 + for col := 0; col < numCols; col++ { + if rg.Float64() < 0.8 { + val := rg.Int63n(500) - 250 + bsi.SetValue(uint64(col), val) + } + } + + querySizes := []int{rg.Intn(10) + 1, rg.Intn(50) + 50, rg.Intn(200) + 100} + for _, querySize := range querySizes { + query := make([]int64, querySize) + for i := range query { + query[i] = rg.Int63n(600) - 300 + } + expected := expectedBSI64BatchEqual(bsi, query) + + for _, parallelism := range []int{0, 1, 2, 4} { + actual := bsi.BatchEqual(parallelism, query) + if !actual.Equals(expected) { + t.Fatalf("run=%d querySize=%d parallelism=%d query=%v expected=%v actual=%v", + run, querySize, parallelism, query, expected.ToArray(), actual.ToArray()) + } + } + } + } +} + +func TestBSI64BatchEqualBitCubePattern(t *testing.T) { + bsi := NewDefaultBSI() + for col := uint64(0); col < 512; col++ { + bsi.SetValue(col, int64(col%256)) + } + + odds := make([]int64, 0, 128) + for v := int64(1); v < 256; v += 2 { + odds = append(odds, v) + } + + expected := expectedBSI64BatchEqual(bsi, odds) + actual := bsi.BatchEqual(0, odds) + assert.True(t, actual.Equals(expected), "expected %v got %v", expected.ToArray(), actual.ToArray()) +} + +func TestBSI64BatchEqualExistenceAuthority(t *testing.T) { + ebm := BitmapOf(1) + plane := BitmapOf(1, 2) + ebmData, err := ebm.MarshalBinary() + if err != nil { + t.Fatal(err) + } + planeData, err := plane.MarshalBinary() + if err != nil { + t.Fatal(err) + } + bsi := NewDefaultBSI() + if err := bsi.UnmarshalBinary([][]byte{ebmData, planeData}); err != nil { + t.Fatal(err) + } + res := bsi.BatchEqual(0, []int64{1}) + assert.True(t, res.Contains(1)) + assert.False(t, res.Contains(2), "column 2 is not in eBM and must not match") + + large := setupLargeBSI(t) + if large == nil { + t.Skip("skipping, large BSI setup failed") + } + for _, vals := range [][]int64{{16}, {55, 57}, {0, 1, 2, 3}} { + res := large.BatchEqual(0, vals) + outside := AndNot(res, large.GetExistenceBitmap()) + assert.True(t, outside.IsEmpty(), "BatchEqual(%v) returned %d columns outside eBM", vals, outside.GetCardinality()) + } +} + +func BenchmarkBSI64BatchEqualLargeAgeFixture(b *testing.B) { + bsi := setupLargeBSI(b) + if bsi == nil { + b.Skip("skipping, large BSI setup failed") + } + b.ResetTimer() + for i := 0; i < b.N; i++ { + res := bsi.BatchEqual(0, []int64{55, 57}) + _ = res + } +} + +func BenchmarkBSI64BatchEqualBigLargeAgeFixture(b *testing.B) { + bsi := setupLargeBSI(b) + if bsi == nil { + b.Skip("skipping, large BSI setup failed") + } + values := []*big.Int{big.NewInt(55), big.NewInt(57)} + b.ResetTimer() + for i := 0; i < b.N; i++ { + res := bsi.BatchEqualBig(0, values) + _ = res + } +} + +func BenchmarkBSI64BatchEqualM128(b *testing.B) { benchmarkBSI64BatchEqualM(b, 128, 1) } +func BenchmarkBSI64BatchEqualM128Scattered(b *testing.B) { benchmarkBSI64BatchEqualM(b, 128, 2) } +func BenchmarkBSI64BatchEqualM200(b *testing.B) { benchmarkBSI64BatchEqualM(b, 200, 1) } + +func benchmarkBSI64BatchEqualM(b *testing.B, m int, stride int64) { + bsi := setupLargeBSI(b) + if bsi == nil { + b.Skip("skipping, large BSI setup failed") + } + vals := make([]int64, m) + for i := range vals { + vals[i] = int64(i)*stride + stride - 1 + } + b.ResetTimer() + for i := 0; i < b.N; i++ { + res := bsi.BatchEqual(0, vals) + _ = res + } +} diff --git a/roaring64/bsi64_batch_equal_values_test.go b/roaring64/bsi64_batch_equal_values_test.go new file mode 100644 index 00000000..d5730937 --- /dev/null +++ b/roaring64/bsi64_batch_equal_values_test.go @@ -0,0 +1,149 @@ +package roaring64 + +import ( + "math/big" + "math/rand" + "sort" + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestBSI64BatchEqualValuesConsistentWithBatchEqual(t *testing.T) { + rg := rand.New(rand.NewSource(909)) + for run := 0; run < 25; run++ { + bsi := NewDefaultBSI() + numCols := rg.Intn(1000) + 50 + for col := 0; col < numCols; col++ { + if rg.Float64() < 0.90 { + bsi.SetValue(uint64(col), rg.Int63n(400)-200) + } + } + + values := []int64{-200, -99, -5, 0, 7, 42, 42, 199} + foundSet := NewBitmap() + for col := 0; col < numCols; col++ { + if col%3 != 0 { + foundSet.Add(uint64(col)) + } + } + + for _, fs := range []*Bitmap{nil, foundSet} { + expected := expectedBSI64BatchEqualValues(bsi, values, fs) + actual := bsi.BatchEqualValues(0, values, fs) + assert.Equal(t, expected, sortedBSI64ValuePairs(actual), "run=%d foundSet=%v", run, fs != nil) + } + } +} + +func TestBSI64BatchEqualValuesHandlesBigWidthFallback(t *testing.T) { + bsi := NewDefaultBSI() + huge := new(big.Int).Lsh(big.NewInt(1), 90) + bsi.SetBigValue(1, huge) + bsi.SetValue(2, -7) + bsi.SetValue(3, 11) + bsi.SetValue(4, -7) + + foundSet := BitmapOf(1, 2, 3) + actual := sortedBSI64ValuePairs(bsi.BatchEqualValues(0, []int64{-7, 11}, foundSet)) + assert.Equal(t, []BSIValuePair{ + {ColumnID: 2, Value: -7}, + {ColumnID: 3, Value: 11}, + }, actual) +} + +func BenchmarkBSI64BatchEqualValuesLargeFixture(b *testing.B) { + bsi, values, foundSet := setupBSI64BatchEqualValuesFixture(b, 100000, 100, 27) + b.ResetTimer() + for i := 0; i < b.N; i++ { + pairs := bsi.BatchEqualValues(0, values, foundSet) + _ = pairs + } +} + +func BenchmarkBSI64BatchEqualGetBigValuesLargeFixture(b *testing.B) { + bsi, values, foundSet := setupBSI64BatchEqualValuesFixture(b, 100000, 100, 27) + b.ResetTimer() + for i := 0; i < b.N; i++ { + matched := bsi.BatchEqual(0, values) + matched.And(foundSet) + columnIDs := matched.ToArray() + bigValues := bsi.GetBigValues(columnIDs) + pairs := make([]BSIValuePair, 0, len(columnIDs)) + for j, columnID := range columnIDs { + if bigValues[j] != nil { + pairs = append(pairs, BSIValuePair{ColumnID: columnID, Value: bigValues[j].Int64()}) + } + } + _ = pairs + } +} + +func BenchmarkBSI64BatchEqualGetValueLoopLargeFixture(b *testing.B) { + bsi, values, foundSet := setupBSI64BatchEqualValuesFixture(b, 100000, 100, 27) + b.ResetTimer() + for i := 0; i < b.N; i++ { + matched := bsi.BatchEqual(0, values) + matched.And(foundSet) + pairs := make([]BSIValuePair, 0, int(matched.GetCardinality())) + iter := matched.Iterator() + for iter.HasNext() { + columnID := iter.Next() + value, ok := bsi.GetValue(columnID) + if ok { + pairs = append(pairs, BSIValuePair{ColumnID: columnID, Value: value}) + } + } + _ = pairs + } +} + +func expectedBSI64BatchEqualValues(bsi *BSI, values []int64, foundSet *Bitmap) []BSIValuePair { + matched := bsi.BatchEqual(0, values) + if foundSet != nil { + matched.And(foundSet) + } + pairs := make([]BSIValuePair, 0, int(matched.GetCardinality())) + iter := matched.Iterator() + for iter.HasNext() { + columnID := iter.Next() + value, ok := bsi.GetValue(columnID) + if ok { + pairs = append(pairs, BSIValuePair{ColumnID: columnID, Value: value}) + } + } + return sortedBSI64ValuePairs(pairs) +} + +func sortedBSI64ValuePairs(pairs []BSIValuePair) []BSIValuePair { + sorted := append([]BSIValuePair(nil), pairs...) + sort.Slice(sorted, func(i, j int) bool { + if sorted[i].ColumnID != sorted[j].ColumnID { + return sorted[i].ColumnID < sorted[j].ColumnID + } + return sorted[i].Value < sorted[j].Value + }) + return sorted +} + +func setupBSI64BatchEqualValuesFixture(tb testing.TB, rows, valueDomain, valueCount int) (*BSI, []int64, *Bitmap) { + tb.Helper() + bsi := NewDefaultBSI() + for row := 0; row < rows; row++ { + value := int64(row%valueDomain) - int64(valueDomain/2) + bsi.SetValue(uint64(row), value) + } + + values := make([]int64, 0, valueCount) + for i := 0; i < valueCount; i++ { + values = append(values, int64((i*7)%valueDomain)-int64(valueDomain/2)) + } + + foundSet := NewBitmap() + for row := 0; row < rows; row++ { + if row%5 != 0 { + foundSet.Add(uint64(row)) + } + } + return bsi, values, foundSet +} diff --git a/roaring64/bsi64_compare_benchmark_test.go b/roaring64/bsi64_compare_benchmark_test.go new file mode 100644 index 00000000..7fa0ce47 --- /dev/null +++ b/roaring64/bsi64_compare_benchmark_test.go @@ -0,0 +1,304 @@ +package roaring64 + +import ( + "math/big" + "math/rand" + "testing" + + "github.com/stretchr/testify/assert" +) + +func expectedBSI64CompareValue(bsi *BSI, op Operation, valueOrStart, end int64, foundSet *Bitmap) *Bitmap { + expected := NewBitmap() + source := bsi.GetExistenceBitmap() + if foundSet != nil { + source = And(source, foundSet) + } + iter := source.Iterator() + for iter.HasNext() { + col := iter.Next() + val, ok := bsi.GetValue(col) + if !ok { + continue + } + switch op { + case LT: + if val < valueOrStart { + expected.Add(col) + } + case LE: + if val <= valueOrStart { + expected.Add(col) + } + case EQ: + if val == valueOrStart { + expected.Add(col) + } + case GE: + if val >= valueOrStart { + expected.Add(col) + } + case GT: + if val > valueOrStart { + expected.Add(col) + } + case RANGE: + if val >= valueOrStart && val <= end { + expected.Add(col) + } + default: + panic("unsupported test operation") + } + } + return expected +} + +func TestBSI64CompareValueConsistentWithGetValue(t *testing.T) { + rg := rand.New(rand.NewSource(84)) + for run := 0; run < 15; run++ { + bsi := NewDefaultBSI() + numCols := rg.Intn(1000) + 10 + for col := 0; col < numCols; col++ { + if rg.Float64() < 0.8 { + bsi.SetValue(uint64(col), rg.Int63n(500)-250) + } + } + + foundSet := NewBitmap() + iter := bsi.GetExistenceBitmap().Iterator() + for iter.HasNext() { + col := iter.Next() + if col%3 != 0 { + foundSet.Add(col) + } + } + + cases := []struct { + op Operation + start int64 + end int64 + }{ + {LT, -17, 0}, + {LE, -17, 0}, + {EQ, -17, 0}, + {GE, -17, 0}, + {GT, -17, 0}, + {RANGE, -25, 25}, + } + for _, tc := range cases { + for _, fs := range []*Bitmap{nil, foundSet} { + expected := expectedBSI64CompareValue(bsi, tc.op, tc.start, tc.end, fs) + actual := bsi.CompareValue(0, tc.op, tc.start, tc.end, fs) + assert.True(t, actual.Equals(expected), "run=%d op=%d foundSet=%v expected=%v actual=%v", + run, tc.op, fs != nil, expected.ToArray(), actual.ToArray()) + } + } + } +} + +func TestBSI64CompareBigValueConsistentWithGetBigValue(t *testing.T) { + rg := rand.New(rand.NewSource(85)) + for run := 0; run < 15; run++ { + bsi := NewDefaultBSI() + numCols := rg.Intn(1000) + 10 + for col := 0; col < numCols; col++ { + if rg.Float64() < 0.8 { + bsi.SetValue(uint64(col), rg.Int63n(500)-250) + } + } + + foundSet := NewBitmap() + iter := bsi.GetExistenceBitmap().Iterator() + for iter.HasNext() { + col := iter.Next() + if col%3 != 0 { + foundSet.Add(col) + } + } + + cases := []struct { + op Operation + start int64 + end int64 + }{ + {LT, -17, 0}, + {LE, -17, 0}, + {EQ, -17, 0}, + {GE, -17, 0}, + {GT, -17, 0}, + {RANGE, -25, 25}, + } + for _, tc := range cases { + for _, fs := range []*Bitmap{nil, foundSet} { + expected := expectedBSI64CompareBigValue(bsi, tc.op, big.NewInt(tc.start), big.NewInt(tc.end), fs) + actual := bsi.CompareBigValue(0, tc.op, big.NewInt(tc.start), big.NewInt(tc.end), fs) + assert.True(t, actual.Equals(expected), "run=%d op=%d foundSet=%v expected=%v actual=%v", + run, tc.op, fs != nil, expected.ToArray(), actual.ToArray()) + } + } + } +} + +func expectedBSI64CompareBigValue(bsi *BSI, op Operation, valueOrStart, end *big.Int, foundSet *Bitmap) *Bitmap { + expected := NewBitmap() + source := bsi.GetExistenceBitmap() + if foundSet != nil { + source = And(source, foundSet) + } + iter := source.Iterator() + for iter.HasNext() { + col := iter.Next() + val, ok := bsi.GetBigValue(col) + if !ok { + continue + } + switch op { + case LT: + if val.Cmp(valueOrStart) < 0 { + expected.Add(col) + } + case LE: + if val.Cmp(valueOrStart) <= 0 { + expected.Add(col) + } + case EQ: + if val.Cmp(valueOrStart) == 0 { + expected.Add(col) + } + case GE: + if val.Cmp(valueOrStart) >= 0 { + expected.Add(col) + } + case GT: + if val.Cmp(valueOrStart) > 0 { + expected.Add(col) + } + case RANGE: + if val.Cmp(valueOrStart) >= 0 && val.Cmp(end) <= 0 { + expected.Add(col) + } + default: + panic("unsupported test operation") + } + } + return expected +} + +func TestBSI64CompareBigValueFallsBackForBigWidth(t *testing.T) { + bsi := NewDefaultBSI() + base := new(big.Int).Lsh(big.NewInt(1), 80) + below := new(big.Int).Sub(base, big.NewInt(1)) + above := new(big.Int).Add(base, big.NewInt(1)) + bsi.SetBigValue(1, below) + bsi.SetBigValue(2, base) + bsi.SetBigValue(3, above) + + eq := bsi.CompareBigValue(0, EQ, base, nil, nil) + assert.True(t, eq.Equals(BitmapOf(2))) + + rng := bsi.CompareBigValue(0, RANGE, base, above, nil) + assert.True(t, rng.Equals(BitmapOf(2, 3))) +} + +func BenchmarkBSI64CompareValueEQLargeAgeFixture(b *testing.B) { + bsi := setupLargeBSI(b) + if bsi == nil { + b.Skip("skipping, large BSI setup failed") + } + b.ResetTimer() + for i := 0; i < b.N; i++ { + res := bsi.CompareValue(0, EQ, 55, 0, nil) + _ = res + } +} + +func BenchmarkBSI64CompareBigValueEQLargeAgeFixture(b *testing.B) { + bsi := setupLargeBSI(b) + if bsi == nil { + b.Skip("skipping, large BSI setup failed") + } + value := big.NewInt(55) + b.ResetTimer() + for i := 0; i < b.N; i++ { + res := bsi.CompareBigValue(0, EQ, value, nil, nil) + _ = res + } +} + +func BenchmarkBSI64CompareValueEQFoundSetLargeAgeFixture(b *testing.B) { + bsi := setupLargeBSI(b) + if bsi == nil { + b.Skip("skipping, large BSI setup failed") + } + foundSet := bsi.CompareValue(0, RANGE, 40, 70, nil) + b.ResetTimer() + for i := 0; i < b.N; i++ { + res := bsi.CompareValue(0, EQ, 55, 0, foundSet) + _ = res + } +} + +func BenchmarkBSI64CompareBigValueEQFoundSetLargeAgeFixture(b *testing.B) { + bsi := setupLargeBSI(b) + if bsi == nil { + b.Skip("skipping, large BSI setup failed") + } + foundSet := bsi.CompareBigValue(0, RANGE, big.NewInt(40), big.NewInt(70), nil) + value := big.NewInt(55) + b.ResetTimer() + for i := 0; i < b.N; i++ { + res := bsi.CompareBigValue(0, EQ, value, nil, foundSet) + _ = res + } +} + +func BenchmarkBSI64CompareValueRangeLargeAgeFixture(b *testing.B) { + bsi := setupLargeBSI(b) + if bsi == nil { + b.Skip("skipping, large BSI setup failed") + } + b.ResetTimer() + for i := 0; i < b.N; i++ { + res := bsi.CompareValue(0, RANGE, 40, 70, nil) + _ = res + } +} + +func BenchmarkBSI64CompareBigValueRangeLargeAgeFixture(b *testing.B) { + bsi := setupLargeBSI(b) + if bsi == nil { + b.Skip("skipping, large BSI setup failed") + } + start := big.NewInt(40) + end := big.NewInt(70) + b.ResetTimer() + for i := 0; i < b.N; i++ { + res := bsi.CompareBigValue(0, RANGE, start, end, nil) + _ = res + } +} + +func BenchmarkBSI64CompareValueGELargeAgeFixture(b *testing.B) { + bsi := setupLargeBSI(b) + if bsi == nil { + b.Skip("skipping, large BSI setup failed") + } + b.ResetTimer() + for i := 0; i < b.N; i++ { + res := bsi.CompareValue(0, GE, 55, 0, nil) + _ = res + } +} + +func BenchmarkBSI64CompareBigValueGELargeAgeFixture(b *testing.B) { + bsi := setupLargeBSI(b) + if bsi == nil { + b.Skip("skipping, large BSI setup failed") + } + value := big.NewInt(55) + b.ResetTimer() + for i := 0; i < b.N; i++ { + res := bsi.CompareBigValue(0, GE, value, nil, nil) + _ = res + } +} diff --git a/roaring64/bsi64_compare_bsi_test.go b/roaring64/bsi64_compare_bsi_test.go new file mode 100644 index 00000000..9d3b7b67 --- /dev/null +++ b/roaring64/bsi64_compare_bsi_test.go @@ -0,0 +1,166 @@ +package roaring64 + +import ( + "math/big" + "math/rand" + "testing" + + "github.com/stretchr/testify/assert" +) + +func expectedBSI64CompareBSI(left *BSI, op Operation, right *BSI, foundSet *Bitmap) *Bitmap { + expected := NewBitmap() + source := And(left.GetExistenceBitmap(), right.GetExistenceBitmap()) + if foundSet != nil { + source.And(foundSet) + } + iter := source.Iterator() + for iter.HasNext() { + col := iter.Next() + leftValue, leftOK := left.GetBigValue(col) + rightValue, rightOK := right.GetBigValue(col) + if !leftOK || !rightOK { + continue + } + compare := leftValue.Cmp(rightValue) + switch op { + case LT: + if compare < 0 { + expected.Add(col) + } + case LE: + if compare <= 0 { + expected.Add(col) + } + case EQ: + if compare == 0 { + expected.Add(col) + } + case GE: + if compare >= 0 { + expected.Add(col) + } + case GT: + if compare > 0 { + expected.Add(col) + } + default: + panic("unsupported test operation") + } + } + return expected +} + +func TestBSI64CompareBSIConsistentWithGetBigValue(t *testing.T) { + rg := rand.New(rand.NewSource(122)) + for run := 0; run < 25; run++ { + left := NewDefaultBSI() + right := NewDefaultBSI() + numCols := rg.Intn(1000) + 50 + for col := 0; col < numCols; col++ { + if rg.Float64() < 0.90 { + left.SetValue(uint64(col), rg.Int63n(2000)-1000) + } + if rg.Float64() < 0.85 { + right.SetValue(uint64(col), rg.Int63n(2000)-1000) + } + } + // Force different bit widths and signs across the two BSIs. + left.SetValue(uint64(numCols+1), 1<<40) + right.SetValue(uint64(numCols+1), -1) + left.SetValue(uint64(numCols+2), -1) + right.SetValue(uint64(numCols+2), 1<<35) + + foundSet := NewBitmap() + source := And(left.GetExistenceBitmap(), right.GetExistenceBitmap()) + iter := source.Iterator() + for iter.HasNext() { + col := iter.Next() + if col%3 != 0 { + foundSet.Add(col) + } + } + + for _, op := range []Operation{LT, LE, EQ, GE, GT} { + for _, fs := range []*Bitmap{nil, foundSet} { + expected := expectedBSI64CompareBSI(left, op, right, fs) + actual := left.CompareBSI(op, right, fs) + assert.True(t, actual.Equals(expected), "run=%d op=%d foundSet=%v expected=%v actual=%v", + run, op, fs != nil, expected.ToArray(), actual.ToArray()) + } + } + } +} + +func TestBSI64CompareBSIExistenceAndResultIsolation(t *testing.T) { + left := NewDefaultBSI() + right := NewDefaultBSI() + left.SetValue(1, 10) + left.SetValue(2, 20) + right.SetValue(2, 15) + right.SetValue(3, 5) + + actual := left.CompareBSI(GT, right, nil) + assert.True(t, actual.Equals(BitmapOf(2))) + + actual.Add(99) + actual.Remove(2) + assert.True(t, left.GetExistenceBitmap().Contains(2)) + assert.True(t, right.GetExistenceBitmap().Contains(2)) + assert.False(t, left.GetExistenceBitmap().Contains(99)) +} + +func TestBSI64CompareBSIBigWidthConsistentWithGetBigValue(t *testing.T) { + left := NewDefaultBSI() + right := NewDefaultBSI() + huge := new(big.Int).Lsh(big.NewInt(1), 90) + hugePlusOne := new(big.Int).Add(huge, big.NewInt(1)) + negativeHuge := new(big.Int).Neg(huge) + negativeHugeMinusOne := new(big.Int).Sub(negativeHuge, big.NewInt(1)) + + left.SetBigValue(1, huge) + right.SetBigValue(1, hugePlusOne) + left.SetBigValue(2, hugePlusOne) + right.SetBigValue(2, huge) + left.SetBigValue(3, negativeHuge) + right.SetBigValue(3, huge) + left.SetBigValue(4, negativeHugeMinusOne) + right.SetBigValue(4, negativeHuge) + left.SetBigValue(5, negativeHuge) + right.SetBigValue(5, negativeHuge) + + assert.True(t, left.CompareBSI(LT, right, nil).Equals(BitmapOf(1, 3, 4))) + assert.True(t, left.CompareBSI(GT, right, nil).Equals(BitmapOf(2))) + assert.True(t, left.CompareBSI(EQ, right, nil).Equals(BitmapOf(5))) + assert.True(t, left.CompareBSI(LE, right, nil).Equals(BitmapOf(1, 3, 4, 5))) + assert.True(t, left.CompareBSI(GE, right, nil).Equals(BitmapOf(2, 5))) +} + +func BenchmarkBSI64CompareBSISameRowBitwise(b *testing.B) { + left, right := setupBSI64CompareBSIFixture(b, 100000) + b.ResetTimer() + for i := 0; i < b.N; i++ { + res := left.CompareBSI(GT, right, nil) + _ = res + } +} + +func BenchmarkBSI64CompareBSISameRowGetBigValue(b *testing.B) { + left, right := setupBSI64CompareBSIFixture(b, 100000) + b.ResetTimer() + for i := 0; i < b.N; i++ { + res := expectedBSI64CompareBSI(left, GT, right, nil) + _ = res + } +} + +func setupBSI64CompareBSIFixture(tb testing.TB, rows int) (*BSI, *BSI) { + tb.Helper() + left := NewDefaultBSI() + right := NewDefaultBSI() + for row := 0; row < rows; row++ { + left.SetValue(uint64(row), int64(row%1000)-500) + right.SetValue(uint64(row), int64((row*7)%1000)-500) + } + return left, right +} diff --git a/roaring64/bsi64_get_big_values_test.go b/roaring64/bsi64_get_big_values_test.go new file mode 100644 index 00000000..38049b2d --- /dev/null +++ b/roaring64/bsi64_get_big_values_test.go @@ -0,0 +1,105 @@ +package roaring64 + +import ( + "math/big" + "math/rand" + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestBSI64GetBigValuesConsistentWithGetBigValue(t *testing.T) { + rg := rand.New(rand.NewSource(864)) + for run := 0; run < 25; run++ { + bsi := NewDefaultBSI() + numCols := rg.Intn(1000) + 50 + for col := 0; col < numCols; col++ { + if rg.Float64() < 0.85 { + bsi.SetValue(uint64(col), rg.Int63n(4000)-2000) + } + } + + columnIDs := make([]uint64, 0, numCols+8) + for col := numCols - 1; col >= 0; col-- { + if col%3 != 0 { + columnIDs = append(columnIDs, uint64(col)) + } + } + columnIDs = append(columnIDs, uint64(numCols+10), 7, 7, 11) + + actual := bsi.GetBigValues(columnIDs) + if len(actual) != len(columnIDs) { + t.Fatalf("run=%d values length = %d, want %d", run, len(actual), len(columnIDs)) + } + for i, columnID := range columnIDs { + expectedValue, expectedOK := bsi.GetBigValue(columnID) + actualValue := actual[i] + if !expectedOK { + assert.Nil(t, actualValue, "run=%d column=%d", run, columnID) + continue + } + if assert.NotNil(t, actualValue, "run=%d column=%d", run, columnID) { + assert.Equal(t, 0, actualValue.Cmp(expectedValue), "run=%d column=%d", run, columnID) + } + } + } +} + +func TestBSI64GetBigValuesHandlesBigWidthAndDuplicates(t *testing.T) { + bsi := NewDefaultBSI() + huge := new(big.Int).Lsh(big.NewInt(1), 90) + hugePlusSeven := new(big.Int).Add(huge, big.NewInt(7)) + negativeHuge := new(big.Int).Neg(hugePlusSeven) + + bsi.SetBigValue(1, hugePlusSeven) + bsi.SetBigValue(2, negativeHuge) + bsi.SetValue(4, 0) + + values := bsi.GetBigValues([]uint64{2, 3, 1, 2, 4}) + assert.Equal(t, 5, len(values)) + assert.Equal(t, 0, values[0].Cmp(negativeHuge)) + assert.Nil(t, values[1]) + assert.Equal(t, 0, values[2].Cmp(hugePlusSeven)) + assert.Equal(t, 0, values[3].Cmp(negativeHuge)) + assert.Equal(t, 0, values[4].Cmp(big.NewInt(0))) + + values[0].SetInt64(12) + assert.Equal(t, 0, values[3].Cmp(negativeHuge), "duplicate result values should be independent") + stored, ok := bsi.GetBigValue(2) + assert.True(t, ok) + assert.Equal(t, 0, stored.Cmp(negativeHuge), "mutating returned values must not alter the BSI") +} + +func BenchmarkBSI64GetBigValuesLargeFixture(b *testing.B) { + bsi, _ := setupBSI64CompareBSIFixture(b, 100000) + columnIDs := bsi64SequentialColumns(100000) + b.ResetTimer() + for i := 0; i < b.N; i++ { + values := bsi.GetBigValues(columnIDs) + _ = values + } +} + +func BenchmarkBSI64GetBigValueLoopLargeFixture(b *testing.B) { + bsi, _ := setupBSI64CompareBSIFixture(b, 100000) + columnIDs := bsi64SequentialColumns(100000) + b.ResetTimer() + for i := 0; i < b.N; i++ { + values := make([]*big.Int, len(columnIDs)) + for j, columnID := range columnIDs { + value, ok := bsi.GetBigValue(columnID) + if ok { + values[j] = value + } + } + _ = values + } +} + +func bsi64SequentialColumns(n int) []uint64 { + columnIDs := make([]uint64, n) + for i := range columnIDs { + columnIDs[i] = uint64(i) + } + return columnIDs +} diff --git a/roaring64/roaring64.go b/roaring64/roaring64.go index adb508fe..143209de 100644 --- a/roaring64/roaring64.go +++ b/roaring64/roaring64.go @@ -1241,6 +1241,7 @@ func (rb *Bitmap) GetSerializedSizeInBytes() uint64 { return rb.highlowcontainer.serializedSizeInBytes() } +// Validate checks whether the bitmap's internal containers are coherent. func (rb *Bitmap) Validate() error { return rb.highlowcontainer.validate() } diff --git a/roaring64/roaringarray64.go b/roaring64/roaringarray64.go index 09c366ff..f974f27d 100644 --- a/roaring64/roaringarray64.go +++ b/roaring64/roaringarray64.go @@ -14,7 +14,9 @@ type roaringArray64 struct { } var ( - ErrKeySortOrder = errors.New("keys were out of order") + // ErrKeySortOrder reports that container keys are out of order. + ErrKeySortOrder = errors.New("keys were out of order") + // ErrCardinalityConstraint reports inconsistent array cardinality metadata. ErrCardinalityConstraint = errors.New("size of arrays was not coherent") ) diff --git a/roaring_test.go b/roaring_test.go index d25476b3..6b31c61f 100644 --- a/roaring_test.go +++ b/roaring_test.go @@ -2682,21 +2682,6 @@ func TestRoaringArray(t *testing.T) { assert.Equal(t, 0, a.size()) }) - t.Run("Test popcount Full", func(t *testing.T) { - res := popcount(uint64(0xffffffffffffffff)) - assert.EqualValues(t, 64, res) - }) - - t.Run("Test popcount Empty", func(t *testing.T) { - res := popcount(0) - assert.EqualValues(t, 0, res) - }) - - t.Run("Test popcount 16", func(t *testing.T) { - res := popcount(0xff00ff) - assert.EqualValues(t, 16, res) - }) - t.Run("Test ArrayContainer Add", func(t *testing.T) { ar := newArrayContainer() ar.iadd(1) @@ -3975,3 +3960,209 @@ func TestValidateFromV1(t *testing.T) { require.NoError(t, v2.UnmarshalBinary(b)) require.NoError(t, v2.Validate()) } + +// bitmapOrBulkMergeFixture holds two variants of a left/right bitmap pair so +// benchmarks can alternate between them and avoid mutating a single receiver +// across iterations. The fixtures below exercise the in-place Bitmap.Or bulk +// merge path (many interleaved source-only keys). +type bitmapOrBulkMergeFixture struct { + lefts [2]*Bitmap + rights [2]*Bitmap + cardinality uint64 +} + +func newBitmapOrBulkMergeFixture(leftKeys, rightKeys []uint16, copyOnWrite bool) bitmapOrBulkMergeFixture { + fixture := bitmapOrBulkMergeFixture{} + for variant := range fixture.lefts { + left := NewBitmap() + right := NewBitmap() + leftLow := uint16(variant * 2) + rightLow := leftLow + 1 + for _, key := range leftKeys { + left.Add(uint32(key)<<16 | uint32(leftLow)) + } + for _, key := range rightKeys { + right.Add(uint32(key)<<16 | uint32(rightLow)) + } + if copyOnWrite { + left.SetCopyOnWrite(true) + right.SetCopyOnWrite(true) + } + fixture.lefts[variant] = left + fixture.rights[variant] = right + } + fixture.cardinality = bitmapOrBulkMergeExpected(fixture.lefts[0], fixture.rights[0]).GetCardinality() + return fixture +} + +func bitmapOrBulkMergeExpected(left, right *Bitmap) *Bitmap { + values := make([]uint32, 0, left.GetCardinality()+right.GetCardinality()) + values = append(values, left.ToArray()...) + values = append(values, right.ToArray()...) + return BitmapOf(values...) +} + +func bitmapOrBulkMergeKeys(start, count, step int) []uint16 { + keys := make([]uint16, count) + for i := range keys { + keys[i] = uint16(start + i*step) + } + return keys +} + +func bitmapOrBulkMergeInterleavedFixture(containers int, copyOnWrite bool) bitmapOrBulkMergeFixture { + return newBitmapOrBulkMergeFixture( + bitmapOrBulkMergeKeys(0, containers, 2), + bitmapOrBulkMergeKeys(1, containers, 2), + copyOnWrite, + ) +} + +func bitmapOrBulkMergeAppendFixture(containers int) bitmapOrBulkMergeFixture { + return newBitmapOrBulkMergeFixture( + bitmapOrBulkMergeKeys(0, containers, 1), + bitmapOrBulkMergeKeys(containers, containers, 1), + false, + ) +} + +func bitmapOrBulkMergeOverlapFixture(containers int) bitmapOrBulkMergeFixture { + keys := bitmapOrBulkMergeKeys(0, containers, 1) + return newBitmapOrBulkMergeFixture(keys, keys, false) +} + +func bitmapOrBulkMergeSingleInteriorFixture(containers int) bitmapOrBulkMergeFixture { + leftKeys := make([]uint16, 0, containers-1) + middle := containers / 2 + for key := 0; key < containers; key++ { + if key != middle { + leftKeys = append(leftKeys, uint16(key)) + } + } + return newBitmapOrBulkMergeFixture(leftKeys, []uint16{uint16(middle)}, false) +} + +func bitmapOrBulkMergeFixtureCases() map[string]bitmapOrBulkMergeFixture { + return map[string]bitmapOrBulkMergeFixture{ + "interleaved-64": bitmapOrBulkMergeInterleavedFixture(64, false), + "interleaved-65": bitmapOrBulkMergeInterleavedFixture(65, false), + "interleaved-1024": bitmapOrBulkMergeInterleavedFixture(1024, false), + "interleaved-4096": bitmapOrBulkMergeInterleavedFixture(4096, false), + "append-only-4096": bitmapOrBulkMergeAppendFixture(4096), + "overlap-4096": bitmapOrBulkMergeOverlapFixture(4096), + "copy-on-write-interleaved-4096": bitmapOrBulkMergeInterleavedFixture(4096, true), + "single-interior-4096": bitmapOrBulkMergeSingleInteriorFixture(4096), + "mixed-key-order": newBitmapOrBulkMergeFixture( + []uint16{2, 4, 6, 8}, + []uint16{1, 4, 5, 8, 9}, + false, + ), + } +} + +func TestBitmapOrBulkMergeFixtures(t *testing.T) { + for name, fixture := range bitmapOrBulkMergeFixtureCases() { + t.Run(name, func(t *testing.T) { + left := fixture.lefts[0] + right := fixture.rights[0] + want := bitmapOrBulkMergeExpected(left, right) + receiver := left.Clone() + receiver.Or(right) + + if !receiver.Equals(want) { + t.Fatalf("unexpected union: got %v, want %v", receiver, want) + } + if receiver.GetCardinality() != fixture.cardinality { + t.Fatalf("unexpected cardinality: got %d, want %d", receiver.GetCardinality(), fixture.cardinality) + } + if err := receiver.Validate(); err != nil { + t.Fatalf("union produced an invalid bitmap: %v", err) + } + }) + } +} + +func TestBitmapOrBulkMergeCopyOnWriteOwnership(t *testing.T) { + fixture := bitmapOrBulkMergeInterleavedFixture(64, true) + left := fixture.lefts[0] + right := fixture.rights[0] + receiver := left.Clone() + receiver.Or(right) + + const ( + leftKey = uint32(0) << 16 + rightKey = uint32(1) << 16 + receiver1 = uint32(10) + receiver2 = uint32(11) + ) + + receiver.Add(rightKey | receiver1) + if right.Contains(rightKey | receiver1) { + t.Fatal("receiver mutation changed a source-only container") + } + right.Add(rightKey | receiver2) + if receiver.Contains(rightKey | receiver2) { + t.Fatal("source mutation changed a receiver source-only container") + } + + receiver.Add(leftKey | receiver1) + if left.Contains(leftKey | receiver1) { + t.Fatal("receiver mutation changed a receiver-only container") + } + left.Add(leftKey | receiver2) + if receiver.Contains(leftKey | receiver2) { + t.Fatal("left mutation changed a receiver container") + } + + if err := receiver.Validate(); err != nil { + t.Fatalf("receiver became invalid after copy-on-write mutations: %v", err) + } + if err := left.Validate(); err != nil { + t.Fatalf("left became invalid after copy-on-write mutations: %v", err) + } + if err := right.Validate(); err != nil { + t.Fatalf("right became invalid after copy-on-write mutations: %v", err) + } +} + +func TestBitmapOrBulkMergeCopyOnWriteTailOwnership(t *testing.T) { + fixture := bitmapOrBulkMergeInterleavedFixture(64, true) + left := fixture.lefts[0] + right := fixture.rights[0] + receiver := left.Clone() + receiver.Or(right) + + tailIndex := right.highlowcontainer.size() - 1 + tailKey := right.highlowcontainer.getKeyAtIndex(tailIndex) + receiverTailIndex := receiver.highlowcontainer.getIndex(tailKey) + if receiverTailIndex < 0 { + t.Fatal("receiver is missing the source-only tail container") + } + if !right.highlowcontainer.needsCopyOnWrite(tailIndex) { + t.Fatal("source-only tail container was not marked copy-on-write") + } + if !receiver.highlowcontainer.needsCopyOnWrite(receiverTailIndex) { + t.Fatal("receiver tail container was not marked copy-on-write") + } + if receiver.highlowcontainer.getContainerAtIndex(receiverTailIndex) != right.highlowcontainer.getContainerAtIndex(tailIndex) { + t.Fatal("source-only tail container was not shared") + } + + receiverValue := uint32(tailKey)<<16 | 10 + sourceValue := uint32(tailKey)<<16 | 11 + receiver.Add(receiverValue) + if right.Contains(receiverValue) { + t.Fatal("receiver tail mutation changed the source") + } + right.Add(sourceValue) + if receiver.Contains(sourceValue) { + t.Fatal("source tail mutation changed the receiver") + } + + if err := receiver.Validate(); err != nil { + t.Fatalf("receiver became invalid after tail mutations: %v", err) + } + if err := right.Validate(); err != nil { + t.Fatalf("source became invalid after tail mutations: %v", err) + } +} diff --git a/roaringarray.go b/roaringarray.go index a1a71784..dfe80c4c 100644 --- a/roaringarray.go +++ b/roaringarray.go @@ -384,6 +384,122 @@ func (ra *roaringArray) insertNewKeyValueAt(i int, key uint16, value container) ra.needCopyOnWrite[i] = false } +// copyOrSourceContainerAt returns the container (and its copy-on-write flag) to +// store for a source-only key. Keys beyond the receiver's last key are the +// trailing suffix: they may be shared under copy-on-write, matching appendCopy. +// Interior source-only keys are always cloned so that later receiver mutations +// cannot leak into the source. +func (ra *roaringArray) copyOrSourceContainerAt(other *roaringArray, index int, receiverLastKey uint16) (container, bool) { + if other.keys[index] > receiverLastKey { + copyOnWrite := (ra.copyOnWrite && other.copyOnWrite) || other.needsCopyOnWrite(index) + if copyOnWrite { + if !other.needsCopyOnWrite(index) { + other.setNeedsCopyOnWrite(index) + } + return other.containers[index], true + } + } + return other.containers[index].clone(), false +} + +// mergeBulk finishes an in-place union (xor == false) or symmetric difference +// (xor == true) once the receiver's structure must change and continuing in +// place would shift the aligned suffix once per changed key -- quadratic when +// many keys are interleaved. It merges the two suffixes forward into fresh +// slices in a single pass instead. +// +// The change that triggers it is a source-only key that must be inserted, or +// (xor only) an aligned pair that cancelled to an empty container. dst is the +// write cursor: the prefix [0, dst) is already final and copied over unchanged. +// left/right are the receiver/source scan positions; the caller advances them +// past an already-consumed aligned-empty pair. For a union a source-only key is +// always inserted, so dst == left; the xor caller may pass dst < left to drop +// the emptied container. +// +// Like every other roaringArray operation it assumes both arrays already hold +// their keys in sorted order; that invariant is enforced at the load boundary +// (Validate), not re-checked here. +func (ra *roaringArray) mergeBulk(other *roaringArray, dst, left, right int, xor bool) { + length1 := ra.size() + length2 := other.size() + receiverLastKey := ra.keys[length1-1] + + // First pass over the keys only (cheap, no container work): count the + // distinct keys of the two suffixes. That is the exact result size for a + // union and, for a xor, a tight upper bound (aligned pairs may cancel). So + // the appends below never reallocate, without grossly over-allocating when + // many aligned containers cancel to empty. + distinct := 0 + l, r := left, right + for l < length1 && r < length2 { + if ra.keys[l] < other.keys[r] { + l++ + } else if ra.keys[l] > other.keys[r] { + r++ + } else { + l++ + r++ + } + distinct++ + } + distinct += (length1 - l) + (length2 - r) + total := dst + distinct + keys := make([]uint16, dst, total) + containers := make([]container, dst, total) + needCopyOnWrite := make([]bool, dst, total) + copy(keys, ra.keys[:dst]) + copy(containers, ra.containers[:dst]) + copy(needCopyOnWrite, ra.needCopyOnWrite[:dst]) + + for left < length1 && right < length2 { + s1 := ra.keys[left] + s2 := other.keys[right] + if s1 < s2 { + keys = append(keys, s1) + containers = append(containers, ra.containers[left]) + needCopyOnWrite = append(needCopyOnWrite, ra.needCopyOnWrite[left]) + left++ + } else if s1 > s2 { + c, cow := ra.copyOrSourceContainerAt(other, right, receiverLastKey) + keys = append(keys, s2) + containers = append(containers, c) + needCopyOnWrite = append(needCopyOnWrite, cow) + right++ + } else { + // Union of two non-empty containers is never empty, so the + // isEmpty check only ever drops a container for xor. + var c container + if xor { + c = ra.getWritableContainerAtIndex(left).ixor(other.containers[right]) + } else { + c = ra.getUnionedWritableContainer(left, other.containers[right]) + } + if !c.isEmpty() { + keys = append(keys, s1) + containers = append(containers, c) + needCopyOnWrite = append(needCopyOnWrite, false) + } + left++ + right++ + } + } + for ; left < length1; left++ { + keys = append(keys, ra.keys[left]) + containers = append(containers, ra.containers[left]) + needCopyOnWrite = append(needCopyOnWrite, ra.needCopyOnWrite[left]) + } + for ; right < length2; right++ { + c, cow := ra.copyOrSourceContainerAt(other, right, receiverLastKey) + keys = append(keys, other.keys[right]) + containers = append(containers, c) + needCopyOnWrite = append(needCopyOnWrite, cow) + } + + ra.keys = keys + ra.containers = containers + ra.needCopyOnWrite = needCopyOnWrite +} + func (ra *roaringArray) remove(key uint16) bool { i := ra.binarySearch(0, int64(len(ra.keys)), key) if i >= 0 { // if a new key diff --git a/roaringcow_test.go b/roaringcow_test.go index c04acae4..6083d919 100644 --- a/roaringcow_test.go +++ b/roaringcow_test.go @@ -1604,21 +1604,6 @@ func TestRoaringArrayCOW(t *testing.T) { assert.Equal(t, 0, a.size()) }) - t.Run("Test popcount Full", func(t *testing.T) { - res := popcount(uint64(0xffffffffffffffff)) - assert.EqualValues(t, 64, res) - }) - - t.Run("Test popcount Empty", func(t *testing.T) { - res := popcount(0) - assert.EqualValues(t, 0, res) - }) - - t.Run("Test popcount 16", func(t *testing.T) { - res := popcount(0xff00ff) - assert.EqualValues(t, 16, res) - }) - t.Run("Test ArrayContainer Add", func(t *testing.T) { ar := newArrayContainer() ar.iadd(1) diff --git a/runcontainer.go b/runcontainer.go index e5ff8857..7c369b75 100644 --- a/runcontainer.go +++ b/runcontainer.go @@ -41,6 +41,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. import ( "errors" "fmt" + "math/bits" "slices" ) @@ -220,7 +221,7 @@ func newRunContainer16FromBitmapContainer(bc *bitmapContainer) *runContainer16 { // wrap up, no more runs return rc } - localRunStart := countTrailingZeros(curWord) + localRunStart := bits.TrailingZeros64(curWord) runStart := localRunStart + 64*longCtr // stuff 1s into number's LSBs curWordWith1s := curWord | (curWord - 1) @@ -239,7 +240,7 @@ func newRunContainer16FromBitmapContainer(bc *bitmapContainer) *runContainer16 { rc.iv[runCount].length = uint16(runEnd) - uint16(runStart) - 1 return rc } - localRunEnd := countTrailingZeros(^curWordWith1s) + localRunEnd := bits.TrailingZeros64(^curWordWith1s) runEnd = localRunEnd + longCtr*64 rc.iv[runCount].start = uint16(runStart) rc.iv[runCount].length = uint16(runEnd) - 1 - uint16(runStart) diff --git a/util.go b/util.go index 031dfa30..e727d6e0 100644 --- a/util.go +++ b/util.go @@ -3,6 +3,7 @@ package roaring import ( "cmp" "math" + "math/bits" "math/rand" "slices" ) @@ -70,7 +71,7 @@ func fillArrayAND(container []uint16, bitmap1, bitmap2 []uint64) { bitset := bitmap1[k] & bitmap2[k] for bitset != 0 { t := bitset & -bitset - container[pos] = uint16((k*64 + int(popcount(t-1)))) + container[pos] = uint16((k*64 + bits.OnesCount64(t-1))) pos = pos + 1 bitset ^= t } @@ -87,7 +88,7 @@ func fillArrayANDNOT(container []uint16, bitmap1, bitmap2 []uint64) { bitset := bitmap1[k] &^ bitmap2[k] for bitset != 0 { t := bitset & -bitset - container[pos] = uint16((k*64 + int(popcount(t-1)))) + container[pos] = uint16((k*64 + bits.OnesCount64(t-1))) pos = pos + 1 bitset ^= t } @@ -104,7 +105,7 @@ func fillArrayXOR(container []uint16, bitmap1, bitmap2 []uint64) { bitset := bitmap1[k] ^ bitmap2[k] for bitset != 0 { t := bitset & -bitset - container[pos] = uint16((k*64 + int(popcount(t-1)))) + container[pos] = uint16((k*64 + bits.OnesCount64(t-1))) pos = pos + 1 bitset ^= t } @@ -205,7 +206,7 @@ func wordCardinalityForBitmapRange(bitmap []uint64, start int, end int) uint64 { firstword := start / 64 endword := (end - 1) / 64 for i := firstword; i <= endword; i++ { - answer += popcount(bitmap[i]) + answer += uint64(bits.OnesCount64(bitmap[i])) } return answer } @@ -215,31 +216,31 @@ func selectBitPosition(w uint64, j int) int { // Divide 64bit part := w & 0xFFFFFFFF - n := popcount(part) - if n <= uint64(j) { + n := bits.OnesCount64(part) + if n <= j { part = w >> 32 seen += 32 - j -= int(n) + j -= n } w = part // Divide 32bit part = w & 0xFFFF - n = popcount(part) - if n <= uint64(j) { + n = bits.OnesCount64(part) + if n <= j { part = w >> 16 seen += 16 - j -= int(n) + j -= n } w = part // Divide 16bit part = w & 0xFF - n = popcount(part) - if n <= uint64(j) { + n = bits.OnesCount64(part) + if n <= j { part = w >> 8 seen += 8 - j -= int(n) + j -= n } w = part From aa3338d1e2f0b2fa321f6b500dc1587dfe318738 Mon Sep 17 00:00:00 2001 From: Perfloop Agent Date: Tue, 4 Aug 2026 01:36:43 +0000 Subject: [PATCH 2/2] perf(BSI): optimize populated run-optimized query results --- BitSliceIndexing/bsi.go | 26 ++++++++++---------------- 1 file changed, 10 insertions(+), 16 deletions(-) diff --git a/BitSliceIndexing/bsi.go b/BitSliceIndexing/bsi.go index 6b2ddc56..36598c2e 100644 --- a/BitSliceIndexing/bsi.go +++ b/BitSliceIndexing/bsi.go @@ -97,9 +97,6 @@ func (b *BSI) SetValue(columnID uint64, value int64) { if b.MaxValue == 0 && b.MinValue == 0 { for i := bits.Len64(uint64(value)) - b.BitCount(); i > 0; i-- { b.bA = append(b.bA, roaring.NewBitmap()) - if b.runOptimized { - b.bA[i].RunOptimize() - } } } @@ -121,9 +118,6 @@ func (b *BSI) SetMany(foundSet *roaring.Bitmap, value int64) { if b.MaxValue == 0 && b.MinValue == 0 { for i := bits.Len64(uint64(value)) - b.BitCount(); i > 0; i-- { b.bA = append(b.bA, roaring.NewBitmap()) - if b.runOptimized { - b.bA[i].RunOptimize() - } } } @@ -191,7 +185,12 @@ func parallelExecutor(parallelism int, t *task, e action, ba = append(ba, bm) } - return roaring.ParOr(0, ba...) + results := roaring.ParOr(0, ba...) + // Optimize the aggregate returned to the caller after it has been populated. + if t.bsi.runOptimized && !results.IsEmpty() { + results.RunOptimize() + } + return results } @@ -241,6 +240,10 @@ func parallelExecutorBSIResults(parallelism int, input *BSI, e bsiAction, foundS } else { results.ParOr(0, ba...) } + // Optimize the aggregate returned to the caller after it has been populated. + if input.runOptimized && !results.eBM.IsEmpty() { + results.RunOptimize() + } return results } @@ -297,9 +300,6 @@ func compareValue(e *task, batch []uint32, resultsChan chan *roaring.Bitmap, wg defer wg.Done() results := roaring.NewBitmap() - if e.bsi.runOptimized { - results.RunOptimize() - } if len(batch) == 0 { resultsChan <- results return @@ -610,9 +610,6 @@ func transpose(e *task, batch []uint32, resultsChan chan *roaring.Bitmap, wg *sy defer wg.Done() results := roaring.NewBitmap() - if e.bsi.runOptimized { - results.RunOptimize() - } for _, cID := range batch { if value, ok := e.bsi.GetValue(uint64(cID)); ok { results.Add(uint32(value)) @@ -1031,9 +1028,6 @@ func transposeWithCounts(input *BSI, batch []uint32, resultsChan chan *BSI, wg * defer wg.Done() results := NewDefaultBSI() - if input.runOptimized { - results.RunOptimize() - } for _, cID := range batch { if value, ok := input.GetValue(uint64(cID)); ok { if val, ok2 := results.GetValue(uint64(value)); !ok2 {