motiejus/zig
1
Fork 0
Code Releases Activity

std.enums: fix EnumSet.init and EnumMap.init for non-exhaustive enums

Browse Source
This commit is contained in:
Simon Brown
2024-03-24 22:29:27 +00:00
committed by GitHub
parent af0668d6c2
commit 5c628312b1
1 changed files with 51 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

64
lib/std/enums.zig
View File

@@ -240,7 +240,7 @@ test nameCast {
}
/// A set of enum elements, backed by a bitfield. If the enum
/// is not dense, a mapping will be constructed from enum values
/// is exhaustive but not dense, a mapping will be constructed from enum values
/// to dense indices. This type does no dynamic allocation and
/// can be copied by value.
pub fn EnumSet(comptime E: type) type {
@@ -263,11 +263,21 @@ pub fn EnumSet(comptime E: type) type {
pub fn init(init_values: EnumFieldStruct(E, bool, false)) Self {
@setEvalBranchQuota(2 * @typeInfo(E).Enum.fields.len);
var result: Self = .{};
inline for (0..Self.len) |i| {
const key = comptime Indexer.keyForIndex(i);
const tag = @tagName(key);
if (@field(init_values, tag)) {
result.bits.set(i);
if (@typeInfo(E).Enum.is_exhaustive) {
inline for (0..Self.len) |i| {
const key = comptime Indexer.keyForIndex(i);
const tag = @tagName(key);
if (@field(init_values, tag)) {
result.bits.set(i);
}
}
} else {
inline for (std.meta.fields(E)) |field| {
const key = @field(E, field.name);
if (@field(init_values, field.name)) {
const i = comptime Indexer.indexOf(key);
result.bits.set(i);
}
}
}
return result;
@@ -416,7 +426,7 @@ pub fn EnumSet(comptime E: type) type {
}
/// A map keyed by an enum, backed by a bitfield and a dense array.
/// If the enum is not dense, a mapping will be constructed from
/// If the enum is exhaustive but not dense, a mapping will be constructed from
/// enum values to dense indices. This type does no dynamic
/// allocation and can be copied by value.
pub fn EnumMap(comptime E: type, comptime V: type) type {
@@ -444,12 +454,23 @@ pub fn EnumMap(comptime E: type, comptime V: type) type {
pub fn init(init_values: EnumFieldStruct(E, ?Value, null)) Self {
@setEvalBranchQuota(2 * @typeInfo(E).Enum.fields.len);
var result: Self = .{};
inline for (0..Self.len) |i| {
const key = comptime Indexer.keyForIndex(i);
const tag = @tagName(key);
if (@field(init_values, tag)) |*v| {
result.bits.set(i);
result.values[i] = v.*;
if (@typeInfo(E).Enum.is_exhaustive) {
inline for (0..Self.len) |i| {
const key = comptime Indexer.keyForIndex(i);
const tag = @tagName(key);
if (@field(init_values, tag)) |*v| {
result.bits.set(i);
result.values[i] = v.*;
}
}
} else {
inline for (std.meta.fields(E)) |field| {
const key = @field(E, field.name);
if (@field(init_values, field.name)) |*v| {
const i = comptime Indexer.indexOf(key);
result.bits.set(i);
result.values[i] = v.*;
}
}
}
return result;
@@ -1222,6 +1243,23 @@ test "EnumSet const iterator" {
try testing.expect(result.eql(diag_move));
}
test "EnumSet non-exhaustive" {
const BitIndices = enum(u4) {
a = 0,
b = 1,
c = 4,
_,
};
const BitField = EnumSet(BitIndices);
var flags = BitField.init(.{ .a = true, .b = true });
flags.insert(.c);
flags.remove(.a);
try testing.expect(!flags.contains(.a));
try testing.expect(flags.contains(.b));
try testing.expect(flags.contains(.c));
}
pub fn EnumIndexer(comptime E: type) type {
// Assumes that the enum fields are sorted in ascending order (optimistic).
// Unsorted enums may require the user to manually increase the quota.