From ad11d3bea03eeb74e43fc518321943db7341e7be Mon Sep 17 00:00:00 2001
From: Andreas Stenmark <astenmark@users.noreply.github.com>
Date: Sun, 30 Nov 2025 21:46:26 +0100
Subject: [PATCH] Fix #5978: choice_bit_set respects bit_set domain

---
 core/math/rand/rand.odin                      | 11 +++++++----
 tests/core/math/rand/test_core_math_rand.odin | 18 ++++++++++++++++++
 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..19f844086 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 {