lib/choir/src/core/dump.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

  1 const std = @import("std");
  2 const alloc_arena = @import("alloc_arena");
  3 
  4 const attribute = @import("attribute.zig");
  5 const Attribute = attribute.Attribute;
  6 const Block = @import("block.zig").Block;
  7 const NamedAttribute = attribute.NamedAttribute;
  8 const Operation = @import("operation/root.zig").Operation;
  9 const Type = @import("type.zig").Type;
 10 const Value = @import("value.zig").Value;
 11 
 12 pub fn operationAlloc(allocator: std.mem.Allocator, op: *Operation) ![]u8 {
 13     var out = std.Io.Writer.Allocating.init(allocator);
 14     defer out.deinit();
 15     try writeOperation(allocator, &out.writer, op);
 16     return try out.toOwnedSlice();
 17 }
 18 
 19 pub fn writeOperation(allocator: std.mem.Allocator, writer: *std.Io.Writer, op: *Operation) !void {
 20     var printer = Printer.init(allocator, writer);
 21     defer printer.deinit();
 22     try printer.writeOperation(op, 0);
 23 }
 24 
 25 const Printer = struct {
 26     writer: *std.Io.Writer,
 27     value_ids: std.AutoHashMap(*const Value, usize),
 28     next_value_id: usize = 0,
 29 
 30     fn init(allocator: std.mem.Allocator, writer: *std.Io.Writer) Printer {
 31         return .{
 32             .writer = writer,
 33             .value_ids = std.AutoHashMap(*const Value, usize).init(allocator),
 34         };
 35     }
 36 
 37     fn deinit(self: *Printer) void {
 38         self.value_ids.deinit();
 39     }
 40 
 41     fn writeOperation(self: *Printer, op: *Operation, indent: usize) anyerror!void {
 42         try self.writeIndent(indent);
 43 
 44         if (op.results.items.len > 0) {
 45             for (op.results.items, 0..) |*result, index| {
 46                 if (index > 0) try self.writer.writeAll(", ");
 47                 try self.writeValue(result);
 48             }
 49             try self.writer.writeAll(" = ");
 50         }
 51 
 52         try self.writer.writeAll(op.name.name);
 53         try self.writer.writeAll("(");
 54         for (op.operand_values, 0..) |operand, index| {
 55             if (index > 0) try self.writer.writeAll(", ");
 56             try self.writeValue(operand);
 57         }
 58         try self.writer.writeAll(")");
 59 
 60         if (op.successors.items.len > 0) {
 61             try self.writer.writeAll(" -> ");
 62             for (op.successors.items, 0..) |successor, index| {
 63                 if (index > 0) try self.writer.writeAll(", ");
 64                 try self.writer.print("^bb{d}", .{successor.id});
 65             }
 66         }
 67 
 68         if (try op.getPropertiesAsAttr()) |properties| {
 69             try self.writer.writeAll(" properties(");
 70             try self.writeAttribute(properties);
 71             try self.writer.writeAll(")");
 72         }
 73 
 74         try self.writeAttrs(op.getRawDictionaryAttrs());
 75 
 76         if (op.results.items.len > 0) {
 77             try self.writer.writeAll(" : ");
 78             for (op.results.items, 0..) |result, index| {
 79                 if (index > 0) try self.writer.writeAll(", ");
 80                 try self.writeType(result.type);
 81             }
 82         }
 83 
 84         if (op.regions.items.len == 0) {
 85             try self.writer.writeAll("\n");
 86             return;
 87         }
 88 
 89         for (op.regions.items) |*region| {
 90             try self.writer.writeAll(" ");
 91             try self.writer.writeAll("{\n");
 92             var block = region.blocks.head;
 93             while (block) |current| : (block = current.next) {
 94                 const block_indent = indent + 1;
 95                 try self.writeBlockHeader(current, block_indent);
 96                 var ops = current.getOperations();
 97                 while (ops.next()) |nested| {
 98                     try self.writeOperation(nested, block_indent + 1);
 99                 }
100             }
101             try self.writeIndent(indent);
102             try self.writer.writeAll("}");
103         }
104         try self.writer.writeAll("\n");
105     }
106 
107     fn writeBlockHeader(self: *Printer, block: *Block, indent: usize) anyerror!void {
108         try self.writeIndent(indent);
109         try self.writer.print("^bb{d}", .{block.id});
110         if (block.arguments.items.len > 0) {
111             try self.writer.writeAll("(");
112             for (block.arguments.items, 0..) |argument, index| {
113                 if (index > 0) try self.writer.writeAll(", ");
114                 try self.writeValue(argument);
115                 try self.writer.writeAll(": ");
116                 try self.writeType(argument.type);
117             }
118             try self.writer.writeAll(")");
119         }
120         try self.writer.writeAll(":\n");
121     }
122 
123     fn writeAttrs(self: *Printer, attrs: []const NamedAttribute) anyerror!void {
124         if (attrs.len == 0) return;
125         try self.writer.writeAll(" {");
126         for (attrs, 0..) |attr, index| {
127             if (index > 0) try self.writer.writeAll(", ");
128             try self.writeNamedAttr(attr);
129         }
130         try self.writer.writeAll("}");
131     }
132 
133     fn writeNamedAttr(self: *Printer, attr: NamedAttribute) anyerror!void {
134         try self.writer.print("{s} = ", .{attr.name});
135         try self.writeAttribute(attr.value);
136     }
137 
138     fn writeAttribute(self: *Printer, attr: Attribute) anyerror!void {
139         if (attr.cast(Attribute.IntegerAttr)) |integer| {
140             try self.writer.print("{d}:i{d}", .{ integer.value, integer.width });
141             return;
142         }
143         if (attr.cast(Attribute.FloatAttr)) |float| {
144             try self.writer.print("{d}:f{d}", .{ float.value, float.width });
145             return;
146         }
147         if (attr.cast(Attribute.BoolAttr)) |bool_attr| {
148             try self.writer.writeAll(if (bool_attr.value) "true" else "false");
149             return;
150         }
151         if (attr.cast(Attribute.StringAttr)) |string| {
152             try self.writeString(string.value);
153             return;
154         }
155         if (attr.cast(Attribute.SymbolRefAttr)) |symbol| {
156             try self.writer.print("@{s}", .{symbol.root_reference});
157             for (symbol.nested_references) |nested| {
158                 try self.writer.print("::@{s}", .{nested});
159             }
160             return;
161         }
162         if (attr.cast(Attribute.StringListAttr)) |list| {
163             try self.writer.writeAll("[");
164             for (list.values, 0..) |value, index| {
165                 if (index > 0) try self.writer.writeAll(", ");
166                 try self.writeString(value);
167             }
168             try self.writer.writeAll("]");
169             return;
170         }
171         if (attr.cast(Attribute.TypeListAttr)) |list| {
172             try self.writer.writeAll("[");
173             for (list.values, 0..) |typ, index| {
174                 if (index > 0) try self.writer.writeAll(", ");
175                 try self.writeType(typ);
176             }
177             try self.writer.writeAll("]");
178             return;
179         }
180         if (attr.cast(Attribute.ArrayAttr)) |array| {
181             try self.writer.writeAll("[");
182             for (array.values, 0..) |value, index| {
183                 if (index > 0) try self.writer.writeAll(", ");
184                 try self.writeAttribute(value);
185             }
186             try self.writer.writeAll("]");
187             return;
188         }
189         if (attr.cast(Attribute.DialectAttr)) |dialect| {
190             try self.writer.print("#attr<{s}>", .{attr.abstract.name});
191             if (dialect.payload.len > 0) {
192                 try self.writer.writeAll("(");
193                 try self.writeString(dialect.payload);
194                 try self.writer.writeAll(")");
195             }
196             return;
197         }
198         try self.writer.print("{f}", .{attr});
199     }
200 
201     fn writeString(self: *Printer, value: []const u8) anyerror!void {
202         try self.writer.writeByte('"');
203         for (value) |byte| {
204             switch (byte) {
205                 '"' => try self.writer.writeAll("\\\""),
206                 '\\' => try self.writer.writeAll("\\\\"),
207                 '\n' => try self.writer.writeAll("\\n"),
208                 '\r' => try self.writer.writeAll("\\r"),
209                 '\t' => try self.writer.writeAll("\\t"),
210                 else => try self.writer.writeByte(byte),
211             }
212         }
213         try self.writer.writeByte('"');
214     }
215 
216     fn writeType(self: *Printer, typ: Type) anyerror!void {
217         try self.writer.print("{f}", .{typ});
218     }
219 
220     fn writeValue(self: *Printer, value: *const Value) anyerror!void {
221         const id = try self.valueId(value);
222         try self.writer.print("%{d}", .{id});
223     }
224 
225     fn valueId(self: *Printer, value: *const Value) anyerror!usize {
226         if (self.value_ids.get(value)) |id| return id;
227         const id = self.next_value_id;
228         self.next_value_id += 1;
229         try self.value_ids.put(value, id);
230         return id;
231     }
232 
233     fn writeIndent(self: *Printer, indent: usize) anyerror!void {
234         for (0..indent) |_| try self.writer.writeAll("  ");
235     }
236 };
237 
238 test "Choir dump prints nested operations with stable local values" {
239     const testing = std.testing;
240     const dialects = @import("../dialects/root.zig");
241     const Context = @import("context/root.zig").Context;
242 
243     var arena = alloc_arena.Arena.init(testing.allocator);
244     defer arena.deinit();
245 
246     var ctx = try Context.init(arena.allocator(), Context.Limits.testing);
247     defer ctx.deinit(arena.allocator());
248 
249     const loc = @import("location.zig").Location.getUnknown();
250     const module = try dialects.BuiltinDialect.ModuleOp.create(&ctx, loc);
251     const module_block = module.getBodyBlock();
252     const index_type = try dialects.ArithDialect.getScalarType(&ctx, .index);
253     var func = try dialects.FuncDialect.FuncOp.create(&ctx, loc, "sample", &.{index_type}, &.{index_type});
254     try module_block.addOperation(func.op);
255 
256     const entry = func.getEntryBlock();
257     var constant = try dialects.ArithDialect.ConstantOp.createInt(&ctx, loc, index_type, 7);
258     try entry.addOperation(constant.op);
259     const ret = try dialects.FuncDialect.ReturnOp.create(&ctx, loc, &.{constant.getResult()});
260     try entry.addOperation(ret.op);
261 
262     const text = try operationAlloc(testing.allocator, module.op);
263     defer testing.allocator.free(text);
264 
265     try testing.expect(std.mem.indexOf(u8, text, "builtin.module()") != null);
266     try testing.expect(std.mem.indexOf(u8, text, "func.func()") != null);
267     try testing.expect(std.mem.indexOf(u8, text, "^bb0(%1: !arith.index)") != null);
268     try testing.expect(std.mem.indexOf(u8, text, "%2 = arith.constant()") != null);
269     try testing.expect(std.mem.indexOf(u8, text, "func.return(%2)") != null);
270     try testing.expect(std.mem.indexOf(u8, text, "sym_name = #attr<func.sym_name>(\"sample\")") != null);
271 }