Commit 62fe4a0ba8

Justin Whear <justin.whear+github@gmail.com>
2021-08-19 21:18:23
std.rand.Random: add enumValue() (#9583)
* add Random.enumValue() * edits suggested by review * applied zig fmt * Rewrite to use std.enums.values Implemented pfgithub's suggestion to rewrite against this function, greatly simplifying the implementation. Co-authored-by: Justin Whear <justin@economicmodeling.com>
1 parent 7e7d67d
Changed files (1)
lib
lib/std/rand.zig
@@ -47,6 +47,19 @@ pub const Random = struct {
         return r.int(u1) != 0;
     }
 
+    /// Returns a random value from an enum, evenly distributed.
+    pub fn enumValue(r: *Random, comptime EnumType: type) EnumType {
+        if (comptime !std.meta.trait.is(.Enum)(EnumType)) {
+            @compileError("Random.enumValue requires an enum type, not a " ++ @typeName(EnumType));
+        }
+
+        // We won't use int -> enum casting because enum elements can have
+        //  arbitrary values.  Instead we'll randomly pick one of the type's values.
+        const values = std.enums.values(EnumType);
+        const index = r.uintLessThan(usize, values.len);
+        return values[index];
+    }
+
     /// Returns a random int `i` such that `0 <= i <= maxInt(T)`.
     /// `i` is evenly distributed.
     pub fn int(r: *Random, comptime T: type) T {
@@ -377,6 +390,23 @@ fn testRandomBoolean() !void {
     try expect(r.random.boolean() == true);
 }
 
+test "Random enum" {
+    try testRandomEnumValue();
+    comptime try testRandomEnumValue();
+}
+fn testRandomEnumValue() !void {
+    const TestEnum = enum {
+        First,
+        Second,
+        Third,
+    };
+    var r = SequentialPrng.init();
+    r.next_value = 0;
+    try expect(r.random.enumValue(TestEnum) == TestEnum.First);
+    try expect(r.random.enumValue(TestEnum) == TestEnum.First);
+    try expect(r.random.enumValue(TestEnum) == TestEnum.First);
+}
+
 test "Random intLessThan" {
     @setEvalBranchQuota(10000);
     try testRandomIntLessThan();