Add `rand.choice_enum`

1 files changed, 20 insertions, 0 deletions

diff --git a/core/math/rand/rand.odin b/core/math/rand/rand.odin
index 14894e82c..560dc8379 100644
--- a/core/math/rand/rand.odin
+++ b/core/math/rand/rand.odin
@@ -834,3 +834,23 @@ choice :: proc(array: $T/[]$E, r: ^Rand = nil) -> (res: E) {
 	}
 	return array[int63_max(n, r)]
 }
+
+
+@(require_results)
+choice_enum :: proc($T: typeid, r: ^Rand = nil) -> T
+	where
+		intrinsics.type_is_enum(T),
+		size_of(T) <= 8,
+		len(T) == cap(T) /* Only allow contiguous enum types */
+{
+	when intrinsics.type_is_unsigned(intrinsics.type_core_type(T)) &&
+	     u64(max(T)) > u64(max(i64)) {
+		i := uint64(r) % u64(len(T))
+		i += u64(min(T))
+		return T(i)
+	} else {
+		i := int63_max(i64(len(T)), r)
+		i += i64(min(T))
+		return T(i)
+	}
+}
\ No newline at end of file