package main_test

import (
	"math/bits"
	"testing"
)

var tests = []struct {
	v    uint64
	want int
}{
	{1, 1}, {2, 1}, {3, 2}, {4, 1}, {5, 2}, {9, 2}, {8, 1}, {16, 1},
         {32, 1}, {64, 1}, {127, 7}, {256, 1}, {65536, 1},
         {4294967296, 1}, {4294967295, 32},
}

var sink int

func BenchmarkCustomOnesCount(b *testing.B) {
	runBenchmarks(b, customOnesCountForPowerOf2)
}

func BenchmarkStdlibOnesCount(b *testing.B) {
	runBenchmarks(b, bits.OnesCount64)
}

func runBenchmarks(b *testing.B, fn func(uint64) int) {
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		for _, tt := range tests {
			sink = fn(tt.v)
			if sink != tt.want {
				b.Errorf("%d: got %d want %d", tt.v, sink, tt.want)
			}
		}
	}
	if sink == 0 {
		b.Fatal("Did not run the benchmark!")
	}
}

func customOnesCountForPowerOf2(x uint64) int {
	// Simulate absolute branch since powerOf2 check.
	if x > 0 && x&(x-1) == 0 {
		return 1
	}
	const m = 1<<64 - 1
	x = x>>1&(m0&m) + x&(m0&m)
	x = x>>2&(m1&m) + x&(m1&m)
	x = (x>>4 + x) & (m2 & m)
	x += x >> 8
	x += x >> 16
	x += x >> 32
	return int(x) & (1<<7 - 1)
}

// Inlining these constants code here to show what
// the code for OnesCount64 will look like after the alteration.
const m0 = 0x5555555555555555 // 01010101 ...
const m1 = 0x3333333333333333 // 00110011 ...
const m2 = 0x0f0f0f0f0f0f0f0f // 00001111 ...
const m3 = 0x00ff00ff00ff00ff // etc.
const m4 = 0x0000ffff0000ffff