diff --git a/lib/std/rand.zig b/lib/std/rand.zig index 66f758cb48..7d967d3715 100644 --- 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();