From 29288d6f1c6f03aa8b5813a5963c7a44114e3572 Mon Sep 17 00:00:00 2001
From: redhat <2292650292@qq.com>
Date: Tue, 15 Apr 2025 14:27:26 +0800
Subject: [PATCH] first commit

---
 bloomfilter.go      | 97 +++++++++++++++++++++++++++++++++++++++++++++
 bloomfilter_test.go | 50 +++++++++++++++++++++++
 go.mod              |  3 ++
 local.go            | 36 +++++++++++++++++
 lua.go              | 31 +++++++++++++++
 option.go           | 45 +++++++++++++++++++++
 redis.go            | 42 ++++++++++++++++++++
 7 files changed, 304 insertions(+)
 create mode 100644 bloomfilter.go
 create mode 100644 bloomfilter_test.go
 create mode 100644 go.mod
 create mode 100644 local.go
 create mode 100644 lua.go
 create mode 100644 option.go
 create mode 100644 redis.go

diff --git a/bloomfilter.go b/bloomfilter.go
new file mode 100644
index 0000000..b9051cf
--- /dev/null
+++ b/bloomfilter.go
@@ -0,0 +1,97 @@
+package bloomfilter
+
+import (
+	"hash"
+	"hash/fnv"
+	"math"
+)
+
+type BloomFilter struct {
+	bitmap BitMapEr
+	m, k   uint64
+	hashFn []hash.Hash64
+	opt    *Options
+}
+
+type BitMapEr interface {
+	Create(uint64)
+	Set([]uint64) error
+	Get([]uint64) (bool, error)
+}
+
+func optimaMapSize(n uint64, fpRate float64) uint64 {
+	// m = -n * ln(fpRate) / (ln(2)^2)
+	size := -float64(n) * math.Log(fpRate) / math.Pow(math.Log(2), 2)
+	return uint64(math.Ceil(size))
+}
+
+func optimaHashCount(n, m uint64) uint64 {
+	// k = (m/n) * ln(2)
+	k := float64(m) / float64(n) * math.Log(2)
+	return uint64(math.Ceil(k))
+}
+
+func (bf *BloomFilter) getHashIndex(data []byte, i int) uint64 {
+	hash64 := bf.hashFn[i%len(bf.hashFn)]
+	hash64.Reset()
+	hash64.Write(data)
+	hash64.Write([]byte{byte(i)})
+
+	return hash64.Sum64() % bf.m
+}
+
+func (bf *BloomFilter) getSums(data []byte) []uint64 {
+	var res []uint64
+	for i := uint64(0); i < bf.k; i++ {
+		res = append(res, bf.getHashIndex(data, int(i)))
+	}
+
+	return res
+}
+
+func NewBloomFilter(n uint64, fpRate float64, opts ...Option) *BloomFilter {
+	options := &Options{}
+	for _, opt := range opts {
+		opt(options)
+	}
+	if err := options.repairOption(); err != nil {
+		panic(err)
+	}
+
+	m := optimaMapSize(n, fpRate)
+	k := optimaHashCount(n, m)
+
+	var bitmap BitMapEr
+	if options.mode == ModeLocal {
+		bitmap = &Local{}
+	} else {
+		bitmap = &Redis{
+			client: options.client,
+			key:    options.key,
+		}
+	}
+
+	bitmap.Create(m)
+
+	hashFn := []hash.Hash64{fnv.New64(), fnv.New64a()}
+
+	return &BloomFilter{
+		bitmap: bitmap,
+		m:      m,
+		k:      k,
+		hashFn: hashFn,
+	}
+}
+
+func (bf *BloomFilter) Add(item []byte) error {
+	return bf.bitmap.Set(bf.getSums(item))
+}
+
+func (bf *BloomFilter) MightContain(item []byte) bool {
+	get, err := bf.bitmap.Get(bf.getSums(item))
+	if err != nil {
+		return false
+	}
+
+	return get
+}
diff --git a/bloomfilter_test.go b/bloomfilter_test.go
new file mode 100644
index 0000000..035592b
--- /dev/null
+++ b/bloomfilter_test.go
@@ -0,0 +1,50 @@
+package bloomfilter
+
+import (
+	"context"
+	"errors"
+	"github.com/go-redis/redis/v8"
+	"testing"
+)
+
+func TestBloomFilter_Local(t *testing.T) {
+	filter := NewBloomFilter(10000, 0.1)
+
+	err := filter.Add([]byte("key1"))
+	if err != nil {
+		t.Error(err)
+	}
+
+	err = filter.Add([]byte("key2"))
+	if err != nil {
+		t.Error(err)
+	}
+
+	t.Error(filter.MightContain([]byte("key1")))
+	t.Error(filter.MightContain([]byte("key2")))
+	t.Error(filter.MightContain([]byte("key13")))
+}
+
+func TestBloomFilter_Redis(t *testing.T) {
+	client := redis.NewClient(&redis.Options{
+		Addr: "192.168.8.1:6379",
+	})
+
+	filter := NewBloomFilter(10000, 0.1, UseRedis(client, "redhatbloom"))
+
+	err := filter.Add([]byte("key1"))
+	if err != nil && !errors.Is(err, redis.Nil) {
+		t.Error(err)
+	}
+
+	err = filter.Add([]byte("key2"))
+	if err != nil && !errors.Is(err, redis.Nil) {
+		t.Error(err)
+	}
+
+	t.Error(filter.MightContain([]byte("key1")))
+	t.Error(filter.MightContain([]byte("key2")))
+	t.Error(filter.MightContain([]byte("key13")))
+
+	client.Del(context.Background(), "redhatbloom")
+}
diff --git a/go.mod b/go.mod
new file mode 100644
index 0000000..76ab88f
--- /dev/null
+++ b/go.mod
@@ -0,0 +1,3 @@
+module git.zhangshuocauc.cn/redhat/goBloomFilter
+
+go 1.24.2
diff --git a/local.go b/local.go
new file mode 100644
index 0000000..efbc6a6
--- /dev/null
+++ b/local.go
@@ -0,0 +1,36 @@
+package bloomfilter
+
+type Local struct {
+	bitmap []int
+}
+
+func (l *Local) getSite(sum uint64) (index uint64, offset uint64) {
+	index = sum >> 5
+	offset = sum & 31
+
+	return
+}
+
+func (l *Local) Create(u uint64) {
+	l.bitmap = make([]int, u/32+1)
+}
+
+func (l *Local) Set(uint64s []uint64) error {
+	for _, u := range uint64s {
+		index, offset := l.getSite(u)
+		l.bitmap[index] |= 1 << offset
+	}
+
+	return nil
+}
+
+func (l *Local) Get(uint64s []uint64) (bool, error) {
+	for _, u := range uint64s {
+		index, offset := l.getSite(u)
+		if l.bitmap[index]&(1<<offset) == 0 {
+			return false, nil
+		}
+	}
+
+	return true, nil
+}
diff --git a/lua.go b/lua.go
new file mode 100644
index 0000000..fec6857
--- /dev/null
+++ b/lua.go
@@ -0,0 +1,31 @@
+package bloomfilter
+
+const LuaBloomBatchGetBits = `
+	local bloomKey = KEYS[1]
+	local bitsCnt = ARGV[1]
+	for i=1,bitsCnt,1 do
+		local offset = ARGV[i+1]
+		local result = redis.call('getbit',bloomKey,offset)
+		if (not result) then
+			error('FALL')
+			return 0
+		end
+		if (result == 0) then
+			return 0
+		end
+	end
+	return 1
+`
+
+const LuaBloomBatchSetBits = `
+	local bloomKey = KEYS[1]
+	local bitsCnt = ARGV[1]
+	for i=1,bitsCnt,1 do
+		local offset = ARGV[i+1]
+		local result = redis.call('setbit',bloomKey,offset,1)
+		if (not result) then
+			error('FALL')
+			return 0
+		end
+	end
+`
diff --git a/option.go b/option.go
new file mode 100644
index 0000000..2361c0e
--- /dev/null
+++ b/option.go
@@ -0,0 +1,45 @@
+package bloomfilter
+
+import (
+	"errors"
+	"github.com/go-redis/redis/v8"
+)
+
+type Mode int
+
+const (
+	ModeLocal Mode = iota
+	ModeRedis
+)
+
+type Options struct {
+	mode   Mode
+	client *redis.Client
+	key    string
+}
+
+type Option func(*Options)
+
+func UseLocal() Option {
+	return func(o *Options) {
+		o.mode = ModeLocal
+	}
+}
+
+func UseRedis(client *redis.Client, key string) Option {
+	return func(o *Options) {
+		o.mode = ModeRedis
+		o.client = client
+		o.key = key
+	}
+}
+
+func (o *Options) repairOption() error {
+	if o.mode == ModeRedis {
+		if o.client == nil || o.key == "" {
+			return errors.New("no addr provided")
+		}
+	}
+
+	return nil
+}
diff --git a/redis.go b/redis.go
new file mode 100644
index 0000000..6e8d5b6
--- /dev/null
+++ b/redis.go
@@ -0,0 +1,42 @@
+package bloomfilter
+
+import (
+	"context"
+	"github.com/go-redis/redis/v8"
+)
+
+type Redis struct {
+	client *redis.Client
+	key    string
+	ctx    context.Context
+}
+
+func (r *Redis) Create(u uint64) {
+	r.ctx = context.Background()
+}
+
+func (r *Redis) Set(uint64s []uint64) error {
+	args := make([]interface{}, 0, len(uint64s)+1)
+	args = append(args, len(uint64s))
+	for _, u := range uint64s {
+		args = append(args, u)
+	}
+
+	return r.client.Eval(r.ctx, LuaBloomBatchSetBits, []string{r.key}, args...).Err()
+}
+
+func (r *Redis) Get(uint64s []uint64) (bool, error) {
+	args := make([]interface{}, 0, len(uint64s)+1)
+	args = append(args, len(uint64s))
+	for _, u := range uint64s {
+		args = append(args, u)
+	}
+
+	eval := r.client.Eval(r.ctx, LuaBloomBatchGetBits, []string{r.key}, args...)
+
+	if res, ok := eval.Val().(int64); ok && res == 1 {
+		return true, nil
+	}
+
+	return false, eval.Err()
+}