From 6aaaca0dc986a40ea4caedf8f1428f1619e59a29 Mon Sep 17 00:00:00 2001
From: Andrzej Liszka <andrzej.liszka@gmail.com>
Date: Thu, 5 Dec 2024 20:08:01 +0100
Subject: [PATCH] clone to buffer

---
 benchmark_opt_test.go | 19 +++++++++++++++++++
 bitmap_opt.go         | 20 ++++++++++++++++++++
 bitmap_opt_test.go    | 44 +++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 83 insertions(+)

diff --git a/benchmark_opt_test.go b/benchmark_opt_test.go
index 9135b57..e272262 100644
--- a/benchmark_opt_test.go
+++ b/benchmark_opt_test.go
@@ -5,6 +5,18 @@ import (
 	"testing"
 )
 
+var prefilled *Bitmap
+var prefilledBuf []byte
+
+// func init() {
+// 	initPrefilled()
+// }
+
+func initPrefilled() {
+	prefilled = Prefill(200_000_000)
+	prefilledBuf = make([]byte, prefilled.LenBytes())
+}
+
 // go test -v -bench BenchmarkPrefillNative -benchmem -run ^$ github.com/weaviate/sroar -cpuprofile cpu.prof
 func BenchmarkPrefillNative(b *testing.B) {
 	for i := 0; i < b.N; i++ {
@@ -12,6 +24,13 @@ func BenchmarkPrefillNative(b *testing.B) {
 	}
 }
 
+// go test -v -bench BenchmarkPrefillCloneToBuf -benchmem -run ^$ github.com/weaviate/sroar -cpuprofile cpu.prof
+func BenchmarkPrefillCloneToBuf(b *testing.B) {
+	for i := 0; i < b.N; i++ {
+		prefilled.CloneToBuf(prefilledBuf)
+	}
+}
+
 // go test -v -bench BenchmarkPrefillFromSortedList -benchmem -run ^$ github.com/weaviate/sroar -cpuprofile cpu.prof
 func BenchmarkPrefillFromSortedList(b *testing.B) {
 	prefillBufferSize := 65_536
diff --git a/bitmap_opt.go b/bitmap_opt.go
index a0547e6..d8b22c4 100644
--- a/bitmap_opt.go
+++ b/bitmap_opt.go
@@ -1,6 +1,7 @@
 package sroar
 
 import (
+	"fmt"
 	"math"
 	"sync"
 )
@@ -746,3 +747,22 @@ func (ra *Bitmap) CapBytes() int {
 	}
 	return cap(ra.data) * 2
 }
+
+func (ra *Bitmap) CloneToBuf(buf []byte) *Bitmap {
+	if len(buf)%2 != 0 {
+		panic(fmt.Sprintf("Buffer size should be even, given %d", len(buf)))
+	}
+
+	a := ra
+	if ra == nil {
+		a = NewBitmap()
+	}
+
+	if alen := a.LenBytes(); alen > len(buf) {
+		panic(fmt.Sprintf("Buffer too small, given %d, required %d", len(buf), alen))
+	}
+
+	abuf := toByteSlice(a.data)
+	copy(buf, abuf)
+	return FromBuffer(buf)
+}
diff --git a/bitmap_opt_test.go b/bitmap_opt_test.go
index d98fd5b..27e4074 100644
--- a/bitmap_opt_test.go
+++ b/bitmap_opt_test.go
@@ -930,6 +930,50 @@ func TestCapBytes(t *testing.T) {
 	})
 }
 
+func TestCloneToBuf(t *testing.T) {
+	t.Run("non-nil bitmap", func(t *testing.T) {
+		bmEmpty := NewBitmap()
+
+		bm1 := NewBitmap()
+		bm1.Set(1)
+
+		bm2 := NewBitmap()
+		bm2.Set(1)
+		bm2.Set(1 + uint64(maxCardinality))
+
+		bm3 := NewBitmap()
+		bm3.Set(1)
+		bm3.Set(1 + uint64(maxCardinality))
+		bm3.Set(1 + uint64(maxCardinality)*2)
+
+		for name, bm := range map[string]*Bitmap{
+			"empty": bmEmpty,
+			"bm1":   bm1,
+			"bm2":   bm2,
+			"bm3":   bm3,
+		} {
+			t.Run(name, func(t *testing.T) {
+				buf := make([]byte, bm.LenBytes())
+				cloned := bm.CloneToBuf(buf)
+
+				require.Equal(t, bm.GetCardinality(), cloned.GetCardinality())
+				require.Equal(t, bm.LenBytes(), cloned.LenBytes())
+			})
+		}
+	})
+
+	t.Run("nil bitmap", func(t *testing.T) {
+		var bmNil *Bitmap
+		bmEmpty := NewBitmap()
+
+		buf := make([]byte, bmEmpty.LenBytes())
+		cloned := bmNil.CloneToBuf(buf)
+
+		require.Equal(t, bmEmpty.GetCardinality(), cloned.GetCardinality())
+		require.Equal(t, bmEmpty.LenBytes(), cloned.LenBytes())
+	})
+}
+
 func TestMergeToSuperset(t *testing.T) {
 	run := func(t *testing.T, bufs [][]uint16) {
 		containerThreshold := uint64(math.MaxUint16 + 1)