lib/choir/src/backends/wasm/emission/plan/planning.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const ir = @import("../../../../core/root.zig");
3 const dialects = @import("../../../../dialects/root.zig");
4 const emission = @import("../root.zig");
5 const plan = @import("root.zig");
6
7 const BuiltinDialect = dialects.BuiltinDialect;
8 const FuncDialect = dialects.FuncDialect;
9 const MemrefDialect = dialects.MemrefDialect;
10 const types = emission.types;
11
12 pub fn fill(
13 module: *ir.Operation,
14 functions: []plan.FunctionPlan,
15 function_index: []u32,
16 values: []plan.ValuePlan,
17 value_index: []u32,
18 facts: emission.Facts,
19 ) emission.Error!plan.Plan {
20 if (!std.mem.eql(u8, module.name.name, BuiltinDialect.ModuleOp.operation_name)) {
21 return error.CodeGenFailed;
22 }
23 if (!plan.index.storageIsValid(functions.len, function_index) or
24 !plan.index.storageIsValid(values.len, value_index))
25 {
26 return error.InputChanged;
27 }
28
29 plan.index.clear(function_index);
30 plan.index.clear(value_index);
31
32 const result = plan.Plan{
33 .functions = functions,
34 .function_index = function_index,
35 .values = values,
36 .value_index = value_index,
37 .import_count = facts.import_count,
38 .definition_count = facts.definition_count,
39 .needs_memory = facts.needs_memory,
40 .sections = .{},
41 };
42 var function_cursor: usize = 0;
43 var value_cursor: usize = 0;
44 var local_count: usize = 0;
45 var import_index: usize = 0;
46 var definition_index = facts.import_count;
47 var nesting_depth: usize = 0;
48 var needs_memory = false;
49
50 const module_op = BuiltinDialect.ModuleOp{ .op = module };
51 var ops = module_op.getBodyBlock().getOperations();
52 while (ops.next()) |op| {
53 if (function_cursor == functions.len) return error.InputChanged;
54 if (!std.mem.eql(u8, op.name.name, FuncDialect.FuncOp.operation_name)) {
55 return error.CodeGenFailed;
56 }
57 const func = FuncDialect.FuncOp{ .op = op };
58 const name = func.getName() orelse return error.CodeGenFailed;
59 needs_memory = needs_memory or typesNeedMemory(func.getInputTypes() orelse &.{}) or
60 typesNeedMemory(func.getResultTypes());
61 const imported = func.isDeclaration();
62 const wasm_index = if (imported) index_value: {
63 const value = try u32Index(import_index);
64 import_index += 1;
65 break :index_value value;
66 } else index_value: {
67 const value = try u32Index(definition_index);
68 definition_index += 1;
69 break :index_value value;
70 };
71 const function_plan = &functions[function_cursor];
72 function_plan.* = .{
73 .name = name,
74 .value_start = try u32Index(value_cursor),
75 .value_count = 0,
76 .param_count = 0,
77 .local_count = 0,
78 .function_index = wasm_index,
79 .body_bytes = 0,
80 };
81 try plan.index.insertFunction(function_index, functions, try u32Index(function_cursor));
82
83 if (!imported) {
84 const args = func.getArguments();
85 function_plan.param_count = try u32Index(args.len);
86 for (args, 0..) |argument, local_index| {
87 try addValue(
88 result,
89 &value_cursor,
90 argument,
91 try u32Index(local_index),
92 &needs_memory,
93 );
94 }
95 var next_local = function_plan.param_count;
96 try addBlockValues(
97 result,
98 &value_cursor,
99 func.getEntryBlock(),
100 true,
101 &next_local,
102 &local_count,
103 &nesting_depth,
104 &needs_memory,
105 0,
106 );
107 function_plan.local_count = next_local - function_plan.param_count;
108 function_plan.value_count = next_local;
109 }
110 function_cursor += 1;
111 }
112
113 if (function_cursor != functions.len or
114 value_cursor != values.len or
115 local_count != facts.local_count or
116 import_index != facts.import_count or
117 definition_index != facts.import_count + facts.definition_count or
118 nesting_depth != facts.nesting_depth or
119 needs_memory != facts.needs_memory)
120 {
121 return error.InputChanged;
122 }
123
124 return result;
125 }
126
127 fn addBlockValues(
128 result: plan.Plan,
129 cursor: *usize,
130 block: *ir.Block,
131 skip_arguments: bool,
132 next_local: *u32,
133 local_count: *usize,
134 nesting_depth: *usize,
135 needs_memory: *bool,
136 depth: usize,
137 ) emission.Error!void {
138 if (depth > emission.max_nesting_depth) return error.NestingLimitExceeded;
139 nesting_depth.* = @max(nesting_depth.*, depth);
140 if (!skip_arguments) {
141 local_count.* = std.math.add(
142 usize,
143 local_count.*,
144 block.arguments.items.len,
145 ) catch return error.CapacityOverflow;
146 for (block.arguments.items) |argument| {
147 try addValue(result, cursor, argument, next_local.*, needs_memory);
148 next_local.* = std.math.add(u32, next_local.*, 1) catch
149 return error.CapacityOverflow;
150 }
151 }
152
153 var ops = block.getOperations();
154 while (ops.next()) |op| {
155 local_count.* = std.math.add(
156 usize,
157 local_count.*,
158 op.results.items.len,
159 ) catch return error.CapacityOverflow;
160 for (op.results.items) |*value| {
161 try addValue(result, cursor, value, next_local.*, needs_memory);
162 next_local.* = std.math.add(u32, next_local.*, 1) catch
163 return error.CapacityOverflow;
164 }
165 if (std.mem.eql(u8, op.name.name, MemrefDialect.LoadOp.operation_name) or
166 std.mem.eql(u8, op.name.name, MemrefDialect.StoreOp.operation_name))
167 {
168 needs_memory.* = true;
169 }
170 for (op.regions.items) |*region| {
171 if (region.blocks.size != 1) return error.InputChanged;
172 const entry = region.getEntryBlock() orelse return error.CodeGenFailed;
173 try addBlockValues(
174 result,
175 cursor,
176 entry,
177 false,
178 next_local,
179 local_count,
180 nesting_depth,
181 needs_memory,
182 depth + 1,
183 );
184 }
185 }
186 }
187
188 fn addValue(
189 result: plan.Plan,
190 cursor: *usize,
191 value: *ir.Value,
192 local_index: u32,
193 needs_memory: *bool,
194 ) emission.Error!void {
195 if (cursor.* == result.values.len) return error.InputChanged;
196 result.values[cursor.*] = .{
197 .value = value,
198 .local_index = local_index,
199 .value_type = try types.wasmTypeForType(value.type),
200 };
201 if (types.typeIsMemref(value.type)) needs_memory.* = true;
202 try plan.index.insertValue(result.value_index, result.values, try u32Index(cursor.*));
203 cursor.* += 1;
204 }
205
206 fn typesNeedMemory(value_types: []const ir.Type) bool {
207 for (value_types) |typ| {
208 if (types.typeIsMemref(typ)) return true;
209 }
210 return false;
211 }
212
213 fn u32Index(value: usize) emission.Error!u32 {
214 return std.math.cast(u32, value) orelse error.CapacityOverflow;
215 }
216
217 test "WASM planning rejects region shape drift" {
218 const allocator = std.testing.allocator;
219 var context = try ir.Context.init(allocator, ir.Context.Limits.testing);
220 defer context.deinit(allocator);
221 try dialects.registerAllDialects(&context);
222
223 const location = ir.Location.getUnknown();
224 const i32_type = try dialects.ArithDialect.getI32Type(&context);
225 const source_module = try BuiltinDialect.ModuleOp.create(&context, location);
226 var function = try FuncDialect.FuncOp.create(
227 &context,
228 location,
229 "nested",
230 &.{i32_type},
231 &.{},
232 );
233 try source_module.getBodyBlock().addOperation(function.op);
234 var if_op = try dialects.ScfDialect.IfOp.create(
235 &context,
236 location,
237 function.getArgument(0),
238 &.{},
239 );
240 try function.getEntryBlock().addOperation(if_op.op);
241 const then_yield = try dialects.ScfDialect.YieldOp.create(&context, location, &.{});
242 try if_op.getThenBlock().addOperation(then_yield.op);
243 const else_yield = try dialects.ScfDialect.YieldOp.create(&context, location, &.{});
244 try if_op.getElseBlock().?.addOperation(else_yield.op);
245 const return_op = try FuncDialect.ReturnOp.create(&context, location, &.{});
246 try function.getEntryBlock().addOperation(return_op.op);
247
248 const limits = try emission.Limits.inspect(source_module.op, .{});
249 _ = try if_op.getThenRegion().addBlock();
250 try std.testing.expectError(
251 error.InputChanged,
252 emission.ModuleEmitter.init(allocator, limits),
253 );
254 }