diff --git a/structs/bitmap/bitmap.go b/structs/bitmap/bitmap.go index 640bc04..9d64a7f 100644 --- a/structs/bitmap/bitmap.go +++ b/structs/bitmap/bitmap.go @@ -45,6 +45,23 @@ func New(maxIndex uint64, opts ...Option) *Bitmap { 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() + } + + dst.data = make([]byte, len(b.data)) + copy(dst.data, b.data) + dst.size = b.size + dst.max = b.max + dst.lockoff = b.lockoff +} + func (b *Bitmap) Set(index uint64) { if index > b.max { return diff --git a/structs/bitmap/bitmap_test.go b/structs/bitmap/bitmap_test.go index b369cf0..d856dce 100644 --- a/structs/bitmap/bitmap_test.go +++ b/structs/bitmap/bitmap_test.go @@ -12,6 +12,21 @@ import ( "go.osspkg.com/casecheck" ) +func TestUnit_Bitmap_CopyTO(t *testing.T) { + src := New(65) + src.Set(1) + src.Set(5) + src.Set(60) + + dst := New(100) + src.CopyTo(dst) + + casecheck.Equal(t, src.size, dst.size) + casecheck.Equal(t, src.data, dst.data) + casecheck.Equal(t, src.lockoff, dst.lockoff) + casecheck.Equal(t, src.max, dst.max) +} + func TestUnit_Bitmap_calcBlockIndex(t *testing.T) { bm := New(65) diff --git a/structs/bloom/bloom.go b/structs/bloom/bloom.go index 9764f2d..69037f1 100644 --- a/structs/bloom/bloom.go +++ b/structs/bloom/bloom.go @@ -99,7 +99,27 @@ func New(opts ...Option) (*Bloom, error) { return b, nil } +func (b *Bloom) CopyTo(dst *Bloom) { + b.mux.Lock() + defer b.mux.Unlock() + + dst.mux.Lock() + defer dst.mux.Unlock() + + b.bits.CopyTo(dst.bits) + dst.size = b.size + + dst.salts = make([][saltSize]byte, len(b.salts)) + copy(dst.salts, b.salts) + + dst.optSize = b.optSize + dst.optRate = b.optRate +} + func (b *Bloom) Dump(w io.Writer) error { + b.mux.RLock() + defer b.mux.RUnlock() + if _, err := w.Write([]byte("OSSPkg:bloom\n")); err != nil { return fmt.Errorf("write header: %w", err) } @@ -118,9 +138,6 @@ func (b *Bloom) Dump(w io.Writer) error { } } - b.mux.RLock() - defer b.mux.RUnlock() - if _, err := w.Write(b.bits.Dump()); err != nil { return fmt.Errorf("write bitmap: %w", err) } @@ -129,6 +146,9 @@ func (b *Bloom) Dump(w io.Writer) error { } func (b *Bloom) Restore(r io.Reader) error { + b.mux.Lock() + defer b.mux.Unlock() + reader := bufio.NewReader(r) head, err := reader.ReadBytes('\n') @@ -174,9 +194,6 @@ func (b *Bloom) Restore(r io.Reader) error { return fmt.Errorf("read bitmap: %w", err) } - b.mux.Lock() - defer b.mux.Unlock() - b.bits.Restore(bm) return nil diff --git a/structs/bloom/bloom_test.go b/structs/bloom/bloom_test.go index dc208b0..07d9e67 100644 --- a/structs/bloom/bloom_test.go +++ b/structs/bloom/bloom_test.go @@ -10,6 +10,7 @@ import ( "crypto/md5" "crypto/sha1" "crypto/sha256" + "fmt" "hash" "hash/fnv" "reflect" @@ -20,7 +21,7 @@ import ( ) func TestUnit_Bloom(t *testing.T) { - bf, err := New(Quantity(1000, 0.00001)) + bf, err := New(Quantity(4, 0.01)) casecheck.NoError(t, err) bf.Add("hello") @@ -34,8 +35,16 @@ func TestUnit_Bloom(t *testing.T) { buf := bytes.NewBuffer(nil) casecheck.NoError(t, bf.Dump(buf)) - //fmt.Println(string(buf.Bytes())) + b1 := buf.Bytes() + + fmt.Println(string(b1)) + casecheck.NoError(t, bf.Restore(buf)) + buf = bytes.NewBuffer(nil) + casecheck.NoError(t, bf.Dump(buf)) + b2 := buf.Bytes() + + casecheck.Equal(t, b1, b2) casecheck.False(t, bf.Contain("users")) casecheck.True(t, bf.Contain("user"))