tiny.choir.backends.wasm.emission.plan.planning

1const std = @import("std");2const ir = @import("../../../../core/root.zig");3const dialects = @import("../../../../dialects/root.zig");4const emission = @import("../root.zig");5const plan = @import("root.zig");67const BuiltinDialect = dialects.BuiltinDialect;8const FuncDialect = dialects.FuncDialect;9const MemrefDialect = dialects.MemrefDialect;10const types = emission.types;1112pub 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) or24        !plan.index.storageIsValid(values.len, value_index))25    {26        return error.InputChanged;27    }2829    plan.index.clear(function_index);30    plan.index.clear(value_index);3132    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;4950    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 &.{}) or60            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));8283        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    }112113    if (function_cursor != functions.len or114        value_cursor != values.len or115        local_count != facts.local_count or116        import_index != facts.import_count or117        definition_index != facts.import_count + facts.definition_count or118        nesting_depth != facts.nesting_depth or119        needs_memory != facts.needs_memory)120    {121        return error.InputChanged;122    }123124    return result;125}126127fn 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) catch149                return error.CapacityOverflow;150        }151    }152153    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) catch163                return error.CapacityOverflow;164        }165        if (std.mem.eql(u8, op.name.name, MemrefDialect.LoadOp.operation_name) or166            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}187188fn 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}205206fn typesNeedMemory(value_types: []const ir.Type) bool {207    for (value_types) |typ| {208        if (types.typeIsMemref(typ)) return true;209    }210    return false;211}212213fn u32Index(value: usize) emission.Error!u32 {214    return std.math.cast(u32, value) orelse error.CapacityOverflow;215}216217test "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);222223    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);247248    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}

1pub const planning = @import("planning.zig");