lib/choir/src/backends/wasm/emission/function/encoder.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const wasm = @import("../../root.zig");
3 const ir = @import("../../../../core/root.zig");
4 const dialects = @import("../../../../dialects/root.zig");
5 const emission = @import("../root.zig");
6
7 const FuncDialect = dialects.FuncDialect;
8 const ScfDialect = dialects.ScfDialect;
9 const instruction = emission.function.instruction;
10 const arithmetic = instruction.arithmetic;
11 const types = emission.types;
12
13 const Terminator = union(enum) {
14 function_return,
15 if_yield: ScfDialect.IfOp,
16 for_yield: ScfDialect.ForOp,
17 };
18
19 const BlockFrame = struct {
20 next: ?*ir.Operation,
21 terminator: Terminator,
22 };
23
24 const IfFrame = struct {
25 operation: ScfDialect.IfOp,
26 then_complete: bool,
27 };
28
29 pub const Frame = union(enum) {
30 block: BlockFrame,
31 if_control: IfFrame,
32 for_control: ScfDialect.ForOp,
33 };
34
35 pub fn writeContent(
36 out: anytype,
37 plan: *const emission.Plan,
38 func: FuncDialect.FuncOp,
39 function_ordinal: usize,
40 local_count: usize,
41 frames: []Frame,
42 ) emission.Error!void {
43 try writeLocalDeclarations(out, plan, function_ordinal, local_count);
44 var encoder = Encoder(@TypeOf(out)){
45 .out = out,
46 .func = func,
47 .frames = frames,
48 .instructions = .{
49 .out = out,
50 .plan = plan,
51 .function_ordinal = function_ordinal,
52 },
53 };
54 try encoder.writeFunctionBody();
55 try out.writeByte(0x0b);
56 }
57
58 fn writeLocalDeclarations(
59 out: anytype,
60 plan: *const emission.Plan,
61 function_ordinal: usize,
62 local_count: usize,
63 ) emission.Error!void {
64 try wasm.binary.writeUleb(out, local_count);
65 const locals = plan.localPlans(function_ordinal);
66 if (locals.len != local_count) return error.InputChanged;
67 for (locals) |local| {
68 try wasm.binary.writeUleb(out, @as(u32, 1));
69 try wasm.binary.writeValueType(out, local.value_type);
70 }
71 }
72
73 fn Encoder(comptime Out: type) type {
74 return struct {
75 out: Out,
76 func: FuncDialect.FuncOp,
77 frames: []Frame,
78 instructions: instruction.Writer(Out),
79 frame_len: usize = 0,
80
81 const Self = @This();
82
83 fn writeFunctionBody(self: *Self) emission.Error!void {
84 try self.pushBlock(self.func.getEntryBlock(), .function_return);
85 while (self.frame_len > 0) {
86 switch (self.frames[self.frame_len - 1]) {
87 .block => try self.stepBlock(),
88 .if_control => try self.finishIf(),
89 .for_control => try self.finishFor(),
90 }
91 }
92 }
93
94 fn stepBlock(self: *Self) emission.Error!void {
95 const block = &self.frames[self.frame_len - 1].block;
96 const operation = block.next orelse {
97 self.frame_len -= 1;
98 return;
99 };
100 block.next = operation.next_op;
101 const name = operation.name.name;
102 if (std.mem.eql(u8, name, FuncDialect.ReturnOp.operation_name)) {
103 try self.writeFunctionReturn(FuncDialect.ReturnOp{ .op = operation });
104 self.frame_len -= 1;
105 return;
106 }
107 if (std.mem.eql(u8, name, ScfDialect.YieldOp.operation_name)) {
108 try self.writeYield(
109 ScfDialect.YieldOp{ .op = operation },
110 block.terminator,
111 );
112 self.frame_len -= 1;
113 return;
114 }
115 if (std.mem.eql(u8, name, ScfDialect.IfOp.operation_name)) {
116 try self.beginIf(ScfDialect.IfOp{ .op = operation });
117 return;
118 }
119 if (std.mem.eql(u8, name, ScfDialect.ForOp.operation_name)) {
120 try self.beginFor(ScfDialect.ForOp{ .op = operation });
121 return;
122 }
123 try self.instructions.writeOperation(operation);
124 }
125
126 fn beginIf(self: *Self, if_op: ScfDialect.IfOp) emission.Error!void {
127 try self.instructions.writeValue(if_op.getCondition());
128 try self.out.writeByte(0x04);
129 try self.out.writeByte(wasm.binary.empty_block_type);
130 try self.push(.{ .if_control = .{
131 .operation = if_op,
132 .then_complete = false,
133 } });
134 try self.pushBlock(if_op.getThenBlock(), .{ .if_yield = if_op });
135 }
136
137 fn finishIf(self: *Self) emission.Error!void {
138 const state = &self.frames[self.frame_len - 1].if_control;
139 if (!state.then_complete) {
140 state.then_complete = true;
141 if (state.operation.getElseBlock()) |else_block| {
142 try self.out.writeByte(0x05);
143 try self.pushBlock(else_block, .{ .if_yield = state.operation });
144 return;
145 }
146 }
147 try self.out.writeByte(0x0b);
148 self.frame_len -= 1;
149 }
150
151 fn beginFor(self: *Self, for_op: ScfDialect.ForOp) emission.Error!void {
152 const induction_local = try self.instructions.localFor(for_op.getInductionVar());
153 try self.instructions.writeValue(for_op.getLowerBound());
154 try self.instructions.writeLocalSet(induction_local);
155
156 const iter_args = for_op.getIterArgs();
157 const init_args = for_op.getInitArgs();
158 if (iter_args.len != init_args.len) return error.CodeGenFailed;
159 for (iter_args, init_args) |iter_arg, init_arg| {
160 try self.instructions.writeValue(init_arg);
161 try self.instructions.writeLocalSet(try self.instructions.localFor(iter_arg));
162 }
163
164 try self.out.writeByte(0x02);
165 try self.out.writeByte(wasm.binary.empty_block_type);
166 try self.out.writeByte(0x03);
167 try self.out.writeByte(wasm.binary.empty_block_type);
168 try self.instructions.writeValue(for_op.getInductionVar());
169 try self.instructions.writeValue(for_op.getUpperBound());
170 const scalar = types.scalarKindForType(for_op.getInductionVar().type) orelse
171 return error.CodeGenFailed;
172 try self.out.writeByte(try arithmetic.opcode.comparison(.lt, scalar));
173 try self.out.writeByte(0x45);
174 try self.out.writeByte(0x0d);
175 try wasm.binary.writeUleb(self.out, @as(u32, 1));
176 try self.push(.{ .for_control = for_op });
177 try self.pushBlock(for_op.getBodyBlock(), .{ .for_yield = for_op });
178 }
179
180 fn finishFor(self: *Self) emission.Error!void {
181 const for_op = self.frames[self.frame_len - 1].for_control;
182 try self.out.writeByte(0x0b);
183 try self.out.writeByte(0x0b);
184 const iter_args = for_op.getIterArgs();
185 if (iter_args.len != for_op.op.getNumResults()) return error.CodeGenFailed;
186 for (iter_args, 0..) |iter_arg, index| {
187 const result = for_op.getResult(index) orelse return error.CodeGenFailed;
188 try self.instructions.writeLocalGet(try self.instructions.localFor(iter_arg));
189 try self.instructions.writeLocalSet(try self.instructions.localFor(result));
190 }
191 self.frame_len -= 1;
192 }
193
194 fn writeFunctionReturn(
195 self: *Self,
196 operation: FuncDialect.ReturnOp,
197 ) emission.Error!void {
198 for (operation.getOperands()) |operand| try self.instructions.writeValue(operand);
199 try self.out.writeByte(0x0f);
200 }
201
202 fn writeYield(
203 self: *Self,
204 operation: ScfDialect.YieldOp,
205 terminator: Terminator,
206 ) emission.Error!void {
207 switch (terminator) {
208 .function_return => return error.CodeGenFailed,
209 .if_yield => |if_op| {
210 const operands = operation.getOperands();
211 if (operands.len != if_op.getNumResults()) return error.CodeGenFailed;
212 for (operands, 0..) |operand, index| {
213 const result = if_op.getResult(index) orelse return error.CodeGenFailed;
214 try self.instructions.writeValue(operand);
215 try self.instructions.writeLocalSet(try self.instructions.localFor(result));
216 }
217 },
218 .for_yield => |for_op| {
219 const operands = operation.getOperands();
220 const iter_args = for_op.getIterArgs();
221 if (operands.len != iter_args.len or
222 operands.len != for_op.op.getNumResults())
223 {
224 return error.CodeGenFailed;
225 }
226 for (operands, 0..) |operand, index| {
227 const result = for_op.getResult(index) orelse return error.CodeGenFailed;
228 try self.instructions.writeValue(operand);
229 try self.instructions.writeLocalSet(try self.instructions.localFor(result));
230 }
231 for (iter_args, 0..) |iter_arg, index| {
232 const result = for_op.getResult(index) orelse return error.CodeGenFailed;
233 try self.instructions.writeLocalGet(try self.instructions.localFor(result));
234 try self.instructions.writeLocalSet(try self.instructions.localFor(iter_arg));
235 }
236 try self.instructions.writeLocalGet(
237 try self.instructions.localFor(for_op.getInductionVar()),
238 );
239 try self.instructions.writeValue(for_op.getStep());
240 const scalar = types.scalarKindForType(for_op.getStep().type) orelse
241 return error.CodeGenFailed;
242 try self.out.writeByte(try arithmetic.opcode.binary(.add, scalar));
243 try self.instructions.writeLocalSet(
244 try self.instructions.localFor(for_op.getInductionVar()),
245 );
246 try self.out.writeByte(0x0c);
247 try wasm.binary.writeUleb(self.out, @as(u32, 0));
248 },
249 }
250 }
251
252 fn pushBlock(
253 self: *Self,
254 block: *ir.Block,
255 terminator: Terminator,
256 ) emission.Error!void {
257 const first = if (block.operations.head) |node|
258 @as(*ir.Operation, @ptrCast(@alignCast(node)))
259 else
260 null;
261 try self.push(.{ .block = .{ .next = first, .terminator = terminator } });
262 }
263
264 fn push(self: *Self, frame: Frame) emission.Error!void {
265 if (self.frame_len == self.frames.len) return error.NestingLimitExceeded;
266 self.frames[self.frame_len] = frame;
267 self.frame_len += 1;
268 }
269 };
270 }