optimize bitmap

Author: markus621

structs/bitmap/bitmap.go

@@ -11,59 +11,66 @@ import (
 
 const (
 	blockSize = 8
+	MaxIndex  = uint64(1 << 34)
 )
 
 type Bitmap struct {
-	data    []byte
-	size    uint64
+	bits []byte
+
+	blocks  uint64
 	max     uint64
-	mux     sync.RWMutex
 	lockoff bool
+
+	mux sync.RWMutex
 }
 
 type Option func(*Bitmap)
 
-func DisableLock() Option {
+func OptDisableLock() Option {
 	return func(o *Bitmap) {
 		o.lockoff = true
 	}
 }
 
-func New(maxIndex uint64, opts ...Option) *Bitmap {
-	size := maxIndex/blockSize + maxIndex%blockSize
-
-	bm := &Bitmap{
-		max:  size * blockSize,
-		size: size,
-		data: make([]byte, size),
+func OptMaxIndex(index uint64) Option {
+	return func(o *Bitmap) {
+		o.max = index
 	}
+}
+
+func New(opts ...Option) *Bitmap {
+	bm := &Bitmap{max: 1}
 
 	for _, opt := range opts {
 		opt(bm)
 	}
 
+	bm.resize(bm.max)
+
 	return bm
 }
 
-func (b *Bitmap) CopyTo(dst *Bitmap) {
-	if !b.lockoff {
-		b.mux.Lock()
-		defer b.mux.Unlock()
-	}
-	if !dst.lockoff {
-		dst.mux.Lock()
-		defer dst.mux.Unlock()
+func (b *Bitmap) resize(index uint64) {
+	size := index / blockSize
+	if index-(size%blockSize) > 0 {
+		size++
 	}
 
-	dst.data = make([]byte, len(b.data))
-	copy(dst.data, b.data)
-	dst.size = b.size
-	dst.max = b.max
-	dst.lockoff = b.lockoff
+	b.bits = append(b.bits, make([]byte, size-b.blocks)...)
+	b.blocks = uint64(len(b.bits))
+	b.max = b.blocks*blockSize - 1
+}
+
+func (b *Bitmap) getBlock(index uint64) uint64 {
+	return index / blockSize
+}
+
+func (b *Bitmap) getBit(index uint64) byte {
+	return 1 << (index - b.getBlock(index)*blockSize)
 }
 
 func (b *Bitmap) Set(index uint64) {
-	if index > b.max {
+	if index > MaxIndex {
 		return
 	}
 
@@ -72,11 +79,15 @@ func (b *Bitmap) Set(index uint64) {
 		defer b.mux.Unlock()
 	}
 
-	b.data[index%b.size] |= 1 << (index % blockSize)
+	if index > b.max {
+		b.resize(index)
+	}
+
+	b.bits[b.getBlock(index)] |= b.getBit(index)
 }
 
 func (b *Bitmap) Del(index uint64) {
-	if index > b.max {
+	if index > b.max || index > MaxIndex {
 		return
 	}
 
@@ -85,11 +96,11 @@ func (b *Bitmap) Del(index uint64) {
 		defer b.mux.Unlock()
 	}
 
-	b.data[index%b.size] &^= 1 << (index % blockSize)
+	b.bits[b.getBlock(index)] &^= b.getBit(index)
 }
 
 func (b *Bitmap) Has(index uint64) bool {
-	if index > b.max {
+	if index > b.max || index > MaxIndex {
 		return false
 	}
 
@@ -98,29 +109,49 @@ func (b *Bitmap) Has(index uint64) bool {
 		defer b.mux.RUnlock()
 	}
 
-	return (b.data[index%b.size] & (1 << (index % blockSize))) > 0
+	return (b.bits[b.getBlock(index)] & b.getBit(index)) > 0
 }
 
-func (b *Bitmap) Dump() []byte {
+func (b *Bitmap) MarshalBinary() ([]byte, error) {
 	if !b.lockoff {
 		b.mux.RLock()
 		defer b.mux.RUnlock()
 	}
 
-	out := make([]byte, b.size)
-	copy(out, b.data)
-	return out
+	out := make([]byte, b.blocks)
+	copy(out, b.bits)
+
+	return out, nil
 }
 
-func (b *Bitmap) Restore(in []byte) {
+func (b *Bitmap) UnmarshalBinary(in []byte) error {
 	if !b.lockoff {
 		b.mux.Lock()
 		defer b.mux.Unlock()
 	}
 
-	b.data = make([]byte, len(in))
-	copy(b.data, in)
+	b.bits = make([]byte, len(in))
+	copy(b.bits, in)
 
-	b.size = uint64(len(in))
-	b.max = b.size * blockSize
+	b.blocks = uint64(len(in))
+	b.max = b.blocks * blockSize
+
+	return nil
+}
+
+func (b *Bitmap) CopyTo(dst *Bitmap) {
+	if !b.lockoff {
+		b.mux.Lock()
+		defer b.mux.Unlock()
+	}
+	if !dst.lockoff {
+		dst.mux.Lock()
+		defer dst.mux.Unlock()
+	}
+
+	dst.bits = make([]byte, len(b.bits))
+	copy(dst.bits, b.bits)
+	dst.blocks = b.blocks
+	dst.max = b.max
+	dst.lockoff = b.lockoff
 }

structs/bitmap/bitmap_test.go

@@ -6,29 +6,41 @@
 package bitmap
 
 import (
-	"math"
+	"fmt"
 	"testing"
 
 	"go.osspkg.com/casecheck"
 )
 
 func TestUnit_Bitmap_CopyTO(t *testing.T) {
-	src := New(65)
+	src := New()
 	src.Set(1)
 	src.Set(5)
 	src.Set(60)
 
-	dst := New(100)
+	dst := New()
 	src.CopyTo(dst)
 
-	casecheck.Equal(t, src.size, dst.size)
-	casecheck.Equal(t, src.data, dst.data)
+	casecheck.Equal(t, src.blocks, dst.blocks)
+	casecheck.Equal(t, src.bits, dst.bits)
 	casecheck.Equal(t, src.lockoff, dst.lockoff)
 	casecheck.Equal(t, src.max, dst.max)
 }
 
-func TestUnit_Bitmap_calcBlockIndex(t *testing.T) {
-	bm := New(65)
+func TestUnit_Bitmap_Resize(t *testing.T) {
+	t.Skip("Only for debug")
+
+	src := New()
+
+	for i := 0; i <= 20; i++ {
+		src.Set(uint64(i))
+		b, _ := src.MarshalBinary()
+		fmt.Printf("(%d) %b `%s`\n", i, b, string(b))
+	}
+}
+
+func TestUnit_Bitmap_Marshaling(t *testing.T) {
+	bm := New()
 
 	for i := 0; i <= 65; i++ {
 		bm.Set(uint64(i))
@@ -37,16 +49,16 @@ func TestUnit_Bitmap_calcBlockIndex(t *testing.T) {
 		casecheck.False(t, bm.Has(uint64(i+1)), "(2) for index: %d", i+1)
 	}
 
-	backup := bm.Dump()
+	backup, _ := bm.MarshalBinary()
 
-	bm.Restore(make([]byte, len(backup)))
+	bm.UnmarshalBinary(make([]byte, len(backup)))
 
 	for i := 0; i <= 65; i++ {
 		casecheck.False(t, bm.Has(uint64(i)), "(3) for index: %d", i)
 		casecheck.False(t, bm.Has(uint64(i+1)), "(4) for index: %d", i+1)
 	}
 
-	bm.Restore(backup)
+	bm.UnmarshalBinary(backup)
 
 	for i := 65; i >= 0; i-- {
 		casecheck.True(t, bm.Has(uint64(i)), "(1) for index: %d", i)
@@ -64,11 +76,12 @@ goarch: amd64
 pkg: go.osspkg.com/algorithms/structs/bitmap
 cpu: 12th Gen Intel(R) Core(TM) i9-12900KF
 Benchmark_Bitmap
-Benchmark_Bitmap-4   	 9942369	       162.7 ns/op	       0 B/op	       0 allocs/op
+Benchmark_Bitmap-4   	 9515085	       182.2 ns/op	     225 B/op	       0 allocs/op
 */
 func Benchmark_Bitmap(b *testing.B) {
-	index := uint64(math.MaxInt16)
-	bm := New(index)
+	index := uint64(1 << 34)
+
+	bm := New(OptMaxIndex(1024))
 
 	b.ReportAllocs()
 	b.ResetTimer()

structs/bloom/bloom.go

@@ -85,7 +85,7 @@ func New(opts ...Option) (*Bloom, error) {
 	m, k := calcOptimalParams(b.optSize, b.optRate)
 
 	b.size = m
-	b.bits = bitmap.New(m, bitmap.DisableLock())
+	b.bits = bitmap.New(bitmap.OptMaxIndex(m), bitmap.OptDisableLock())
 	b.salts = make([][saltSize]byte, k)
 
 	for i := 0; i < int(k); i++ {
@@ -138,7 +138,12 @@ func (b *Bloom) Dump(w io.Writer) error {
 		}
 	}
 
-	if _, err := w.Write(b.bits.Dump()); err != nil {
+	bb, err := b.bits.MarshalBinary()
+	if err != nil {
+		return fmt.Errorf("marshal bitset: %w", err)
+	}
+
+	if _, err = w.Write(bb); err != nil {
 		return fmt.Errorf("write bitmap: %w", err)
 	}
 
@@ -194,9 +199,7 @@ func (b *Bloom) Restore(r io.Reader) error {
 		return fmt.Errorf("read bitmap: %w", err)
 	}
 
-	b.bits.Restore(bm)
-
-	return nil
+	return b.bits.UnmarshalBinary(bm)
 }
 
 func (b *Bloom) Add(arg any) {