tiny.choir.backends.wasm.emission.function.encoder

1const std = @import("std");2const wasm = @import("../../root.zig");3const ir = @import("../../../../core/root.zig");4const dialects = @import("../../../../dialects/root.zig");5const emission = @import("../root.zig");67const FuncDialect = dialects.FuncDialect;8const ScfDialect = dialects.ScfDialect;9const instruction = emission.function.instruction;10const arithmetic = instruction.arithmetic;11const types = emission.types;1213const Terminator = union(enum) {14    function_return,15    if_yield: ScfDialect.IfOp,16    for_yield: ScfDialect.ForOp,17};1819const BlockFrame = struct {20    next: ?*ir.Operation,21    terminator: Terminator,22};2324const IfFrame = struct {25    operation: ScfDialect.IfOp,26    then_complete: bool,27};2829pub const Frame = union(enum) {30    block: BlockFrame,31    if_control: IfFrame,32    for_control: ScfDialect.ForOp,33};3435pub 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}5758fn 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}7273fn 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,8081        const Self = @This();8283        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        }9394        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        }125126        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        }136137        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        }150151        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);155156            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            }163164            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) orelse171                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        }179180        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        }193194        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        }201202        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 or222                        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) orelse241                        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        }251252        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            else260                null;261            try self.push(.{ .block = .{ .next = first, .terminator = terminator } });262        }263264        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}

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