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 @@
1const std = @import("std.zig");
1const std = @import("std");
22const builtin = @import("builtin");
33const assert = std.debug.assert;
44const meta = std.meta;
......@@ -6,24 +6,57 @@ const mem = std.mem;
66const Allocator = mem.Allocator;
77const 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.
1010/// Instead of storing a single list of items, MultiArrayList
11/// stores separate lists for each field of the struct.
12/// This allows for memory savings if the struct has padding,
13/// and also improves cache usage if only some fields are needed
14/// for a computation. The primary API for accessing fields is
11/// stores separate lists for each field of the struct or
12/// lists of tags and bare unions.
13/// This allows for memory savings if the struct or union has padding,
14/// 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
1516/// the `slice()` function, which computes the start pointers
1617/// for the array of each field. From the slice you can call
1718/// `.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 {
1921 return struct {
20 bytes: [*]align(@alignOf(S)) u8 = undefined,
22 bytes: [*]align(@alignOf(T)) u8 = undefined,
2123 len: usize = 0,
2224 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
2861 /// A MultiArrayList.Slice contains cached start pointers for each field in the list.
2962 /// 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 {
4982 return casted_ptr[0..self.len];
5083 }
5184
52 pub fn set(self: Slice, index: usize, elem: S) void {
53 inline for (fields) |field_info| {
54 self.items(@field(Field, field_info.name))[index] = @field(elem, field_info.name);
85 pub fn set(self: *Slice, index: usize, elem: T) void {
86 const e = switch (@typeInfo(T)) {
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);
5593 }
5694 }
5795
58 pub fn get(self: Slice, index: usize) S {
59 var elem: S = undefined;
60 inline for (fields) |field_info| {
61 @field(elem, field_info.name) = self.items(@field(Field, field_info.name))[index];
96 pub fn get(self: Slice, index: usize) T {
97 var result: Elem = undefined;
98 inline for (fields, 0..) |field_info, i| {
99 @field(result, field_info.name) = self.items(@intToEnum(Field, i))[index];
62100 }
63 return elem;
101 return switch (@typeInfo(T)) {
102 .Struct => result,
103 .Union => Elem.toT(result.tags, result.data),
104 else => unreachable,
105 };
64106 }
65107
66108 pub fn toMultiArrayList(self: Slice) Self {
......@@ -68,8 +110,8 @@ pub fn MultiArrayList(comptime S: type) type {
68110 return .{};
69111 }
70112 const unaligned_ptr = self.ptrs[sizes.fields[0]];
71 const aligned_ptr = @alignCast(@alignOf(S), unaligned_ptr);
72 const casted_ptr = @ptrCast([*]align(@alignOf(S)) u8, aligned_ptr);
113 const aligned_ptr = @alignCast(@alignOf(Elem), unaligned_ptr);
114 const casted_ptr = @ptrCast([*]align(@alignOf(Elem)) u8, aligned_ptr);
73115 return .{
74116 .bytes = casted_ptr,
75117 .len = self.len,
......@@ -85,7 +127,7 @@ pub fn MultiArrayList(comptime S: type) type {
85127
86128 /// This function is used in the debugger pretty formatters in tools/ to fetch the
87129 /// 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 {
89131 _ = self;
90132 _ = child;
91133 _ = field;
......@@ -95,8 +137,8 @@ pub fn MultiArrayList(comptime S: type) type {
95137
96138 const Self = @This();
97139
98 const fields = meta.fields(S);
99 /// `sizes.bytes` is an array of @sizeOf each S field. Sorted by alignment, descending.
140 const fields = meta.fields(Elem);
141 /// `sizes.bytes` is an array of @sizeOf each T field. Sorted by alignment, descending.
100142 /// `sizes.fields` is an array mapping from `sizes.bytes` array index to field index.
101143 const sizes = blk: {
102144 const Data = struct {
......@@ -169,24 +211,25 @@ pub fn MultiArrayList(comptime S: type) type {
169211 }
170212
171213 /// Overwrite one array element with new data.
172 pub fn set(self: *Self, index: usize, elem: S) void {
173 return self.slice().set(index, elem);
214 pub fn set(self: *Self, index: usize, elem: T) void {
215 var slices = self.slice();
216 slices.set(index, elem);
174217 }
175218
176219 /// 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 {
178221 return self.slice().get(index);
179222 }
180223
181224 /// 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 {
183226 try self.ensureUnusedCapacity(gpa, 1);
184227 self.appendAssumeCapacity(elem);
185228 }
186229
187230 /// Extend the list by 1 element, but asserting `self.capacity`
188231 /// 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 {
190233 assert(self.len < self.capacity);
191234 self.len += 1;
192235 self.set(self.len - 1, elem);
......@@ -213,7 +256,7 @@ pub fn MultiArrayList(comptime S: type) type {
213256 /// Remove and return the last element from the list.
214257 /// Asserts the list has at least one item.
215258 /// Invalidates pointers to fields of the removed element.
216 pub fn pop(self: *Self) S {
259 pub fn pop(self: *Self) T {
217260 const val = self.get(self.len - 1);
218261 self.len -= 1;
219262 return val;
......@@ -222,7 +265,7 @@ pub fn MultiArrayList(comptime S: type) type {
222265 /// Remove and return the last element from the list, or
223266 /// return `null` if list is empty.
224267 /// Invalidates pointers to fields of the removed element, if any.
225 pub fn popOrNull(self: *Self) ?S {
268 pub fn popOrNull(self: *Self) ?T {
226269 if (self.len == 0) return null;
227270 return self.pop();
228271 }
......@@ -231,7 +274,7 @@ pub fn MultiArrayList(comptime S: type) type {
231274 /// after and including the specified index back by one and
232275 /// sets the given index to the specified element. May reallocate
233276 /// 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 {
235278 try self.ensureUnusedCapacity(gpa, 1);
236279 self.insertAssumeCapacity(index, elem);
237280 }
......@@ -240,10 +283,15 @@ pub fn MultiArrayList(comptime S: type) type {
240283 /// Shifts all elements after and including the specified index
241284 /// back by one and sets the given index to the specified element.
242285 /// 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 {
244287 assert(self.len < self.capacity);
245288 assert(index <= self.len);
246289 self.len += 1;
290 const entry = switch (@typeInfo(T)) {
291 .Struct => elem,
292 .Union => Elem.fromT(elem),
293 else => unreachable,
294 };
247295 const slices = self.slice();
248296 inline for (fields, 0..) |field_info, field_index| {
249297 const field_slice = slices.items(@intToEnum(Field, field_index));
......@@ -251,7 +299,7 @@ pub fn MultiArrayList(comptime S: type) type {
251299 while (i > index) : (i -= 1) {
252300 field_slice[i] = field_slice[i - 1];
253301 }
254 field_slice[index] = @field(elem, field_info.name);
302 field_slice[index] = @field(entry, field_info.name);
255303 }
256304 }
257305
......@@ -304,7 +352,7 @@ pub fn MultiArrayList(comptime S: type) type {
304352
305353 const other_bytes = gpa.alignedAlloc(
306354 u8,
307 @alignOf(S),
355 @alignOf(Elem),
308356 capacityInBytes(new_len),
309357 ) catch {
310358 const self_slice = self.slice();
......@@ -375,7 +423,7 @@ pub fn MultiArrayList(comptime S: type) type {
375423 assert(new_capacity >= self.len);
376424 const new_bytes = try gpa.alignedAlloc(
377425 u8,
378 @alignOf(S),
426 @alignOf(Elem),
379427 capacityInBytes(new_capacity),
380428 );
381429 if (self.len == 0) {
......@@ -453,12 +501,12 @@ pub fn MultiArrayList(comptime S: type) type {
453501 return elem_bytes * capacity;
454502 }
455503
456 fn allocatedBytes(self: Self) []align(@alignOf(S)) u8 {
504 fn allocatedBytes(self: Self) []align(@alignOf(Elem)) u8 {
457505 return self.bytes[0..capacityInBytes(self.capacity)];
458506 }
459507
460508 fn FieldType(comptime field: Field) type {
461 return meta.fieldInfo(S, field).type;
509 return meta.fieldInfo(Elem, field).type;
462510 }
463511
464512 const Entry = entry: {
......@@ -479,7 +527,7 @@ pub fn MultiArrayList(comptime S: type) type {
479527 };
480528 /// This function is used in the debugger pretty formatters in tools/ to fetch the
481529 /// 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 {
483531 _ = self;
484532 _ = child;
485533 _ = field;
......@@ -719,3 +767,58 @@ test "insert elements" {
719767 try testing.expectEqualSlices(u8, &[_]u8{ 1, 2 }, list.items(.a));
720768 try testing.expectEqualSlices(u32, &[_]u32{ 2, 3 }, list.items(.b));
721769}
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 {
4141 };
4242}
4343
44fn addNode(p: *Parse, elem: Ast.NodeList.Elem) Allocator.Error!Node.Index {
44fn addNode(p: *Parse, elem: Ast.Node) Allocator.Error!Node.Index {
4545 const result = @intCast(Node.Index, p.nodes.len);
4646 try p.nodes.append(p.gpa, elem);
4747 return result;
4848}
4949
50fn setNode(p: *Parse, i: usize, elem: Ast.NodeList.Elem) Node.Index {
50fn setNode(p: *Parse, i: usize, elem: Ast.Node) Node.Index {
5151 p.nodes.set(i, elem);
5252 return @intCast(Node.Index, i);
5353}
src/translate_c/ast.zig+1-1
......@@ -845,7 +845,7 @@ const Context = struct {
845845 };
846846 }
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 {
849849 const result = @intCast(NodeIndex, c.nodes.len);
850850 try c.nodes.append(c.gpa, elem);
851851 return result;