From 62fc9479b242525207caa6e4756b31c0c4124f5b Mon Sep 17 00:00:00 2001
From: Aliaksandr Valialkin <valyala@victoriametrics.com>
Date: Sat, 8 Jun 2024 01:15:58 +0200
Subject: [PATCH] wip

---
 lib/logstorage/bitmap.go             | 18 +++++-----
 lib/logstorage/bitmap_test.go        |  8 +++++
 lib/logstorage/bitmap_timing_test.go | 51 +++++++++++++++++++++++++---
 3 files changed, 64 insertions(+), 13 deletions(-)

diff --git a/lib/logstorage/bitmap.go b/lib/logstorage/bitmap.go
index ac1c10695..350ea0cc8 100644
--- a/lib/logstorage/bitmap.go
+++ b/lib/logstorage/bitmap.go
@@ -98,13 +98,6 @@ func (bm *bitmap) areAllBitsSet() bool {
 	return true
 }
 
-func (bm *bitmap) isSetBit(i int) bool {
-	wordIdx := uint(i) / 64
-	wordOffset := uint(i) % 64
-	word := bm.a[wordIdx]
-	return (word & (1 << wordOffset)) != 0
-}
-
 func (bm *bitmap) andNot(x *bitmap) {
 	if bm.bitsLen != x.bitsLen {
 		logger.Panicf("BUG: cannot merge bitmaps with distinct lengths; %d vs %d", bm.bitsLen, x.bitsLen)
@@ -127,6 +120,13 @@ func (bm *bitmap) or(x *bitmap) {
 	}
 }
 
+func (bm *bitmap) isSetBit(i int) bool {
+	wordIdx := uint(i) / 64
+	wordOffset := uint(i) % 64
+	word := bm.a[wordIdx]
+	return (word & (1 << wordOffset)) != 0
+}
+
 // forEachSetBit calls f for each set bit and clears that bit if f returns false
 func (bm *bitmap) forEachSetBit(f func(idx int) bool) {
 	a := bm.a
@@ -143,7 +143,7 @@ func (bm *bitmap) forEachSetBit(f func(idx int) bool) {
 			}
 			idx := i*64 + j
 			if idx >= bitsLen {
-				break
+				return
 			}
 			if !f(idx) {
 				wordNew &= ^mask
@@ -178,7 +178,7 @@ func (bm *bitmap) forEachSetBitReadonly(f func(idx int)) {
 			}
 			idx := i*64 + j
 			if idx >= bitsLen {
-				break
+				return
 			}
 			f(idx)
 		}
diff --git a/lib/logstorage/bitmap_test.go b/lib/logstorage/bitmap_test.go
index 1b5f2ba8b..41cce6974 100644
--- a/lib/logstorage/bitmap_test.go
+++ b/lib/logstorage/bitmap_test.go
@@ -6,6 +6,8 @@ import (
 
 func TestBitmap(t *testing.T) {
 	for i := 0; i < 100; i++ {
+		bitsLen := i
+
 		bm := getBitmap(i)
 		if bm.bitsLen != i {
 			t.Fatalf("unexpected bits length: %d; want %d", bm.bitsLen, i)
@@ -41,6 +43,9 @@ func TestBitmap(t *testing.T) {
 			}
 			nextIdx++
 		})
+		if nextIdx != bitsLen {
+			t.Fatalf("unexpected number of bits set; got %d; want %d", nextIdx, bitsLen)
+		}
 
 		if !bm.areAllBitsSet() {
 			t.Fatalf("all bits must be set for bitmap with %d bits", i)
@@ -71,6 +76,9 @@ func TestBitmap(t *testing.T) {
 			}
 			nextIdx += 2
 		})
+		if nextIdx < bitsLen {
+			t.Fatalf("unexpected number of bits visited; got %d; want %d", nextIdx, bitsLen)
+		}
 
 		// Clear all the bits
 		bm.forEachSetBit(func(_ int) bool {
diff --git a/lib/logstorage/bitmap_timing_test.go b/lib/logstorage/bitmap_timing_test.go
index 9f364290e..5e0a36f6a 100644
--- a/lib/logstorage/bitmap_timing_test.go
+++ b/lib/logstorage/bitmap_timing_test.go
@@ -4,6 +4,35 @@ import (
 	"testing"
 )
 
+func BenchmarkBitmapIsSetBit(b *testing.B) {
+	const bitsLen = 64 * 1024
+
+	b.Run("no-zero-bits", func(b *testing.B) {
+		bm := getBitmap(bitsLen)
+		bm.setBits()
+		benchmarkBitmapIsSetBit(b, bm)
+		putBitmap(bm)
+	})
+	b.Run("half-zero-bits", func(b *testing.B) {
+		bm := getBitmap(bitsLen)
+		bm.setBits()
+		bm.forEachSetBit(func(idx int) bool {
+			return idx%2 == 0
+		})
+		benchmarkBitmapIsSetBit(b, bm)
+		putBitmap(bm)
+	})
+	b.Run("one-set-bit", func(b *testing.B) {
+		bm := getBitmap(bitsLen)
+		bm.setBits()
+		bm.forEachSetBit(func(idx int) bool {
+			return idx == bitsLen/2
+		})
+		benchmarkBitmapIsSetBit(b, bm)
+		putBitmap(bm)
+	})
+}
+
 func BenchmarkBitmapForEachSetBitReadonly(b *testing.B) {
 	const bitsLen = 64 * 1024
 
@@ -86,19 +115,33 @@ func BenchmarkBitmapForEachSetBit(b *testing.B) {
 	})
 }
 
+func benchmarkBitmapIsSetBit(b *testing.B, bm *bitmap) {
+	bitsLen := bm.bitsLen
+	b.SetBytes(int64(bitsLen))
+	b.ReportAllocs()
+	b.RunParallel(func(pb *testing.PB) {
+		n := 0
+		for pb.Next() {
+			for i := 0; i < bitsLen; i++ {
+				if bm.isSetBit(i) {
+					n++
+				}
+			}
+		}
+		GlobalSink.Add(uint64(n))
+	})
+}
+
 func benchmarkBitmapForEachSetBitReadonly(b *testing.B, bm *bitmap) {
 	b.SetBytes(int64(bm.bitsLen))
 	b.ReportAllocs()
 	b.RunParallel(func(pb *testing.PB) {
-		bmLocal := getBitmap(bm.bitsLen)
 		n := 0
 		for pb.Next() {
-			bmLocal.copyFrom(bm)
-			bmLocal.forEachSetBitReadonly(func(_ int) {
+			bm.forEachSetBitReadonly(func(_ int) {
 				n++
 			})
 		}
-		putBitmap(bmLocal)
 		GlobalSink.Add(uint64(n))
 	})
 }