From 56f61545cee120445488ec06015e8f8dd778c7d7 Mon Sep 17 00:00:00 2001
From: Ian Gulliver <git@flamingcow.io>
Date: Sat, 20 Nov 2021 16:56:07 -1000
Subject: [PATCH] Signed biased distribution generator

---
 gen/rand.go       | 16 +++++++++++++++-
 test/rand_test.go | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/gen/rand.go b/gen/rand.go
index 1f3c523..ca1bcbd 100644
--- a/gen/rand.go
+++ b/gen/rand.go
@@ -6,5 +6,19 @@ import "math/rand"
 func RandBiasedUint64() uint64 {
 	// The shift-right by up to 64 (shifting it to 0) makes up for randomness
         // lost by setting the high bit.
-	return (rand.Uint64() | 0x8000000000000000) >> rand.Int31n(65)
+	return (rand.Uint64() | 0x8000000000000000) >> rand.Intn(65)
+}
+
+// Generate a random int64 with an even distribution of the tuple
+// {sign, bits.Len64(math.Abs())}
+func RandBiasedInt64() int64 {
+	shift := rand.Intn(127)
+	// [0,62]: positive
+	// [63,63]: zero
+	// [64,126]: negative
+	if shift < 64 {
+		return int64((rand.Uint64() | 0x8000000000000000) >> (shift + 1))
+	} else {
+		return int64((rand.Uint64() | 0x8000000000000000) >> (shift - 63)) * -1
+	}
 }
diff --git a/test/rand_test.go b/test/rand_test.go
index 5821bab..ea9abfe 100644
--- a/test/rand_test.go
+++ b/test/rand_test.go
@@ -31,3 +31,39 @@ func TestRandBiasedUint64(t *testing.T) {
 		t.Fatalf("Variance greater than allowed: %d > %d (max=%d min=%d)", max - min, max / 10, max, min)
 	}
 }
+
+func TestRandBiasedInt64(t *testing.T) {
+	// [0,62]: positive
+	// [63,63]: zero
+	// [64,126]: negative
+	buckets := [127]uint64{}
+
+	for i := 0; i < 1000000; i++ {
+		val := gen.RandBiasedInt64()
+		switch {
+		case val == 0:
+			buckets[63]++
+		case val < 0:
+			buckets[63 + bits.Len64(uint64(val * -1))]++
+		case val > 0:
+			buckets[63 - bits.Len64(uint64(val))]++
+		}
+	}
+
+	var max uint64 = 0
+	var min uint64 = math.MaxUint64
+
+	for _, count := range buckets {
+		if count > max {
+			max = count
+		}
+
+		if count < min {
+			min = count
+		}
+	}
+
+	if max - min > max / 10 {
+		t.Fatalf("Variance greater than allowed: %d > %d (max=%d min=%d)", max - min, max / 10, max, min)
+	}
+}