lib/choir/src/bytecode/image.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

  1 const std = @import("std");
  2 const bytecode = @import("root.zig");
  3 const Reader = bytecode.Reader;
  4 
  5 pub const Limits = struct { bytes: u32, entities: u32, depth: u16 };
  6 pub const Range = struct { start: u32, count: u32 };
  7 
  8 /// A counted sequence in the existing bytecode format, borrowing immutable bytes.
  9 pub const Ids = struct {
 10     bytes: []const u8,
 11     count: u32,
 12 
 13     pub fn at(self: Ids, index: u32) ?u32 {
 14         if (index >= self.count) return null;
 15         var reader = Reader.init(self.bytes);
 16         for (0..index) |_| _ = reader.readULEBToU32() catch return null;
 17         return reader.readULEBToU32() catch null;
 18     }
 19 };
 20 
 21 pub const Operation = struct {
 22     name: []const u8,
 23     location: u32,
 24     parent_block: ?u32,
 25     results: Range,
 26     operands: Ids,
 27     attributes: []const u8,
 28     attribute_count: u32,
 29     properties: ?u32,
 30     successors: Ids,
 31     region_count: u32,
 32     encoded: []const u8,
 33 };
 34 
 35 pub const Region = struct { operation: u32, position: u32, block_count: u32 };
 36 pub const Block = struct { region: u32, position: u32, arguments: Range, operation_count: u32 };
 37 pub const Value = struct {
 38     owner: union(enum) { operation: u32, block: u32 },
 39     position: u32,
 40     type: u32,
 41     location: u32,
 42 };
 43 
 44 /// All ordinals and table records refer to this one bytecode image.
 45 pub const View = struct {
 46     bytes: []const u8,
 47     operations: []const Operation,
 48     regions: []const Region,
 49     blocks: []const Block,
 50     values: []const Value,
 51     strings: []const []const u8,
 52     dialects: []const bytecode.DialectEntry,
 53     types: []const []const u8,
 54     attributes: []const []const u8,
 55     locations: []const []const u8,
 56     resources: []const bytecode.Resource,
 57 
 58     pub fn region(self: View, operation: u32, position: u32) ?u32 {
 59         for (self.regions, 0..) |item, ordinal| {
 60             if (item.operation == operation and item.position == position) return @intCast(ordinal);
 61         }
 62         return null;
 63     }
 64 
 65     pub fn block(self: View, region_ordinal: u32, position: u32) ?u32 {
 66         for (self.blocks, 0..) |item, ordinal| {
 67             if (item.region == region_ordinal and item.position == position) {
 68                 return @intCast(ordinal);
 69             }
 70         }
 71         return null;
 72     }
 73 
 74     pub fn successor(self: View, operation: u32, position: u32) ?u32 {
 75         if (operation >= self.operations.len) return null;
 76         const op = self.operations[operation];
 77         const parent = op.parent_block orelse return null;
 78         const local = op.successors.at(position) orelse return null;
 79         return self.block(self.blocks[parent].region, local);
 80     }
 81 };
 82 
 83 /// Owns only lookup metadata. The caller keeps the immutable bytecode bytes alive.
 84 pub const Index = opaque {
 85     pub fn create(allocator: std.mem.Allocator, bytes: []const u8, limits: Limits) !*Index {
 86         if (bytes.len > limits.bytes) return error.ImageLimit;
 87         const state = try allocator.create(Data);
 88         state.* = .{ .allocator = allocator, .bytes = bytes, .limits = limits };
 89         errdefer state.destroy();
 90         try state.read();
 91         return @ptrCast(state);
 92     }
 93 
 94     pub fn destroy(self: *Index) void {
 95         data(self).destroy();
 96     }
 97 
 98     pub fn view(self: *const Index) View {
 99         const state: *const Data = @ptrCast(@alignCast(self));
100         return .{
101             .bytes = state.bytes,
102             .operations = state.operations.items,
103             .regions = state.regions.items,
104             .blocks = state.blocks.items,
105             .values = state.values.items,
106             .strings = state.strings.items,
107             .dialects = state.dialects.items,
108             .types = state.types.items,
109             .attributes = state.attributes.items,
110             .locations = state.locations.items,
111             .resources = state.resources.items,
112         };
113     }
114 };
115 
116 const Data = struct {
117     allocator: std.mem.Allocator,
118     bytes: []const u8,
119     limits: Limits,
120     entities: u32 = 0,
121     operations: std.ArrayList(Operation) = .empty,
122     regions: std.ArrayList(Region) = .empty,
123     blocks: std.ArrayList(Block) = .empty,
124     values: std.ArrayList(Value) = .empty,
125     strings: std.ArrayList([]const u8) = .empty,
126     dialects: std.ArrayList(bytecode.DialectEntry) = .empty,
127     types: std.ArrayList([]const u8) = .empty,
128     attributes: std.ArrayList([]const u8) = .empty,
129     locations: std.ArrayList([]const u8) = .empty,
130     resources: std.ArrayList(bytecode.Resource) = .empty,
131 
132     fn destroy(self: *Data) void {
133         inline for (.{
134             "operations", "regions",    "blocks",    "values",    "strings", "dialects",
135             "types",      "attributes", "locations", "resources",
136         }) |field| @field(self, field).deinit(self.allocator);
137         self.allocator.destroy(self);
138     }
139 
140     fn append(self: *Data, comptime field: []const u8, value: anytype) !u32 {
141         if (self.entities == self.limits.entities) return error.ImageLimit;
142         const list = &@field(self, field);
143         const ordinal: u32 = @intCast(list.items.len);
144         try list.append(self.allocator, value);
145         self.entities += 1;
146         return ordinal;
147     }
148 
149     fn count(self: *const Data, reader: *Reader) !u32 {
150         const result = try reader.readULEBToU32();
151         if (result > self.limits.entities) return error.ImageLimit;
152         return result;
153     }
154 
155     fn read(self: *Data) !void {
156         const sections = try bytecode.inspectSections(self.bytes);
157         try self.readStrings(sections.strings orelse return error.MissingStrings);
158         if (sections.dialects) |bytes| try self.readDialects(bytes);
159         if (sections.types) |bytes| try self.readTable(bytes, "types", readType);
160         if (sections.attrs) |bytes| try self.readTable(bytes, "attributes", readAttribute);
161         if (sections.locs) |bytes| try self.readTable(bytes, "locations", readLocation);
162         if (sections.resources) |bytes| try self.readResources(bytes);
163         var reader = Reader.init(sections.module orelse return error.MissingModule);
164         try self.readOperation(&reader, null, 0, 0);
165         try atEnd(reader);
166     }
167 
168     fn readStrings(self: *Data, bytes: []const u8) !void {
169         var reader = Reader.init(bytes);
170         const length = try self.count(&reader);
171         for (0..length) |_| _ = try self.append("strings", try reader.readBytesWithLen());
172         try atEnd(reader);
173     }
174 
175     fn readDialects(self: *Data, bytes: []const u8) !void {
176         var reader = Reader.init(bytes);
177         const length = try self.count(&reader);
178         for (0..length) |_| {
179             const name = try self.string(&reader);
180             _ = try self.append("dialects", bytecode.DialectEntry{
181                 .name = name,
182                 .version = try reader.readULEBToU32(),
183                 .flags = try reader.readULEBToU32(),
184             });
185         }
186         try atEnd(reader);
187     }
188 
189     fn readTable(
190         self: *Data,
191         bytes: []const u8,
192         comptime field: []const u8,
193         comptime parse: anytype,
194     ) !void {
195         var reader = Reader.init(bytes);
196         const length = try self.count(&reader);
197         for (0..length) |_| {
198             const start = reader.offset;
199             try parse(self, &reader);
200             _ = try self.append(field, bytes[start..reader.offset]);
201         }
202         try atEnd(reader);
203     }
204 
205     fn string(self: *const Data, reader: *Reader) ![]const u8 {
206         return self.strings.items[try reference(reader, self.strings.items.len)];
207     }
208 
209     fn readType(self: *Data, reader: *Reader) !void {
210         switch (try tag(bytecode.TypeKind, reader)) {
211             .builtin_scalar => {
212                 _ = try tag(bytecode.BuiltinScalarKind, reader);
213                 _ = try reader.readByte();
214             },
215             .dialect, .dialect_only => |kind| {
216                 _ = try reference(reader, self.dialects.items.len);
217                 if (kind == .dialect) _ = try self.string(reader);
218                 if (try boolean(reader)) _ = try self.string(reader);
219             },
220         }
221     }
222 
223     fn readAttribute(self: *Data, reader: *Reader) !void {
224         switch (try tag(bytecode.AttrKind, reader)) {
225             .integer => {
226                 _ = try reader.readSLEB();
227                 _ = try reader.readByte();
228                 _ = try boolean(reader);
229             },
230             .float_ => {
231                 _ = try reader.readU64();
232                 _ = try reader.readByte();
233             },
234             .bool_ => _ = try boolean(reader),
235             .string => _ = try reader.readBytesWithLen(),
236             .symbol_ref => {
237                 _ = try reader.readBytesWithLen();
238                 try self.readStringsList(reader);
239             },
240             .string_list => try self.readStringsList(reader),
241             .type_list => _ = try self.ids(reader, self.types.items.len),
242             .array => _ = try self.ids(reader, self.attributes.items.len),
243             .dialect => {
244                 _ = try reference(reader, self.dialects.items.len);
245                 _ = try self.string(reader);
246                 _ = try reader.readBytesWithLen();
247             },
248         }
249     }
250 
251     fn readStringsList(self: *const Data, reader: *Reader) !void {
252         const length = try self.count(reader);
253         for (0..length) |_| _ = try reader.readBytesWithLen();
254     }
255 
256     fn readLocation(self: *Data, reader: *Reader) !void {
257         const previous = self.locations.items.len;
258         switch (try tag(bytecode.LocationKind, reader)) {
259             .unknown => {},
260             .file => {
261                 _ = try self.string(reader);
262                 _ = try reader.readULEBToU32();
263                 _ = try reader.readULEBToU32();
264             },
265             .file_range => {
266                 _ = try self.string(reader);
267                 const start = try reader.readULEB();
268                 _ = try reader.readULEBToU32();
269                 _ = try reader.readULEBToU32();
270                 if (start > try reader.readULEB()) return error.InvalidTable;
271                 _ = try reader.readULEBToU32();
272                 _ = try reader.readULEBToU32();
273             },
274             .name => {
275                 _ = try self.string(reader);
276                 if (try boolean(reader)) _ = try reference(reader, previous);
277             },
278             .fused => _ = try self.ids(reader, previous),
279             .call_site => {
280                 _ = try reference(reader, previous);
281                 _ = try reference(reader, previous);
282             },
283         }
284     }
285 
286     fn readResources(self: *Data, bytes: []const u8) !void {
287         var reader = Reader.init(bytes);
288         const length = try self.count(&reader);
289         for (0..length) |_| {
290             var value: bytecode.Resource = undefined;
291             inline for (@typeInfo(bytecode.Resource).@"struct".field_names) |field| {
292                 @field(value, field) = try reader.readBytesWithLen();
293             }
294             _ = try self.append("resources", value);
295         }
296         try atEnd(reader);
297     }
298 
299     fn ids(self: *const Data, reader: *Reader, bound: usize) !Ids {
300         const length = try self.count(reader);
301         const start = reader.offset;
302         for (0..length) |_| _ = try reference(reader, bound);
303         return .{ .count = length, .bytes = reader.bytes[start..reader.offset] };
304     }
305 
306     fn readOperation(
307         self: *Data,
308         reader: *Reader,
309         parent: ?u32,
310         blocks: u32,
311         depth: u16,
312     ) anyerror!void {
313         if (depth > self.limits.depth) return error.ImageLimit;
314         const start = reader.offset;
315         const ordinal = try self.append("operations", @as(Operation, undefined));
316         const name = try self.string(reader);
317         const location = try reference(reader, self.locations.items.len);
318         const types = try self.ids(reader, self.types.items.len);
319         const operands = try self.ids(reader, self.values.items.len);
320         const attributes = try self.readNamedAttributes(reader);
321         const properties = if (try boolean(reader))
322             try reference(reader, self.attributes.items.len)
323         else
324             null;
325         const successors = try self.ids(reader, blocks);
326         const regions = try self.count(reader);
327         const results = try self.defineResults(types, ordinal, location);
328         for (0..regions) |position| try self.readRegion(reader, ordinal, @intCast(position), depth);
329         self.operations.items[ordinal] = .{
330             .name = name,
331             .location = location,
332             .parent_block = parent,
333             .results = results,
334             .operands = operands,
335             .attributes = attributes.bytes,
336             .attribute_count = attributes.count,
337             .properties = properties,
338             .successors = successors,
339             .region_count = regions,
340             .encoded = reader.bytes[start..reader.offset],
341         };
342     }
343 
344     fn readNamedAttributes(self: *Data, reader: *Reader) !Ids {
345         const length = try self.count(reader);
346         const start = reader.offset;
347         var previous: ?[]const u8 = null;
348         for (0..length) |_| {
349             const name = try self.string(reader);
350             if (previous) |prior| {
351                 if (!std.mem.lessThan(u8, prior, name)) return error.InvalidTable;
352             }
353             previous = name;
354             _ = try reference(reader, self.attributes.items.len);
355         }
356         return .{ .count = length, .bytes = reader.bytes[start..reader.offset] };
357     }
358 
359     fn defineResults(self: *Data, types: Ids, operation: u32, location: u32) !Range {
360         const range = Range{ .start = @intCast(self.values.items.len), .count = types.count };
361         var reader = Reader.init(types.bytes);
362         for (0..types.count) |position| {
363             _ = try self.append("values", Value{
364                 .owner = .{ .operation = operation },
365                 .position = @intCast(position),
366                 .type = try reader.readULEBToU32(),
367                 .location = location,
368             });
369         }
370         return range;
371     }
372 
373     fn readRegion(self: *Data, reader: *Reader, operation: u32, position: u32, depth: u16) !void {
374         const length = try self.count(reader);
375         const ordinal = try self.append("regions", Region{
376             .operation = operation,
377             .position = position,
378             .block_count = length,
379         });
380         for (0..length) |index| try self.readBlock(reader, ordinal, @intCast(index), length, depth);
381     }
382 
383     fn readBlock(
384         self: *Data,
385         reader: *Reader,
386         region: u32,
387         position: u32,
388         blocks: u32,
389         depth: u16,
390     ) !void {
391         const ordinal = try self.append("blocks", @as(Block, undefined));
392         const length = try self.count(reader);
393         const range = Range{ .start = @intCast(self.values.items.len), .count = length };
394         for (0..length) |index| {
395             _ = try self.append("values", Value{
396                 .owner = .{ .block = ordinal },
397                 .position = @intCast(index),
398                 .type = try reference(reader, self.types.items.len),
399                 .location = try reference(reader, self.locations.items.len),
400             });
401         }
402         const operations = try self.count(reader);
403         self.blocks.items[ordinal] = .{
404             .region = region,
405             .position = position,
406             .arguments = range,
407             .operation_count = operations,
408         };
409         if (operations > 0 and depth == std.math.maxInt(u16)) return error.ImageLimit;
410         for (0..operations) |_| try self.readOperation(reader, ordinal, blocks, depth + 1);
411     }
412 };
413 
414 fn reference(reader: *Reader, bound: usize) !u32 {
415     const ordinal = try reader.readULEBToU32();
416     if (ordinal >= bound) return error.InvalidTable;
417     return ordinal;
418 }
419 
420 fn boolean(reader: *Reader) !bool {
421     return switch (try reader.readByte()) {
422         0 => false,
423         1 => true,
424         else => error.InvalidTable,
425     };
426 }
427 
428 fn tag(comptime T: type, reader: *Reader) !T {
429     return std.enums.fromInt(T, try reader.readByte()) orelse error.InvalidTable;
430 }
431 
432 fn atEnd(reader: Reader) !void {
433     if (reader.offset != reader.bytes.len) return error.InvalidTable;
434 }
435 
436 fn data(index: *Index) *Data {
437     return @ptrCast(@alignCast(index));
438 }
439 
440 fn testBytes(allocator: std.mem.Allocator) ![]u8 {
441     const ir = @import("../core/root.zig");
442     var context = try ir.Context.init(allocator, ir.Context.Limits.testing);
443     defer context.deinit(allocator);
444     try context.allowUnregistered();
445     var state = ir.Operation.State.init("test.function", .getFile("test", 1, 1));
446     state.addRegion();
447     const root = try context.createOperation(state);
448     const block = try root.getRegion(0).?.addBlock();
449     const typ = try context.getDialectTypeFromNameWithKey("test.word", "32");
450     const argument = try block.addArgument(typ, .getFile("argument", 2, 3));
451     var child_state = ir.Operation.State.init("test.add", .getFile("operation", 4, 5));
452     child_state.addOperands(&.{ argument, argument });
453     child_state.addTypes(&.{typ});
454     child_state.addAttributes(&.{.{ .name = "answer", .value = try context.getI64Attr(42) }});
455     const child = try context.createOperation(child_state);
456     try block.addOperation(child);
457     const resources = [_]bytecode.Resource{.{
458         .namespace = "test",
459         .name = "data",
460         .type_id = "bytes",
461         .data = "resource",
462     }};
463     return bytecode.encodeModuleWithResources(allocator, root, &resources);
464 }
465 
466 const test_limits = Limits{ .bytes = 65536, .entities = 1024, .depth = 32 };
467 
468 test "bytecode image exposes immutable entities after originating Context destruction" {
469     const allocator = std.testing.allocator;
470     const bytes = try testBytes(allocator);
471     defer allocator.free(bytes);
472     const index = try Index.create(allocator, bytes, test_limits);
473     defer index.destroy();
474     const view = index.view();
475     try std.testing.expectEqual(2, view.operations.len);
476     try std.testing.expectEqual(1, view.blocks.len);
477     try std.testing.expectEqual(1, view.regions.len);
478     try std.testing.expectEqual(2, view.values.len);
479     try std.testing.expectEqualStrings("test.function", view.operations[0].name);
480     try std.testing.expectEqualStrings("test.add", view.operations[1].name);
481     try std.testing.expectEqual(0, view.operations[1].parent_block.?);
482     try std.testing.expectEqual(0, view.operations[1].operands.at(0).?);
483     try std.testing.expectEqual(0, view.operations[1].operands.at(1).?);
484     try std.testing.expectEqual(1, view.operations[1].results.start);
485     try std.testing.expectEqual(0, view.region(0, 0).?);
486     try std.testing.expectEqual(0, view.block(0, 0).?);
487     try std.testing.expectEqualStrings("resource", view.resources[0].data);
488     var location = Reader.init(view.locations[view.values[0].location]);
489     try std.testing.expectEqual(
490         bytecode.LocationKind.file,
491         try tag(bytecode.LocationKind, &location),
492     );
493     try std.testing.expectEqualStrings("argument", view.strings[try location.readULEBToU32()]);
494     try std.testing.expectEqual(2, try location.readULEBToU32());
495     try std.testing.expectEqual(3, try location.readULEBToU32());
496     var attribute = Reader.init(view.attributes[0]);
497     try std.testing.expectEqual(bytecode.AttrKind.integer, try tag(bytecode.AttrKind, &attribute));
498     try std.testing.expectEqual(42, try attribute.readSLEB());
499 }
500 
501 fn allocationScenario(allocator: std.mem.Allocator, bytes: []const u8) !void {
502     const index = try Index.create(allocator, bytes, test_limits);
503     defer index.destroy();
504     try std.testing.expectEqual(2, index.view().operations.len);
505 }
506 
507 test "bytecode image bounds malformed input and cleans every indexing allocation failure" {
508     const allocator = std.testing.allocator;
509     const bytes = try testBytes(allocator);
510     defer allocator.free(bytes);
511     try std.testing.checkAllAllocationFailures(allocator, allocationScenario, .{bytes});
512     var limits = test_limits;
513     limits.entities = 1;
514     try std.testing.expectError(error.ImageLimit, Index.create(allocator, bytes, limits));
515     limits = test_limits;
516     limits.depth = 0;
517     try std.testing.expectError(error.ImageLimit, Index.create(allocator, bytes, limits));
518     try std.testing.expectError(
519         error.InvalidSection,
520         Index.create(allocator, bytes[0 .. bytes.len - 1], test_limits),
521     );
522 }