Merge pull request #5979 from astenmark/fix-5978-choice-bit-set

Fix #5978: make choice_bit_set respect bit_set domain

2 files changed, 25 insertions, 4 deletions

diff --git a/core/math/rand/rand.odin b/core/math/rand/rand.odin
index 4ffcc595e..41b066255 100644
--- a/core/math/rand/rand.odin
+++ b/core/math/rand/rand.odin
@@ -1199,11 +1199,14 @@ choice_bit_set :: proc(set: $T/bit_set[$E], gen := context.random_generator) ->
 		return {}, false
 	}
 
-	core_set := transmute(intrinsics.type_bit_set_underlying_type(T))set
+	target := int_max(total_set, gen)
 
-	for target := int_max(total_set, gen); target > 0; target -= 1 {
-		core_set &= core_set - 1
+	for value in set {
+		if target == 0 {
+			return value, true
+		}
+		target -= 1
 	}
 
-	return E(intrinsics.count_trailing_zeros(core_set)), true
+	return {}, false
 }
diff --git a/tests/core/math/rand/test_core_math_rand.odin b/tests/core/math/rand/test_core_math_rand.odin
index 814a1b9f8..5c2f4af84 100644
--- a/tests/core/math/rand/test_core_math_rand.odin
+++ b/tests/core/math/rand/test_core_math_rand.odin
@@ -78,6 +78,24 @@ rand_issue_5881 :: proc(t:^testing.T, rng: Generator) {
 	expect_quaternion_sign_uniformity(t, rng, 200_000)
 }
 
+@(test)
+test_issue_5978 :: proc(t:^testing.T) {
+	// Tests issue #5978 https://github.com/odin-lang/Odin/issues/5978
+
+	s := bit_set[1 ..= 5]{1, 5}
+
+	cases := []struct {
+		seed: u64,
+		expected: int,
+	}{ {13, 1}, {27, 5} }
+
+	for c in cases {
+		rand.reset(c.seed)
+		i, _ := rand.choice_bit_set(s)
+		testing.expectf(t, i == c.expected, "choice_bit_set returned %v with seed %v, expected %v", i, c.seed, c.expected)
+	}
+}
+
 // Helper: compute chi-square statistic for counts vs equal-expected across k bins
 @(private = "file")
 chi_square_equal :: proc(counts: []int) -> f64 {