bloom.go

/*
Package bloom provides data structures and methods for creating Bloom filters.

A Bloom filter is a representation of a set of _n_ items, where the main
requirement is to make membership queries; _i.e._, whether an item is a
member of a set.

A Bloom filter has two parameters: _m_, a maximum size (typically a reasonably large
multiple of the cardinality of the set to represent) and _k_, the number of hashing
functions on elements of the set. (The actual hashing functions are important, too,
but this is not a parameter for this implementation). A Bloom filter is backed by
a BitSet; a key is represented in the filter by setting the bits at each value of the
hashing functions (modulo _m_). Set membership is done by _testing_ whether the
bits at each value of the hashing functions (again, modulo _m_) are set. If so,
the item is in the set. If the item is actually in the set, a Bloom filter will
never fail (the true positive rate is 1.0); but it is susceptible to false
positives. The art is to choose _k_ and _m_ correctly.

In this implementation, the hashing functions used is murmurhash,
a non-cryptographic hashing function.

This implementation accepts keys for setting as testing as []byte. Thus, to
add a string item, "Love":

    uint n = 1000
    filter := bloom.New(20*n, 5) // load of 20, 5 keys
    filter.Add([]byte("Love"))

Similarly, to test if "Love" is in bloom:

    if filter.Test([]byte("Love"))

For numeric data, I recommend that you look into the binary/encoding library. But,
for example, to add a uint32 to the filter:

    i := uint32(100)
    n1 := make([]byte,4)
    binary.BigEndian.PutUint32(n1,i)
    f.Add(n1)

Finally, there is a method to estimate the false positive rate of a particular
Bloom filter for a set of size _n_:

    if filter.EstimateFalsePositiveRate(1000) > 0.001

Given the particular hashing scheme, it's best to be empirical about this. Note
that estimating the FP rate will clear the Bloom filter.
*/
package bloom

import (
	"math"
)

// A BitSetProvider is a set to store bit
type BitSetProvider interface {
	// Set bit to bitset
	Set(offset uint) error
	// Test offset is 1
	Test(offset uint) (bool, error)
	// TestBatch test offset array all 1
	TestBatch(offset []uint) (bool, error)
	// TestBatchOffset test offset array all 1
	TestBatchOffset(offsetArray [][]uint) ([]bool, error)
	// SetBatch set bit array to bitset
	SetBatch(offset []uint) error
	// New with m bits init
	New(m uint)
}

var (
	TestBatch = true
	AddBatch  = true
)

// A BloomFilter is a representation of a set of _n_ items, where the main
// requirement is to make membership queries; _i.e._, whether an item is a
// member of a set.
type BloomFilter struct {
	m uint
	k uint
	b BitSetProvider
}

func max(x, y uint) uint {
	if x > y {
		return x
	}
	return y
}

// New creates a new Bloom filter with _m_ bits and _k_ hashing functions
// We force _m_ and _k_ to be at least one to avoid panics.
func New(m uint, k uint, b BitSetProvider) *BloomFilter {
	return &BloomFilter{max(1, m), max(1, k), b}
}

// baseHashes returns the four hash values of data that are used to create k
// hashes
func baseHashes(data []byte) [4]uint64 {
	var d digest128 // murmur hashing
	hash1, hash2, hash3, hash4 := d.sum256(data)
	return [4]uint64{
		hash1, hash2, hash3, hash4,
	}
}

// location returns the ith hashed location using the four base hash values
func location(h [4]uint64, i uint) uint64 {
	ii := uint64(i)
	return h[ii%2] + ii*h[2+(((ii+(ii%2))%4)/2)]
}

// location returns the ith hashed location using the four base hash values
func (f *BloomFilter) location(h [4]uint64, i uint) uint {
	return uint(location(h, i) % uint64(f.m))
}

// EstimateParameters estimates requirements for m and k.
// Based on https://bitbucket.org/ww/bloom/src/829aa19d01d9/bloom.go
// used with permission.
func EstimateParameters(n uint, p float64) (m uint, k uint) {
	m = uint(math.Ceil(-1 * float64(n) * math.Log(p) / math.Pow(math.Log(2), 2)))
	k = uint(math.Ceil(math.Log(2) * float64(m) / float64(n)))
	return
}

// NewWithEstimates creates a new Bloom filter for about n items with fp
// false positive rate
func NewWithEstimates(n uint, fp float64, provider BitSetProvider) *BloomFilter {
	m, k := EstimateParameters(n, fp)
	return New(m, k, provider)
}

// Cap returns the capacity, _m_, of a Bloom filter
func (f *BloomFilter) Cap() uint {
	return f.m
}

// K returns the number of hash functions used in the BloomFilter
func (f *BloomFilter) K() uint {
	return f.k
}

// BitSet returns the underlying bitset for this filter.
func (f *BloomFilter) BitSet() *BitSetProvider {
	return &f.b
}

// Add data to the Bloom Filter. Returns the filter (allows chaining)
func (f *BloomFilter) Add(data []byte) error {
	h := baseHashes(data)
	var offset []uint
	for i := uint(0); i < f.k; i++ {
		if AddBatch {
			offset = append(offset, f.location(h, i))
			continue
		}
		err := f.b.Set(f.location(h, i))
		if err != nil {
			return err
		}
	}

	if AddBatch {
		return f.b.SetBatch(offset)
	}

	return nil
}

// AddString to the Bloom Filter. Returns the filter (allows chaining)
func (f *BloomFilter) AddString(data string) error {
	return f.Add([]byte(data))
}

func (f *BloomFilter) AddStrings(data []*string) error {
	var offset []uint
	for i := range data {
		dataBytes := []byte(*data[i])
		h := baseHashes(dataBytes)
		for j := uint(0); j < f.k; j++ {
			offset = append(offset, f.location(h, j))
		}
	}

	return f.b.SetBatch(offset)
}

// Test returns true if the data is in the BloomFilter, false otherwise.
// If true, the result might be a false positive. If false, the data
// is definitely not in the set.
func (f *BloomFilter) Test(data []byte) (bool, error) {
	h := baseHashes(data)
	var offset []uint
	for i := uint(0); i < f.k; i++ {
		if TestBatch {
			offset = append(offset, f.location(h, i))
			continue
		}
		result, err := f.b.Test(f.location(h, i))
		if err != nil {
			return false, err
		}

		if !result {
			return false, nil
		}
	}

	if TestBatch {
		return f.b.TestBatch(offset)
	}

	return true, nil
}

func (f *BloomFilter) getOffsetByData(data []byte) []uint {
	h := baseHashes(data)
	var offset []uint
	for i := uint(0); i < f.k; i++ {
		offset = append(offset, f.location(h, i))
	}
	return offset
}

// TestString returns true if the string is in the BloomFilter, false otherwise.
// If true, the result might be a false positive. If false, the data
// is definitely not in the set.
func (f *BloomFilter) TestString(data string) (bool, error) {
	return f.Test([]byte(data))
}

func (f *BloomFilter) TestStrings(data []string) ([]bool, error) {
	var result []bool
	var offsets [][]uint
	for i := range data {
		offsets = append(offsets, f.getOffsetByData([]byte(data[i])))
	}
	boolResults, err := f.b.TestBatchOffset(offsets)
	if err != nil {
		return result, err
	}

	tempResult := true
	count := 1
	for i := range boolResults {
		if uint(count)%f.k == 0 {
			result = append(result, tempResult)
			tempResult = true
			count++
			continue
		}

		if tempResult && !boolResults[i] {
			tempResult = false
		}
		count++
	}

	return result, nil
}

// TestLocations returns true if all locations are set in the BloomFilter, false
// otherwise.
func (f *BloomFilter) TestLocations(locs []uint64) (bool, error) {
	for i := 0; i < len(locs); i++ {
		result, err := f.b.Test(uint(locs[i] % uint64(f.m)))
		if err != nil {
			return false, err
		}

		if !result {
			return false, nil
		}
	}

	return true, nil
}

// TestAndAdd is the equivalent to calling Test(data) then Add(data).
// Returns the result of Test.
func (f *BloomFilter) TestAndAdd(data []byte) (bool, error) {
	present := true
	h := baseHashes(data)
	for i := uint(0); i < f.k; i++ {
		l := f.location(h, i)
		result, err := f.b.Test(l)
		if err != nil {
			return false, err
		}
		if !result {
			present = false
		}

		err = f.b.Set(l)
		if err != nil {
			return false, err
		}
	}

	return present, nil
}

// TestAndAddString is the equivalent to calling Test(string) then Add(string).
// Returns the result of Test.
func (f *BloomFilter) TestAndAddString(data string) (bool, error) {
	return f.TestAndAdd([]byte(data))
}

// TestOrAdd is the equivalent to calling Test(data) then if not present Add(data).
// Returns the result of Test.
func (f *BloomFilter) TestOrAdd(data []byte) (bool, error) {
	present := true
	h := baseHashes(data)
	for i := uint(0); i < f.k; i++ {
		l := f.location(h, i)
		result, err := f.b.Test(l)
		if err != nil {
			return false, err
		}

		if !result {
			present = false
			err = f.b.Set(l)
			if err != nil {
				return false, err
			}
		}
	}

	return present, nil
}

// TestOrAddString is the equivalent to calling Test(string) then if not present Add(string).
// Returns the result of Test.
func (f *BloomFilter) TestOrAddString(data string) (bool, error) {
	return f.TestOrAdd([]byte(data))
}

// Approximating the number of items
// https://en.wikipedia.org/wiki/Bloom_filter#Approximating_the_number_of_items_in_a_Bloom_filter
func (f *BloomFilter) ApproximatedSize() uint32 {
	x := float64(f.m)
	m := float64(f.Cap())
	k := float64(f.K())
	size := -1 * m / k * math.Log(1-x/m) / math.Log(math.E)
	return uint32(math.Floor(size + 0.5)) // round
}

// Locations returns a list of hash locations representing a data item.
func Locations(data []byte, k uint) []uint64 {
	locs := make([]uint64, k)

	// calculate locations
	h := baseHashes(data)
	for i := uint(0); i < k; i++ {
		locs[i] = location(h, i)
	}

	return locs
}