2 bloom filters

main.go

package main

import (
	"bufio"
	"encoding/binary"
	"fmt"
	"math"
	"os"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/bits-and-blooms/bloom/v3"
)

func main() {
	rate := 0.01
	if envRate := os.Getenv("ERROR_RATE"); envRate != "" {
		rate, _ = strconv.ParseFloat(envRate, 64)
	}

	fmt.Printf("Initializing bloom filters with ERROR_RATE=%.2f)\n", rate)

	even := bloom.NewWithEstimates(uint(math.Pow(2, 31)), rate)
	odd := bloom.NewWithEstimates(uint(math.Pow(2, 31)), rate)
	fmt.Println("Populating filters...")
	start := time.Now()

	var wg sync.WaitGroup
	workers := runtime.NumCPU()
	perWorker := uint64(math.MaxUint32) / uint64(workers)

	for w := range workers {
		wg.Add(1)
		wStart := uint64(w) * perWorker
		wEnd := uint64(w+1) * perWorker
		if w == workers-1 {
			wEnd = uint64(math.MaxUint32) + 1
		}
		go func(s, e uint64) {
			defer wg.Done()
			buf := make([]byte, 4)
			for i := s; i < e; i++ {
				binary.BigEndian.PutUint32(buf, uint32(i))
				if i%2 == 0 {
					even.Add(buf)
				} else {
					odd.Add(buf)
				}
			}
		}(wStart, wEnd)
	}
	wg.Wait()

	fmt.Printf("Completed in %d seconds\n", int(time.Since(start).Seconds()))
	fmt.Printf("Filter storage: %d MB\n\n", (even.Cap()+odd.Cap())/8/1024/1024)

	test(even, odd)

	// If stdin is a terminal, start REPL
	stat, _ := os.Stdin.Stat()
	if (stat.Mode() & os.ModeCharDevice) != 0 {
		repl(even, odd)
	}
}

func repl(even, odd *bloom.BloomFilter) {
	scanner := bufio.NewScanner(os.Stdin)
	fmt.Println("Enter a number (Ctrl-D to exit):")

	for {
		fmt.Print("> ")
		if !scanner.Scan() {
			if err := scanner.Err(); err != nil {
				fmt.Fprintf(os.Stderr, "\nError: %v\n", err)
			} else {
				fmt.Println()
			}
			break
		}

		input := strings.TrimSpace(scanner.Text())
		if input == "" {
			continue
		}

		num, err := strconv.ParseUint(input, 10, 32)
		if err != nil {
			fmt.Println("Invalid input")
			continue
		}

		fmt.Println(check(uint32(num), even, odd))
	}
}

func check(num uint32, even, odd *bloom.BloomFilter) string {
	buf := make([]byte, 4)
	binary.BigEndian.PutUint32(buf, num)

	inEven := even.Test(buf)
	inOdd := odd.Test(buf)

	if inEven && !inOdd {
		return "Even"
	} else if inOdd && !inEven {
		return "Odd"
	} else if inEven && inOdd {
		return "Unknown"
	} else {
		panic("")
	}
}

func test(even, odd *bloom.BloomFilter) {
	correct, unknown, total := 0, 0, 1000
	fmt.Printf("Self test (%d numbers)\n", total)

	for i := 0; i < total; i++ {
		num := uint32(i) * (math.MaxUint32 / uint32(total))
		result := check(num, even, odd)
		isEven := num%2 == 0

		if result == "Unknown" {
			unknown++
		} else if (isEven && result == "Even") || (result == "Odd") {
			correct++
		}
	}

	fmt.Printf("  correct: %d (%.1f%%)\n", correct, float64(correct)/float64(total)*100)
	fmt.Printf("  unknown: %d (%.1f%%)\n", unknown, float64(unknown)/float64(total)*100)
}