x86 assembly dsl in zig

asm_dsl.zig

comptime {
    const builtin = @import("builtin");
    std.debug.assert(builtin.target.cpu.arch == .x86);
}

pub const Instruction = union(enum) {
    mov_: Mov,
    int_: Int,
    ret_: Ret,
    lea_: Lea,

    pub const Mov = struct {
        dest: Register,
        src: Register_Or_Immediate,

        fn encode(instruction: Mov) [5]u8 {
            var bytes: [5]u8 = undefined;

            switch (instruction.src) {
                .immediate => |imm| {
                    bytes[0] = switch (instruction.dest) {
                        .eax => '\xB8',
                        .ecx => '\xB9',
                        .edx => '\xBA',
                        .ebx => '\xBB',
                        .esp => '\xBC',
                        .ebp => '\xBD',
                        .esi => '\xBE',
                        .edi => '\xBF',
                    };
                    bytes[1..][0..4].* = @bitCast(imm);
                },
                else => @panic("Idk how to encode this"),
            }

            return bytes;
        }
    };

    pub const Lea = struct {
        dest: Register,
        src: Immediate,

        pub fn encode(instruction: Lea) [6]u8 {
            var bytes: [6]u8 = undefined;
            bytes[0] = '\x8D';
            // hardcode modifier and rm because idk what I'm doing
            bytes[1] = (0b00 << 6) |
                @as(u8, @intCast((@as(u32, @intFromEnum(instruction.dest)) << 3))) |
                @intFromEnum(Register.ebp);
            bytes[2..][0..4].* = @bitCast(instruction.src);
            return bytes;
        }
    };

    pub const Int = enum(u8) {
        _,

        pub fn encode(instruction: Int) [2]u8 {
            return .{ '\xCD', @intFromEnum(instruction) };
        }
    };

    pub const Ret = struct {
        pub fn encode(instruction: Ret) [1]u8 {
            _ = instruction;
            return .{'\xC3'};
        }
    };

    pub fn lea(dest: Register, src: Immediate) Instruction {
        return .{ .lea_ = .{ .dest = dest, .src = src } };
    }

    pub fn ret() Instruction {
        return .{ .ret_ = .{} };
    }

    pub fn int(imm: u8) Instruction {
        return .{ .int_ = @enumFromInt(imm) };
    }

    pub fn mov(dest: Register, src: Register_Or_Immediate) Instruction {
        return .{ .mov_ = .{ .dest = dest, .src = src } };
    }

    pub fn encode(instruction: Instruction, writer: anytype) !void {
        switch (instruction) {
            inline else => |inst| try writer.writeAll(&inst.encode()),
        }
    }
};

pub const Register_Or_Immediate = union(enum) {
    register: Register,
    immediate: Immediate,

    pub const eax: Register_Or_Immediate = .{ .register = .eax };
    pub const ecx: Register_Or_Immediate = .{ .register = .ecx };
    pub const edx: Register_Or_Immediate = .{ .register = .edx };
    pub const ebx: Register_Or_Immediate = .{ .register = .ebx };
    pub const esp: Register_Or_Immediate = .{ .register = .esp };
    pub const ebp: Register_Or_Immediate = .{ .register = .ebp };
    pub const esi: Register_Or_Immediate = .{ .register = .esi };
    pub const edi: Register_Or_Immediate = .{ .register = .edi };

    pub fn imm(v: Immediate) Register_Or_Immediate {
        return .{ .immediate = v };
    }
};

pub const Immediate = u32;

pub const Register = enum {
    eax,
    ecx,
    edx,
    ebx,
    esp,
    ebp,
    esi,
    edi,
};

pub fn encodeMany(code: []const Instruction, writer: anytype) !void {
    for (code) |inst| {
        try inst.encode(writer);
    }
}

pub fn exec(code: []const Instruction) !void {
    const memory = try std.posix.mmap(
        null,
        code.len * 16,
        std.posix.PROT.WRITE,
        .{
            .TYPE = .PRIVATE,
            .ANONYMOUS = true,
        },
        -1,
        0,
    );
    defer std.posix.munmap(memory);

    var stream = std.io.fixedBufferStream(memory);
    try encodeMany(code, stream.writer());
    try std.posix.mprotect(memory, std.posix.PROT.READ | std.posix.PROT.EXEC);
    const func: *const fn () callconv(.c) void = @ptrCast(memory);
    func();
}

const std = @import("std");

example.zig

pub fn main() !void {
    const msg: []const u8 = "Hello, World!\n";
    const code: []const asm_dsl.Instruction = &.{
        .mov(.eax, .imm(4)),
        .mov(.ebx, .imm(1)),
        .lea(.ecx, @intFromPtr(msg.ptr)),
        .mov(.edx, .imm(msg.len)),
        .int(0x80),
        .ret(),
    };
    try asm_dsl.exec(code);
}

const std = @import("std");
const asm_dsl = @import("asm_dsl.zig");