-
Notifications
You must be signed in to change notification settings - Fork 661
/
bloom_filter.go
77 lines (59 loc) · 1.65 KB
/
bloom_filter.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
// Copyright (C) 2019-2022, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.
package bloom
import (
"errors"
"sync"
"github.com/spaolacci/murmur3"
streakKnife "github.com/holiman/bloomfilter/v2"
)
var errMaxBytes = errors.New("too large")
type Filter interface {
// Add adds to filter, assumed thread safe
Add(...[]byte)
// Check checks filter, assumed thread safe
Check([]byte) bool
}
func New(maxN uint64, p float64, maxBytes uint64) (Filter, error) {
neededBytes := bytesSteakKnifeFilter(maxN, p)
if neededBytes > maxBytes {
return nil, errMaxBytes
}
return newSteakKnifeFilter(maxN, p)
}
type steakKnifeFilter struct {
lock sync.RWMutex
filter *streakKnife.Filter
}
func bytesSteakKnifeFilter(maxN uint64, p float64) uint64 {
m := streakKnife.OptimalM(maxN, p)
k := streakKnife.OptimalK(m, maxN)
// This is pulled from bloomFilter.newBits and bloomfilter.newRandKeys. The
// calculation is the size of the bitset which would be created from this
// filter.
mSize := (m + 63) / 64
totalSize := mSize + k
return totalSize * 8 // 8 == sizeof(uint64))
}
func newSteakKnifeFilter(maxN uint64, p float64) (Filter, error) {
m := streakKnife.OptimalM(maxN, p)
k := streakKnife.OptimalK(m, maxN)
filter, err := streakKnife.New(m, k)
return &steakKnifeFilter{filter: filter}, err
}
func (f *steakKnifeFilter) Add(bl ...[]byte) {
f.lock.Lock()
defer f.lock.Unlock()
for _, b := range bl {
h := murmur3.New64()
_, _ = h.Write(b)
f.filter.Add(h)
}
}
func (f *steakKnifeFilter) Check(b []byte) bool {
f.lock.RLock()
defer f.lock.RUnlock()
h := murmur3.New64()
_, _ = h.Write(b)
return f.filter.Contains(h)
}