zig

commit 62fe4a0ba8a08b9477e772841d3ae128937f0752 (tree)
parent 7e7d67d8eed45bcf3908edd2f4ca864144fffad5
Author: Justin Whear <justin.whear+github@gmail.com>
Date:   Thu, 19 Aug 2021 12:18:23 -0700

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>
Diffstat:
M
lib/std/rand.zig
 | 
30
++++++++++++++++++++++++++++++

1 file changed, 30 insertions(+), 0 deletions(-)
diff --git a/lib/std/rand.zig b/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();