From a967130bae71f1af57678cbfd7c4c5dafcb811cb Mon Sep 17 00:00:00 2001
From: Konstantin Bogomolov <bogomolov@google.com>
Date: Mon, 24 Jun 2024 09:39:16 -0700
Subject: [PATCH] Add Reset method to Bitmap.

Useful for fully clearing bitmap efficiently without reallocation.

PiperOrigin-RevId: 646125349
---
 pkg/bitmap/bitmap.go      | 8 ++++++++
 pkg/bitmap/bitmap_test.go | 7 +++++++
 2 files changed, 15 insertions(+)

diff --git a/pkg/bitmap/bitmap.go b/pkg/bitmap/bitmap.go
index b9eae562d..96bb3ca79 100644
--- a/pkg/bitmap/bitmap.go
+++ b/pkg/bitmap/bitmap.go
@@ -250,6 +250,14 @@ func (b *Bitmap) FlipRange(begin, end uint32) {
 	}
 }
 
+// Reset zeroes the entire bitmap.
+func (b *Bitmap) Reset() {
+	b.numOnes = 0
+	for i := range b.bitBlock {
+		b.bitBlock[i] = 0
+	}
+}
+
 // ForEach calls `f` for each set bit in the range [start, end).
 //
 // If f returns false, ForEach stops the iteration.
diff --git a/pkg/bitmap/bitmap_test.go b/pkg/bitmap/bitmap_test.go
index 082495aff..1712aa6da 100644
--- a/pkg/bitmap/bitmap_test.go
+++ b/pkg/bitmap/bitmap_test.go
@@ -292,6 +292,13 @@ func TestBitmapNumOnes(t *testing.T) {
 	if bitmapOnes != uint32(20) {
 		t.Errorf("After ClearRange, GetNumOnes() returns: %v, wanted: %v", bitmapOnes, 20)
 	}
+
+	// Reset to fully clear bitmap.
+	bitmap.Reset()
+	bitmapOnes = bitmap.GetNumOnes()
+	if bitmapOnes != uint32(0) {
+		t.Errorf("After Reset, GetNumOnes() returns: %v, wanted: %v", bitmapOnes, 0)
+	}
 }
 
 type bitmapForEachTestcase struct {