authorgravatar for si@sjbrown.co.ukSimon Brown <si@sjbrown.co.uk> 2024-03-24 22:29:27+00:00
committergravatar for noreply@github.comGitHub <noreply@github.com> 2024-03-24 22:29:27+00:00
log5c628312b16ce972cc3108ed44eed47960e17af7
tree80bfded3ff2d6f94b6d07f2ac7e13ad7bd4825b5
parentaf0668d6c2c7696e695cf0553c1ae2d593e9effe
signaturebadge-check Signed by PGP key B5690EEEBB952194

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

1 files changed, 51 insertions(+), 13 deletions(-)

lib/std/enums.zig+51-13
......@@ -240,7 +240,7 @@ test nameCast {
240240}
241241
242242/// A set of enum elements, backed by a bitfield. If the enum
243/// is not dense, a mapping will be constructed from enum values
243/// is exhaustive but not dense, a mapping will be constructed from enum values
244244/// to dense indices. This type does no dynamic allocation and
245245/// can be copied by value.
246246pub fn EnumSet(comptime E: type) type {
......@@ -263,11 +263,21 @@ pub fn EnumSet(comptime E: type) type {
263263 pub fn init(init_values: EnumFieldStruct(E, bool, false)) Self {
264264 @setEvalBranchQuota(2 * @typeInfo(E).Enum.fields.len);
265265 var result: Self = .{};
266 inline for (0..Self.len) |i| {
267 const key = comptime Indexer.keyForIndex(i);
268 const tag = @tagName(key);
269 if (@field(init_values, tag)) {
270 result.bits.set(i);
266 if (@typeInfo(E).Enum.is_exhaustive) {
267 inline for (0..Self.len) |i| {
268 const key = comptime Indexer.keyForIndex(i);
269 const tag = @tagName(key);
270 if (@field(init_values, tag)) {
271 result.bits.set(i);
272 }
273 }
274 } else {
275 inline for (std.meta.fields(E)) |field| {
276 const key = @field(E, field.name);
277 if (@field(init_values, field.name)) {
278 const i = comptime Indexer.indexOf(key);
279 result.bits.set(i);
280 }
271281 }
272282 }
273283 return result;
......@@ -416,7 +426,7 @@ pub fn EnumSet(comptime E: type) type {
416426}
417427
418428/// A map keyed by an enum, backed by a bitfield and a dense array.
419/// If the enum is not dense, a mapping will be constructed from
429/// If the enum is exhaustive but not dense, a mapping will be constructed from
420430/// enum values to dense indices. This type does no dynamic
421431/// allocation and can be copied by value.
422432pub fn EnumMap(comptime E: type, comptime V: type) type {
......@@ -444,12 +454,23 @@ pub fn EnumMap(comptime E: type, comptime V: type) type {
444454 pub fn init(init_values: EnumFieldStruct(E, ?Value, null)) Self {
445455 @setEvalBranchQuota(2 * @typeInfo(E).Enum.fields.len);
446456 var result: Self = .{};
447 inline for (0..Self.len) |i| {
448 const key = comptime Indexer.keyForIndex(i);
449 const tag = @tagName(key);
450 if (@field(init_values, tag)) |*v| {
451 result.bits.set(i);
452 result.values[i] = v.*;
457 if (@typeInfo(E).Enum.is_exhaustive) {
458 inline for (0..Self.len) |i| {
459 const key = comptime Indexer.keyForIndex(i);
460 const tag = @tagName(key);
461 if (@field(init_values, tag)) |*v| {
462 result.bits.set(i);
463 result.values[i] = v.*;
464 }
465 }
466 } else {
467 inline for (std.meta.fields(E)) |field| {
468 const key = @field(E, field.name);
469 if (@field(init_values, field.name)) |*v| {
470 const i = comptime Indexer.indexOf(key);
471 result.bits.set(i);
472 result.values[i] = v.*;
473 }
453474 }
454475 }
455476 return result;
......@@ -1222,6 +1243,23 @@ test "EnumSet const iterator" {
12221243 try testing.expect(result.eql(diag_move));
12231244}
12241245
1246test "EnumSet non-exhaustive" {
1247 const BitIndices = enum(u4) {
1248 a = 0,
1249 b = 1,
1250 c = 4,
1251 _,
1252 };
1253 const BitField = EnumSet(BitIndices);
1254
1255 var flags = BitField.init(.{ .a = true, .b = true });
1256 flags.insert(.c);
1257 flags.remove(.a);
1258 try testing.expect(!flags.contains(.a));
1259 try testing.expect(flags.contains(.b));
1260 try testing.expect(flags.contains(.c));
1261}
1262
12251263pub fn EnumIndexer(comptime E: type) type {
12261264 // Assumes that the enum fields are sorted in ascending order (optimistic).
12271265 // Unsorted enums may require the user to manually increase the quota.