From 7ab11097be7aa28c4268a08ac59630f1981c88b0 Mon Sep 17 00:00:00 2001
From: Subhomoy Haldar <hungrybluedev@gmail.com>
Date: Thu, 29 Jun 2023 20:11:27 +0100
Subject: [PATCH] rand: fix edge case, when bit length is 31 and 63, add tests
 for `rand.intn(2147483647)!` etc (#18714)

---
 vlib/rand/rand.v                | 16 ++++++++++++----
 vlib/rand/random_numbers_test.v | 20 ++++++++++++++++++++
 2 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/vlib/rand/rand.v b/vlib/rand/rand.v
index 9bbf6498e6..4f5ce401dd 100644
--- a/vlib/rand/rand.v
+++ b/vlib/rand/rand.v
@@ -55,7 +55,7 @@ pub fn (mut rng PRNG) u32n(max u32) !u32 {
 	// the closest power of two. Then we loop until we find
 	// an int in the required range
 	bit_len := bits.len_32(max)
-	if bit_len == 32 {
+	if _unlikely_(bit_len == 32) {
 		for {
 			value := rng.u32()
 			if value < max {
@@ -63,7 +63,11 @@ pub fn (mut rng PRNG) u32n(max u32) !u32 {
 			}
 		}
 	} else {
-		mask := (u32(1) << (bit_len + 1)) - 1
+		mask := if _unlikely_(bit_len == 31) {
+			u32(0x7FFFFFFF)
+		} else {
+			(u32(1) << (bit_len + 1)) - 1
+		}
 		for {
 			value := rng.u32() & mask
 			if value < max {
@@ -81,7 +85,7 @@ pub fn (mut rng PRNG) u64n(max u64) !u64 {
 		return error('max must be positive integer')
 	}
 	bit_len := bits.len_64(max)
-	if bit_len == 64 {
+	if _unlikely_(bit_len == 64) {
 		for {
 			value := rng.u64()
 			if value < max {
@@ -89,7 +93,11 @@ pub fn (mut rng PRNG) u64n(max u64) !u64 {
 			}
 		}
 	} else {
-		mask := (u64(1) << (bit_len + 1)) - 1
+		mask := if _unlikely_(bit_len == 63) {
+			u64(0x7FFFFFFFFFFFFFFF)
+		} else {
+			(u64(1) << (bit_len + 1)) - 1
+		}
 		for {
 			value := rng.u64() & mask
 			if value < max {
diff --git a/vlib/rand/random_numbers_test.v b/vlib/rand/random_numbers_test.v
index cf63b9b05b..1ddbd80b80 100644
--- a/vlib/rand/random_numbers_test.v
+++ b/vlib/rand/random_numbers_test.v
@@ -454,3 +454,23 @@ fn test_element2() {
 		assert 4 != e
 	}
 }
+
+fn test_proper_masking() {
+	under32 := []int{len: 10, init: index * 0 + rand.intn(1073741823)!}
+	assert under32 != [0].repeat(10)
+
+	over32 := []int{len: 10, init: index * 0 + rand.intn(1073741824)!}
+	assert over32 != [0].repeat(10)
+
+	under64 := []i64{len: 10, init: index * 0 + rand.i64n(i64(4611686018427387903))!}
+	assert under64 != [i64(0)].repeat(10)
+
+	over64 := []i64{len: 10, init: index * 0 + rand.i64n(i64(4611686018427387904))!}
+	assert over64 != [i64(0)].repeat(10)
+
+	almost_full32 := []int{len: 10, init: index * 0 + rand.intn(2147483647)!}
+	assert almost_full32 != [0].repeat(10)
+
+	almost_full64 := []i64{len: 10, init: index * 0 + rand.i64n(i64(9223372036854775807))!}
+	assert almost_full64 != [i64(0)].repeat(10)
+}