Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Mar 24, 2024
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
64 changes: 51 additions & 13 deletions lib/std/enums.zig
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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;
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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.
Expand Down