From ac1e73e249f8ce06bc7c89d2bdc4359b0399236c Mon Sep 17 00:00:00 2001 From: mlugg Date: Thu, 31 Jul 2025 10:10:40 +0100 Subject: [PATCH] std.enums: fix `EnumIndexer` branch quota It's quite silly to have this override which nonetheless makes assumptions about the input type. Encode the actual complexity of the sort. Also, simplify the sorting logic, and fix a bug (grab min and max *after* the sort, not *before*!) --- lib/std/enums.zig | 58 +++++++++++++++++++++++++++++------------------ 1 file changed, 36 insertions(+), 22 deletions(-) diff --git a/lib/std/enums.zig b/lib/std/enums.zig index e34cf17e07..e47af4aca3 100644 --- a/lib/std/enums.zig +++ b/lib/std/enums.zig @@ -1317,9 +1317,9 @@ test "EnumSet non-exhaustive" { } 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. - @setEvalBranchQuota(3 * @typeInfo(E).@"enum".fields.len + eval_branch_quota_cushion); + // n log n for `std.mem.sortUnstable` call below. + const fields_len = @typeInfo(E).@"enum".fields.len; + @setEvalBranchQuota(3 * fields_len * std.math.log2(@max(fields_len, 1)) + eval_branch_quota_cushion); if (!@typeInfo(E).@"enum".is_exhaustive) { const BackingInt = @typeInfo(E).@"enum".tag_type; @@ -1354,10 +1354,6 @@ pub fn EnumIndexer(comptime E: type) type { }; } - const const_fields = @typeInfo(E).@"enum".fields; - var fields = const_fields[0..const_fields.len].*; - const fields_len = fields.len; - if (fields_len == 0) { return struct { pub const Key = E; @@ -1373,22 +1369,17 @@ pub fn EnumIndexer(comptime E: type) type { }; } + var fields: [fields_len]EnumField = @typeInfo(E).@"enum".fields[0..].*; + + std.mem.sortUnstable(EnumField, &fields, {}, struct { + fn lessThan(ctx: void, lhs: EnumField, rhs: EnumField) bool { + ctx; + return lhs.value < rhs.value; + } + }.lessThan); + const min = fields[0].value; - const max = fields[fields.len - 1].value; - - const SortContext = struct { - fields: []EnumField, - - pub fn lessThan(comptime ctx: @This(), comptime a: usize, comptime b: usize) bool { - return ctx.fields[a].value < ctx.fields[b].value; - } - - pub fn swap(comptime ctx: @This(), comptime a: usize, comptime b: usize) void { - return std.mem.swap(EnumField, &ctx.fields[a], &ctx.fields[b]); - } - }; - std.sort.insertionContext(0, fields_len, SortContext{ .fields = &fields }); - + const max = fields[fields_len - 1].value; if (max - min == fields.len - 1) { return struct { pub const Key = E; @@ -1538,6 +1529,29 @@ test "EnumIndexer empty" { try testing.expectEqual(0, Indexer.count); } +test "EnumIndexer large dense unsorted" { + @setEvalBranchQuota(500_000); // many `comptimePrint`s + // Make an enum with 500 fields with values in *descending* order. + const E = @Type(.{ .@"enum" = .{ + .tag_type = u32, + .fields = comptime fields: { + var fields: [500]EnumField = undefined; + for (&fields, 0..) |*f, i| f.* = .{ + .name = std.fmt.comptimePrint("f{d}", .{i}), + .value = 500 - i, + }; + break :fields &fields; + }, + .decls = &.{}, + .is_exhaustive = true, + } }); + const Indexer = EnumIndexer(E); + try testing.expectEqual(E.f0, Indexer.keyForIndex(499)); + try testing.expectEqual(E.f499, Indexer.keyForIndex(0)); + try testing.expectEqual(499, Indexer.indexOf(.f0)); + try testing.expectEqual(0, Indexer.indexOf(.f499)); +} + test values { const E = enum { X,