authorgravatar for gethwilliams@googlemail.comGethDW <gethwilliams@googlemail.com> 2023-03-23 10:00:10+00:00
committergravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2023-04-10 12:15:05-04:00
log2c639d657002ac66749d08c4977cbb201d113ce1
tree67f2fbc5bbaed7406c1fefd78e4796c1e4a60cbb
parent88dfb1381820d7a79ed301af4d728510353e0c10

std.MultiArrayList: add support for tagged unions.

3 files changed, 144 insertions(+), 41 deletions(-)

lib/std/multi_array_list.zig+141-38
...@@ -1,4 +1,4 @@...@@ -1,4 +1,4 @@
1const std = @import("std.zig");1const std = @import("std");
2const builtin = @import("builtin");2const builtin = @import("builtin");
3const assert = std.debug.assert;3const assert = std.debug.assert;
4const meta = std.meta;4const meta = std.meta;
...@@ -6,24 +6,57 @@ const mem = std.mem;...@@ -6,24 +6,57 @@ const mem = std.mem;
6const Allocator = mem.Allocator;6const Allocator = mem.Allocator;
7const testing = std.testing;7const testing = std.testing;
88
9/// A MultiArrayList stores a list of a struct type.9/// A MultiArrayList stores a list of a struct or tagged union type.
10/// Instead of storing a single list of items, MultiArrayList10/// Instead of storing a single list of items, MultiArrayList
11/// stores separate lists for each field of the struct.11/// stores separate lists for each field of the struct or
12/// This allows for memory savings if the struct has padding,12/// lists of tags and bare unions.
13/// and also improves cache usage if only some fields are needed13/// This allows for memory savings if the struct or union has padding,
14/// for a computation. The primary API for accessing fields is14/// and also improves cache usage if only some fields or or just tags
15/// are needed for a computation. The primary API for accessing fields is
15/// the `slice()` function, which computes the start pointers16/// the `slice()` function, which computes the start pointers
16/// for the array of each field. From the slice you can call17/// for the array of each field. From the slice you can call
17/// `.items(.<field_name>)` to obtain a slice of field values.18/// `.items(.<field_name>)` to obtain a slice of field values.
18pub fn MultiArrayList(comptime S: type) type {19/// For unions you can call `.items(.tags)` or `.items(.data)`.
20pub fn MultiArrayList(comptime T: type) type {
19 return struct {21 return struct {
20 bytes: [*]align(@alignOf(S)) u8 = undefined,22 bytes: [*]align(@alignOf(T)) u8 = undefined,
21 len: usize = 0,23 len: usize = 0,
22 capacity: usize = 0,24 capacity: usize = 0,
2325
24 pub const Elem = S;26 const Elem = switch (@typeInfo(T)) {
27 .Struct => T,
28 .Union => |u| struct {
29 pub const Bare =
30 @Type(.{ .Union = .{
31 .layout = u.layout,
32 .tag_type = null,
33 .fields = u.fields,
34 .decls = &.{},
35 } });
36 pub const Tag =
37 u.tag_type orelse @compileError("MultiArrayList does not support untagged unions");
38 tags: Tag,
39 data: Bare,
40
41 pub fn fromT(outer: T) @This() {
42 const tag = meta.activeTag(outer);
43 return .{
44 .tags = tag,
45 .data = switch (tag) {
46 inline else => |t| @unionInit(Bare, @tagName(t), @field(outer, @tagName(t))),
47 },
48 };
49 }
50 pub fn toT(tag: Tag, bare: Bare) T {
51 return switch (tag) {
52 inline else => |t| @unionInit(T, @tagName(t), @field(bare, @tagName(t))),
53 };
54 }
55 },
56 else => @compileError("MultiArrayList only supports structs and tagged unions"),
57 };
2558
26 pub const Field = meta.FieldEnum(S);59 pub const Field = meta.FieldEnum(Elem);
2760
28 /// A MultiArrayList.Slice contains cached start pointers for each field in the list.61 /// A MultiArrayList.Slice contains cached start pointers for each field in the list.
29 /// These pointers are not normally stored to reduce the size of the list in memory.62 /// These pointers are not normally stored to reduce the size of the list in memory.
...@@ -49,18 +82,27 @@ pub fn MultiArrayList(comptime S: type) type {...@@ -49,18 +82,27 @@ pub fn MultiArrayList(comptime S: type) type {
49 return casted_ptr[0..self.len];82 return casted_ptr[0..self.len];
50 }83 }
5184
52 pub fn set(self: Slice, index: usize, elem: S) void {85 pub fn set(self: *Slice, index: usize, elem: T) void {
53 inline for (fields) |field_info| {86 const e = switch (@typeInfo(T)) {
54 self.items(@field(Field, field_info.name))[index] = @field(elem, field_info.name);87 .Struct => elem,
88 .Union => Elem.fromT(elem),
89 else => unreachable,
90 };
91 inline for (fields, 0..) |field_info, i| {
92 self.items(@intToEnum(Field, i))[index] = @field(e, field_info.name);
55 }93 }
56 }94 }
5795
58 pub fn get(self: Slice, index: usize) S {96 pub fn get(self: Slice, index: usize) T {
59 var elem: S = undefined;97 var result: Elem = undefined;
60 inline for (fields) |field_info| {98 inline for (fields, 0..) |field_info, i| {
61 @field(elem, field_info.name) = self.items(@field(Field, field_info.name))[index];99 @field(result, field_info.name) = self.items(@intToEnum(Field, i))[index];
62 }100 }
63 return elem;101 return switch (@typeInfo(T)) {
102 .Struct => result,
103 .Union => Elem.toT(result.tags, result.data),
104 else => unreachable,
105 };
64 }106 }
65107
66 pub fn toMultiArrayList(self: Slice) Self {108 pub fn toMultiArrayList(self: Slice) Self {
...@@ -68,8 +110,8 @@ pub fn MultiArrayList(comptime S: type) type {...@@ -68,8 +110,8 @@ pub fn MultiArrayList(comptime S: type) type {
68 return .{};110 return .{};
69 }111 }
70 const unaligned_ptr = self.ptrs[sizes.fields[0]];112 const unaligned_ptr = self.ptrs[sizes.fields[0]];
71 const aligned_ptr = @alignCast(@alignOf(S), unaligned_ptr);113 const aligned_ptr = @alignCast(@alignOf(Elem), unaligned_ptr);
72 const casted_ptr = @ptrCast([*]align(@alignOf(S)) u8, aligned_ptr);114 const casted_ptr = @ptrCast([*]align(@alignOf(Elem)) u8, aligned_ptr);
73 return .{115 return .{
74 .bytes = casted_ptr,116 .bytes = casted_ptr,
75 .len = self.len,117 .len = self.len,
...@@ -85,7 +127,7 @@ pub fn MultiArrayList(comptime S: type) type {...@@ -85,7 +127,7 @@ pub fn MultiArrayList(comptime S: type) type {
85127
86 /// This function is used in the debugger pretty formatters in tools/ to fetch the128 /// This function is used in the debugger pretty formatters in tools/ to fetch the
87 /// child field order and entry type to facilitate fancy debug printing for this type.129 /// child field order and entry type to facilitate fancy debug printing for this type.
88 fn dbHelper(self: *Slice, child: *S, field: *Field, entry: *Entry) void {130 fn dbHelper(self: *Slice, child: *Elem, field: *Field, entry: *Entry) void {
89 _ = self;131 _ = self;
90 _ = child;132 _ = child;
91 _ = field;133 _ = field;
...@@ -95,8 +137,8 @@ pub fn MultiArrayList(comptime S: type) type {...@@ -95,8 +137,8 @@ pub fn MultiArrayList(comptime S: type) type {
95137
96 const Self = @This();138 const Self = @This();
97139
98 const fields = meta.fields(S);140 const fields = meta.fields(Elem);
99 /// `sizes.bytes` is an array of @sizeOf each S field. Sorted by alignment, descending.141 /// `sizes.bytes` is an array of @sizeOf each T field. Sorted by alignment, descending.
100 /// `sizes.fields` is an array mapping from `sizes.bytes` array index to field index.142 /// `sizes.fields` is an array mapping from `sizes.bytes` array index to field index.
101 const sizes = blk: {143 const sizes = blk: {
102 const Data = struct {144 const Data = struct {
...@@ -169,24 +211,25 @@ pub fn MultiArrayList(comptime S: type) type {...@@ -169,24 +211,25 @@ pub fn MultiArrayList(comptime S: type) type {
169 }211 }
170212
171 /// Overwrite one array element with new data.213 /// Overwrite one array element with new data.
172 pub fn set(self: *Self, index: usize, elem: S) void {214 pub fn set(self: *Self, index: usize, elem: T) void {
173 return self.slice().set(index, elem);215 var slices = self.slice();
216 slices.set(index, elem);
174 }217 }
175218
176 /// Obtain all the data for one array element.219 /// Obtain all the data for one array element.
177 pub fn get(self: Self, index: usize) S {220 pub fn get(self: Self, index: usize) T {
178 return self.slice().get(index);221 return self.slice().get(index);
179 }222 }
180223
181 /// Extend the list by 1 element. Allocates more memory as necessary.224 /// Extend the list by 1 element. Allocates more memory as necessary.
182 pub fn append(self: *Self, gpa: Allocator, elem: S) !void {225 pub fn append(self: *Self, gpa: Allocator, elem: T) !void {
183 try self.ensureUnusedCapacity(gpa, 1);226 try self.ensureUnusedCapacity(gpa, 1);
184 self.appendAssumeCapacity(elem);227 self.appendAssumeCapacity(elem);
185 }228 }
186229
187 /// Extend the list by 1 element, but asserting `self.capacity`230 /// Extend the list by 1 element, but asserting `self.capacity`
188 /// is sufficient to hold an additional item.231 /// is sufficient to hold an additional item.
189 pub fn appendAssumeCapacity(self: *Self, elem: S) void {232 pub fn appendAssumeCapacity(self: *Self, elem: T) void {
190 assert(self.len < self.capacity);233 assert(self.len < self.capacity);
191 self.len += 1;234 self.len += 1;
192 self.set(self.len - 1, elem);235 self.set(self.len - 1, elem);
...@@ -213,7 +256,7 @@ pub fn MultiArrayList(comptime S: type) type {...@@ -213,7 +256,7 @@ pub fn MultiArrayList(comptime S: type) type {
213 /// Remove and return the last element from the list.256 /// Remove and return the last element from the list.
214 /// Asserts the list has at least one item.257 /// Asserts the list has at least one item.
215 /// Invalidates pointers to fields of the removed element.258 /// Invalidates pointers to fields of the removed element.
216 pub fn pop(self: *Self) S {259 pub fn pop(self: *Self) T {
217 const val = self.get(self.len - 1);260 const val = self.get(self.len - 1);
218 self.len -= 1;261 self.len -= 1;
219 return val;262 return val;
...@@ -222,7 +265,7 @@ pub fn MultiArrayList(comptime S: type) type {...@@ -222,7 +265,7 @@ pub fn MultiArrayList(comptime S: type) type {
222 /// Remove and return the last element from the list, or265 /// Remove and return the last element from the list, or
223 /// return `null` if list is empty.266 /// return `null` if list is empty.
224 /// Invalidates pointers to fields of the removed element, if any.267 /// Invalidates pointers to fields of the removed element, if any.
225 pub fn popOrNull(self: *Self) ?S {268 pub fn popOrNull(self: *Self) ?T {
226 if (self.len == 0) return null;269 if (self.len == 0) return null;
227 return self.pop();270 return self.pop();
228 }271 }
...@@ -231,7 +274,7 @@ pub fn MultiArrayList(comptime S: type) type {...@@ -231,7 +274,7 @@ pub fn MultiArrayList(comptime S: type) type {
231 /// after and including the specified index back by one and274 /// after and including the specified index back by one and
232 /// sets the given index to the specified element. May reallocate275 /// sets the given index to the specified element. May reallocate
233 /// and invalidate iterators.276 /// and invalidate iterators.
234 pub fn insert(self: *Self, gpa: Allocator, index: usize, elem: S) !void {277 pub fn insert(self: *Self, gpa: Allocator, index: usize, elem: T) !void {
235 try self.ensureUnusedCapacity(gpa, 1);278 try self.ensureUnusedCapacity(gpa, 1);
236 self.insertAssumeCapacity(index, elem);279 self.insertAssumeCapacity(index, elem);
237 }280 }
...@@ -240,10 +283,15 @@ pub fn MultiArrayList(comptime S: type) type {...@@ -240,10 +283,15 @@ pub fn MultiArrayList(comptime S: type) type {
240 /// Shifts all elements after and including the specified index283 /// Shifts all elements after and including the specified index
241 /// back by one and sets the given index to the specified element.284 /// back by one and sets the given index to the specified element.
242 /// Will not reallocate the array, does not invalidate iterators.285 /// Will not reallocate the array, does not invalidate iterators.
243 pub fn insertAssumeCapacity(self: *Self, index: usize, elem: S) void {286 pub fn insertAssumeCapacity(self: *Self, index: usize, elem: T) void {
244 assert(self.len < self.capacity);287 assert(self.len < self.capacity);
245 assert(index <= self.len);288 assert(index <= self.len);
246 self.len += 1;289 self.len += 1;
290 const entry = switch (@typeInfo(T)) {
291 .Struct => elem,
292 .Union => Elem.fromT(elem),
293 else => unreachable,
294 };
247 const slices = self.slice();295 const slices = self.slice();
248 inline for (fields, 0..) |field_info, field_index| {296 inline for (fields, 0..) |field_info, field_index| {
249 const field_slice = slices.items(@intToEnum(Field, field_index));297 const field_slice = slices.items(@intToEnum(Field, field_index));
...@@ -251,7 +299,7 @@ pub fn MultiArrayList(comptime S: type) type {...@@ -251,7 +299,7 @@ pub fn MultiArrayList(comptime S: type) type {
251 while (i > index) : (i -= 1) {299 while (i > index) : (i -= 1) {
252 field_slice[i] = field_slice[i - 1];300 field_slice[i] = field_slice[i - 1];
253 }301 }
254 field_slice[index] = @field(elem, field_info.name);302 field_slice[index] = @field(entry, field_info.name);
255 }303 }
256 }304 }
257305
...@@ -304,7 +352,7 @@ pub fn MultiArrayList(comptime S: type) type {...@@ -304,7 +352,7 @@ pub fn MultiArrayList(comptime S: type) type {
304352
305 const other_bytes = gpa.alignedAlloc(353 const other_bytes = gpa.alignedAlloc(
306 u8,354 u8,
307 @alignOf(S),355 @alignOf(Elem),
308 capacityInBytes(new_len),356 capacityInBytes(new_len),
309 ) catch {357 ) catch {
310 const self_slice = self.slice();358 const self_slice = self.slice();
...@@ -375,7 +423,7 @@ pub fn MultiArrayList(comptime S: type) type {...@@ -375,7 +423,7 @@ pub fn MultiArrayList(comptime S: type) type {
375 assert(new_capacity >= self.len);423 assert(new_capacity >= self.len);
376 const new_bytes = try gpa.alignedAlloc(424 const new_bytes = try gpa.alignedAlloc(
377 u8,425 u8,
378 @alignOf(S),426 @alignOf(Elem),
379 capacityInBytes(new_capacity),427 capacityInBytes(new_capacity),
380 );428 );
381 if (self.len == 0) {429 if (self.len == 0) {
...@@ -453,12 +501,12 @@ pub fn MultiArrayList(comptime S: type) type {...@@ -453,12 +501,12 @@ pub fn MultiArrayList(comptime S: type) type {
453 return elem_bytes * capacity;501 return elem_bytes * capacity;
454 }502 }
455503
456 fn allocatedBytes(self: Self) []align(@alignOf(S)) u8 {504 fn allocatedBytes(self: Self) []align(@alignOf(Elem)) u8 {
457 return self.bytes[0..capacityInBytes(self.capacity)];505 return self.bytes[0..capacityInBytes(self.capacity)];
458 }506 }
459507
460 fn FieldType(comptime field: Field) type {508 fn FieldType(comptime field: Field) type {
461 return meta.fieldInfo(S, field).type;509 return meta.fieldInfo(Elem, field).type;
462 }510 }
463511
464 const Entry = entry: {512 const Entry = entry: {
...@@ -479,7 +527,7 @@ pub fn MultiArrayList(comptime S: type) type {...@@ -479,7 +527,7 @@ pub fn MultiArrayList(comptime S: type) type {
479 };527 };
480 /// This function is used in the debugger pretty formatters in tools/ to fetch the528 /// This function is used in the debugger pretty formatters in tools/ to fetch the
481 /// child field order and entry type to facilitate fancy debug printing for this type.529 /// child field order and entry type to facilitate fancy debug printing for this type.
482 fn dbHelper(self: *Self, child: *S, field: *Field, entry: *Entry) void {530 fn dbHelper(self: *Self, child: *Elem, field: *Field, entry: *Entry) void {
483 _ = self;531 _ = self;
484 _ = child;532 _ = child;
485 _ = field;533 _ = field;
...@@ -719,3 +767,58 @@ test "insert elements" {...@@ -719,3 +767,58 @@ test "insert elements" {
719 try testing.expectEqualSlices(u8, &[_]u8{ 1, 2 }, list.items(.a));767 try testing.expectEqualSlices(u8, &[_]u8{ 1, 2 }, list.items(.a));
720 try testing.expectEqualSlices(u32, &[_]u32{ 2, 3 }, list.items(.b));768 try testing.expectEqualSlices(u32, &[_]u32{ 2, 3 }, list.items(.b));
721}769}
770
771test "union" {
772 const ally = testing.allocator;
773
774 const Foo = union(enum) {
775 a: u32,
776 b: []const u8,
777 };
778
779 var list = MultiArrayList(Foo){};
780 defer list.deinit(ally);
781
782 try testing.expectEqual(@as(usize, 0), list.items(.tags).len);
783
784 try list.ensureTotalCapacity(ally, 2);
785
786 list.appendAssumeCapacity(.{ .a = 1 });
787 list.appendAssumeCapacity(.{ .b = "zigzag" });
788
789 try testing.expectEqualSlices(meta.Tag(Foo), list.items(.tags), &.{ .a, .b });
790 try testing.expectEqual(@as(usize, 2), list.items(.tags).len);
791
792 list.appendAssumeCapacity(.{ .b = "foobar" });
793 try testing.expectEqualStrings("zigzag", list.items(.data)[1].b);
794 try testing.expectEqualStrings("foobar", list.items(.data)[2].b);
795
796 // Add 6 more things to force a capacity increase.
797 for (0..6) |i| {
798 try list.append(ally, .{ .a = @intCast(u32, 4 + i) });
799 }
800
801 try testing.expectEqualSlices(
802 meta.Tag(Foo),
803 &.{ .a, .b, .b, .a, .a, .a, .a, .a, .a },
804 list.items(.tags),
805 );
806 try testing.expectEqual(list.get(0), .{ .a = 1 });
807 try testing.expectEqual(list.get(1), .{ .b = "zigzag" });
808 try testing.expectEqual(list.get(2), .{ .b = "foobar" });
809 try testing.expectEqual(list.get(3), .{ .a = 4 });
810 try testing.expectEqual(list.get(4), .{ .a = 5 });
811 try testing.expectEqual(list.get(5), .{ .a = 6 });
812 try testing.expectEqual(list.get(6), .{ .a = 7 });
813 try testing.expectEqual(list.get(7), .{ .a = 8 });
814 try testing.expectEqual(list.get(8), .{ .a = 9 });
815
816 list.shrinkAndFree(ally, 3);
817
818 try testing.expectEqual(@as(usize, 3), list.items(.tags).len);
819 try testing.expectEqualSlices(meta.Tag(Foo), list.items(.tags), &.{ .a, .b, .b });
820
821 try testing.expectEqual(list.get(0), .{ .a = 1 });
822 try testing.expectEqual(list.get(1), .{ .b = "zigzag" });
823 try testing.expectEqual(list.get(2), .{ .b = "foobar" });
824}
lib/std/zig/Parse.zig+2-2
...@@ -41,13 +41,13 @@ fn listToSpan(p: *Parse, list: []const Node.Index) !Node.SubRange {...@@ -41,13 +41,13 @@ fn listToSpan(p: *Parse, list: []const Node.Index) !Node.SubRange {
41 };41 };
42}42}
4343
44fn addNode(p: *Parse, elem: Ast.NodeList.Elem) Allocator.Error!Node.Index {44fn addNode(p: *Parse, elem: Ast.Node) Allocator.Error!Node.Index {
45 const result = @intCast(Node.Index, p.nodes.len);45 const result = @intCast(Node.Index, p.nodes.len);
46 try p.nodes.append(p.gpa, elem);46 try p.nodes.append(p.gpa, elem);
47 return result;47 return result;
48}48}
4949
50fn setNode(p: *Parse, i: usize, elem: Ast.NodeList.Elem) Node.Index {50fn setNode(p: *Parse, i: usize, elem: Ast.Node) Node.Index {
51 p.nodes.set(i, elem);51 p.nodes.set(i, elem);
52 return @intCast(Node.Index, i);52 return @intCast(Node.Index, i);
53}53}
src/translate_c/ast.zig+1-1
...@@ -845,7 +845,7 @@ const Context = struct {...@@ -845,7 +845,7 @@ const Context = struct {
845 };845 };
846 }846 }
847847
848 fn addNode(c: *Context, elem: std.zig.Ast.NodeList.Elem) Allocator.Error!NodeIndex {848 fn addNode(c: *Context, elem: std.zig.Ast.Node) Allocator.Error!NodeIndex {
849 const result = @intCast(NodeIndex, c.nodes.len);849 const result = @intCast(NodeIndex, c.nodes.len);
850 try c.nodes.append(c.gpa, elem);850 try c.nodes.append(c.gpa, elem);
851 return result;851 return result;