lib/choir/src/backends/wasm/emission/inspection.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 
  6 const BuiltinDialect = dialects.BuiltinDialect;
  7 const FuncDialect = dialects.FuncDialect;
  8 const MemrefDialect = dialects.MemrefDialect;
  9 const types = emission.types;
 10 
 11 pub fn inspect(module: *ir.Operation, options: emission.Options) emission.Error!emission.Facts {
 12     if (!std.mem.eql(u8, module.name.name, BuiltinDialect.ModuleOp.operation_name)) {
 13         return error.CodeGenFailed;
 14     }
 15 
 16     var facts = emission.Facts{
 17         .function_count = 0,
 18         .import_count = 0,
 19         .definition_count = 0,
 20         .value_count = 0,
 21         .local_count = 0,
 22         .nesting_depth = 0,
 23         .needs_memory = false,
 24     };
 25 
 26     const module_op = BuiltinDialect.ModuleOp{ .op = module };
 27     var ops = module_op.getBodyBlock().getOperations();
 28     while (ops.next()) |operation| {
 29         if (!std.mem.eql(u8, operation.name.name, FuncDialect.FuncOp.operation_name)) {
 30             return error.CodeGenFailed;
 31         }
 32         const func = FuncDialect.FuncOp{ .op = operation };
 33         _ = func.getName() orelse return error.CodeGenFailed;
 34         facts.function_count = try addCount(facts.function_count, 1);
 35         const inputs = func.getInputTypes() orelse &.{};
 36         try inspectTypes(inputs, &facts.needs_memory);
 37         try inspectTypes(func.getResultTypes(), &facts.needs_memory);
 38         if (func.isDeclaration()) {
 39             facts.import_count = try addCount(facts.import_count, 1);
 40             continue;
 41         }
 42 
 43         facts.definition_count = try addCount(facts.definition_count, 1);
 44         const args = func.getArguments();
 45         facts.value_count = try addCount(facts.value_count, args.len);
 46         for (args) |argument| _ = try types.wasmTypeForType(argument.type);
 47         try inspectBlock(func.getEntryBlock(), true, 0, &facts);
 48     }
 49 
 50     try requireU32(facts.function_count);
 51     try requireU32(facts.import_count);
 52     try requireU32(facts.definition_count);
 53     try requireU32(facts.value_count);
 54     try requireU32(facts.local_count);
 55     if (options.entry != null) {
 56         var symbols = try emission.symbols.Iterator.init(module, options);
 57         var found = false;
 58         while (try symbols.next()) |symbol| {
 59             if (symbol.role == .provided) {
 60                 found = true;
 61                 break;
 62             }
 63         }
 64         if (!found) return error.FunctionNotFound;
 65     }
 66     return facts;
 67 }
 68 
 69 fn inspectBlock(
 70     block: *ir.Block,
 71     skip_arguments: bool,
 72     depth: usize,
 73     facts: *emission.Facts,
 74 ) emission.Error!void {
 75     if (depth > emission.max_nesting_depth) return error.NestingLimitExceeded;
 76     facts.nesting_depth = @max(facts.nesting_depth, depth);
 77     if (!skip_arguments) {
 78         facts.value_count = try addCount(facts.value_count, block.arguments.items.len);
 79         facts.local_count = try addCount(facts.local_count, block.arguments.items.len);
 80         for (block.arguments.items) |argument| _ = try types.wasmTypeForType(argument.type);
 81     }
 82 
 83     var ops = block.getOperations();
 84     while (ops.next()) |operation| {
 85         facts.value_count = try addCount(facts.value_count, operation.results.items.len);
 86         facts.local_count = try addCount(facts.local_count, operation.results.items.len);
 87         for (operation.results.items) |result| _ = try types.wasmTypeForType(result.type);
 88         if (std.mem.eql(u8, operation.name.name, MemrefDialect.LoadOp.operation_name) or
 89             std.mem.eql(u8, operation.name.name, MemrefDialect.StoreOp.operation_name))
 90         {
 91             facts.needs_memory = true;
 92         }
 93         for (operation.regions.items) |*region| {
 94             if (region.blocks.size != 1) return error.CodeGenFailed;
 95             const entry = region.getEntryBlock() orelse return error.CodeGenFailed;
 96             try inspectBlock(entry, false, depth + 1, facts);
 97         }
 98     }
 99 }
100 
101 fn inspectTypes(value_types: []const ir.Type, needs_memory: *bool) emission.Error!void {
102     try requireU32(value_types.len);
103     for (value_types) |typ| {
104         _ = try types.wasmTypeForType(typ);
105         if (types.typeIsMemref(typ)) needs_memory.* = true;
106     }
107 }
108 
109 fn addCount(lhs: usize, rhs: usize) emission.Error!usize {
110     return std.math.add(usize, lhs, rhs) catch error.CapacityOverflow;
111 }
112 
113 fn requireU32(value: usize) emission.Error!void {
114     if (value > std.math.maxInt(u32)) return error.CapacityOverflow;
115 }