improve performance

bucket.go

@@ -21,41 +21,46 @@ const (
 // insert a fingerprint into a bucket. Returns true if there was enough space and insertion succeeded.
 // Note it allows inserting the same fingerprint multiple times.
 func (b *bucket) insert(fp fingerprint) bool {
-	for i, tfp := range b {
-		if tfp == nullFp {
-			b[i] = fp
-			return true
-		}
+	if i := b.index(nullFp); i != 4 {
+		b[i] = fp
+		return true
 	}
 	return false
 }
 
 // delete a fingerprint from a bucket.
 // Returns true if the fingerprint was present and successfully removed.
 func (b *bucket) delete(fp fingerprint) bool {
-	for i, tfp := range b {
-		if tfp == fp {
-			b[i] = nullFp
-			return true
-		}
+	if i := b.index(fp); i != 4 {
+		b[i] = nullFp
+		return true
 	}
 	return false
 }
 
 func (b *bucket) contains(needle fingerprint) bool {
-	for _, fp := range b {
-		if fp == needle {
-			return true
-		}
+	return b.index(needle) != 4
+}
+
+func (b *bucket) index(needle fingerprint) uint8 {
+	if b[0] == needle {
+		return 0
 	}
-	return false
+	if b[1] == needle {
+		return 1
+	}
+	if b[2] == needle {
+		return 2
+	}
+	if b[3] == needle {
+		return 3
+	}
+	return 4
 }
 
 // reset deletes all fingerprints in the bucket.
 func (b *bucket) reset() {
-	for i := range b {
-		b[i] = nullFp
-	}
+	*b = [bucketSize]fingerprint{nullFp, nullFp, nullFp, nullFp}
 }
 
 func (b *bucket) String() string {

cuckoofilter.go

@@ -1,10 +1,10 @@
 package cuckoo
 
 import (
-	"bytes"
 	"encoding/binary"
 	"fmt"
-	"math/rand"
+
+	"github.com/zeebo/wyhash"
 )
 
 // maxCuckooKickouts is the maximum number of times reinsert
@@ -18,6 +18,7 @@ type Filter struct {
 	// Bit mask set to len(buckets) - 1. As len(buckets) is always a power of 2,
 	// applying this mask mimics the operation x % len(buckets).
 	bucketIndexMask uint
+	rng             wyhash.RNG
 }
 
 // NewFilter returns a new cuckoofilter suitable for the given number of elements.
@@ -72,7 +73,7 @@ func (cf *Filter) Insert(data []byte) bool {
 	if cf.insert(fp, i2) {
 		return true
 	}
-	return cf.reinsert(fp, randi(i1, i2))
+	return cf.reinsert(fp, randi(&cf.rng, i1, i2))
 }
 
 func (cf *Filter) insert(fp fingerprint, i uint) bool {
@@ -85,7 +86,7 @@ func (cf *Filter) insert(fp fingerprint, i uint) bool {
 
 func (cf *Filter) reinsert(fp fingerprint, i uint) bool {
 	for k := 0; k < maxCuckooKickouts; k++ {
-		j := rand.Intn(bucketSize)
+		j := cf.rng.Intn(bucketSize)
 		// Swap fingerprint with bucket entry.
 		cf.buckets[i][j], fp = fp, cf.buckets[i][j]
 
@@ -127,15 +128,13 @@ const bytesPerBucket = bucketSize * fingerprintSizeBits / 8
 
 // Encode returns a byte slice representing a Cuckoofilter.
 func (cf *Filter) Encode() []byte {
-	res := new(bytes.Buffer)
-	res.Grow(len(cf.buckets) * bytesPerBucket)
-
+	buf := make([]byte, 0, len(cf.buckets)*bytesPerBucket)
 	for _, b := range cf.buckets {
 		for _, fp := range b {
-			binary.Write(res, binary.LittleEndian, fp)
+			buf = binary.LittleEndian.AppendUint16(buf, uint16(fp))
 		}
 	}
-	return res.Bytes()
+	return buf
 }
 
 // Decode returns a Cuckoofilter from a byte slice created using Encode.
@@ -150,13 +149,13 @@ func Decode(data []byte) (*Filter, error) {
 	if getNextPow2(uint64(numBuckets)) != uint(numBuckets) {
 		return nil, fmt.Errorf("numBuckets must to be a power of 2, got %d", numBuckets)
 	}
-	var count uint
-	buckets := make([]bucket, numBuckets)
-	reader := bytes.NewReader(data)
 
+	var count, pos uint
+	buckets := make([]bucket, numBuckets)
 	for i, b := range buckets {
 		for j := range b {
-			binary.Read(reader, binary.LittleEndian, &buckets[i][j])
+			buckets[i][j] = fingerprint(binary.LittleEndian.Uint16(data[pos : pos+2]))
+			pos += 2
 			if buckets[i][j] != nullFp {
 				count++
 			}

go.mod

@@ -3,6 +3,7 @@ module github.com/panmari/cuckoofilter
 go 1.15
 
 require (
-	github.com/dgryski/go-metro v0.0.0-20200812162917-85c65e2d0165
 	github.com/google/go-cmp v0.5.9
+	github.com/zeebo/wyhash v0.0.1
+	github.com/zeebo/xxh3 v1.0.2
 )

go.sum

@@ -1,4 +1,10 @@
-github.com/dgryski/go-metro v0.0.0-20200812162917-85c65e2d0165 h1:BS21ZUJ/B5X2UVUbczfmdWH7GapPWAhxcMsDnjJTU1E=
-github.com/dgryski/go-metro v0.0.0-20200812162917-85c65e2d0165/go.mod h1:c9O8+fpSOX1DM8cPNSkX/qsBWdkD4yd2dpciOWQjpBw=
 github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38=
 github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
+github.com/klauspost/cpuid/v2 v2.0.9 h1:lgaqFMSdTdQYdZ04uHyN2d/eKdOMyi2YLSvlQIBFYa4=
+github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg=
+github.com/zeebo/assert v1.3.0 h1:g7C04CbJuIDKNPFHmsk4hwZDO5O+kntRxzaUoNXj+IQ=
+github.com/zeebo/assert v1.3.0/go.mod h1:Pq9JiuJQpG8JLJdtkwrJESF0Foym2/D9XMU5ciN/wJ0=
+github.com/zeebo/wyhash v0.0.1 h1:VEByEMek3iHhV65CgG3SRAWVtg/6TcmbEKj5jPOKDrc=
+github.com/zeebo/wyhash v0.0.1/go.mod h1:Ti+OwfNtM5AZiYAL0kOPIfliqDP5c0VtOnnMAqzuuZk=
+github.com/zeebo/xxh3 v1.0.2 h1:xZmwmqxHZA8AI603jOQ0tMqmBr9lPeFwGg6d+xy9DC0=
+github.com/zeebo/xxh3 v1.0.2/go.mod h1:5NWz9Sef7zIDm2JHfFlcQvNekmcEl9ekUZQQKCYaDcA=

util.go

@@ -2,24 +2,32 @@ package cuckoo
 
 import (
 	"encoding/binary"
-	"math/rand"
+	"math/bits"
 
-	metro "github.com/dgryski/go-metro"
+	"github.com/zeebo/wyhash"
+	"github.com/zeebo/xxh3"
 )
 
+var altHash [maxFingerprint + 1]uint
+
+func init() {
+	b := make([]byte, 2)
+	for i := 0; i < maxFingerprint+1; i++ {
+		binary.LittleEndian.PutUint16(b, uint16(i))
+		altHash[i] = (uint(xxh3.Hash(b)))
+	}
+}
+
 // randi returns either i1 or i2 randomly.
-func randi(i1, i2 uint) uint {
-	if rand.Int31()%2 == 0 {
+func randi(rng *wyhash.RNG, i1, i2 uint) uint {
+	if rng.Uint64()&1 == 0 {
 		return i1
 	}
 	return i2
 }
 
 func getAltIndex(fp fingerprint, i uint, bucketIndexMask uint) uint {
-	b := make([]byte, 2)
-	binary.LittleEndian.PutUint16(b, uint16(fp))
-	hash := uint(metro.Hash64(b, 1337))
-	return (i ^ hash) & bucketIndexMask
+	return (i ^ altHash[fp]) & bucketIndexMask
 }
 
 func getFingerprint(hash uint64) fingerprint {
@@ -32,21 +40,13 @@ func getFingerprint(hash uint64) fingerprint {
 
 // getIndexAndFingerprint returns the primary bucket index and fingerprint to be used
 func getIndexAndFingerprint(data []byte, bucketIndexMask uint) (uint, fingerprint) {
-	hash := metro.Hash64(data, 1337)
+	hash := xxh3.Hash(data)
 	f := getFingerprint(hash)
 	// Use least significant bits for deriving index.
 	i1 := uint(hash) & bucketIndexMask
 	return i1, f
 }
 
 func getNextPow2(n uint64) uint {
-	n--
-	n |= n >> 1
-	n |= n >> 2
-	n |= n >> 4
-	n |= n >> 8
-	n |= n >> 16
-	n |= n >> 32
-	n++
-	return uint(n)
+	return uint(1 << bits.Len64(n-1))
 }