tapir/primitives/bloom.go

package primitives

import (
	"crypto/sha256"
	"sync"
)

// BloomFilter implements a bloom filter
type BloomFilter struct {
	B    []bool
	lock sync.Mutex
}

// Init constructs a bloom filter of size m
func (bf *BloomFilter) Init(m int16) {
	bf.B = make([]bool, m)
}

// Hash transforms a message to a set of bit flips
// Supports up to m == 65535
func (bf *BloomFilter) Hash(msg []byte) []int {
	hash := sha256.Sum256(msg)

	pos1a := (int(hash[0]) + int(hash[1]) + int(hash[2]) + int(hash[3])) % 0xFF
	pos1b := (int(hash[4]) + int(hash[5]) + int(hash[6]) + int(hash[7])) % 0xFF
	pos1 := ((pos1a << 8) + pos1b) & (0xFFFF % len(bf.B))

	pos2a := (int(hash[8]) + int(hash[9]) + int(hash[10]) + int(hash[11])) % 0xFF
	pos2b := (int(hash[12]) + int(hash[13]) + int(hash[14]) + int(hash[15])) % 0xFF
	pos2 := ((pos2a << 8) + pos2b) & (0xFFFF % len(bf.B))

	pos3a := (int(hash[16]) + int(hash[17]) + int(hash[18]) + int(hash[19])) % 0xFF
	pos3b := (int(hash[20]) + int(hash[21]) + int(hash[22]) + int(hash[23])) % 0xFF
	pos3 := ((pos3a << 8) + pos3b) & (0xFFFF % len(bf.B))

	pos4a := (int(hash[24]) + int(hash[25]) + int(hash[26]) + int(hash[27])) % 0xFF
	pos4b := (int(hash[28]) + int(hash[29]) + int(hash[30]) + int(hash[31])) % 0xFF
	pos4 := ((pos4a << 8) + pos4b) & (0xFFFF % len(bf.B))

	return []int{pos1, pos2, pos3, pos4}
}

// Insert updates the BloomFilter (suitable for concurrent use)
func (bf *BloomFilter) Insert(msg []byte) {
	pos := bf.Hash(msg)
	bf.lock.Lock()
	defer bf.lock.Unlock()
	bf.B[pos[0]] = true
	bf.B[pos[1]] = true
	bf.B[pos[2]] = true
	bf.B[pos[3]] = true
}

// Check returns true if the messages might be in the BloomFilter
// (No false positives, possible false negatives due to the probabilistic nature of the filter)
func (bf *BloomFilter) Check(msg []byte) bool {
	pos := bf.Hash(msg)
	if bf.B[pos[0]] && bf.B[pos[1]] && bf.B[pos[2]] && bf.B[pos[3]] {
		return true
	}
	return false
}