const std = @import("std");
const elf = std.elf;
const Allocator = std.mem.Allocator;
const root = @import("root");
const compile = root.compile;

const Register = enum(u5) {
    // zig fmt: off
    zero, ra, sp, gp, tp,
    t0, t1, t2,
    s0, s1,
    a0, a1, a2, a3, a4, a5, a6, a7,
    s2, s3, s4, s5, s6, s7, s8, s9, s10, s11,
    t3, t4, t5, t6,
    // zig fmt: on

    const fp = .s0;

    fn x(number: u5) @This() {
        return @enumFromInt(number);
    }
};

/// TODO: Panic on calls that become reserved hints, e.g. addw with rd=x0
const Instruction = packed union {
    r: R,
    i: I,
    s: S,
    b: B,
    u: U,
    j: J,

    const Opcode = u7;

    const R = packed struct(u32) {
        opcode: Opcode,
        rd: Register,
        funct3: u3,
        rs1: Register,
        rs2: Register,
        funct7: u7,

        fn init(opcode: Opcode, rd: Register, funct3: u3, rs1: Register, rs2: Register, funct7: u7) Self {
            return .{ .r = .{
                .opcode = opcode,
                .rd = rd,
                .funct3 = funct3,
                .rs1 = rs1,
                .rs2 = rs2,
                .funct7 = funct7,
            } };
        }
    };
    /// rd = rs1 + rs2
    fn add(rd: Register, rs1: Register, rs2: Register) Self {
        return R.init(0b0110011, rd, 0, rs1, rs2, 0);
    }
    /// rd = rs1 + rs2
    fn addw(rd: Register, rs1: Register, rs2: Register) Self {
        return R.init(0b0111011, rd, 0, rs1, rs2, 0);
    }
    /// rd = rs1 - rs2
    fn sub(rd: Register, rs1: Register, rs2: Register) Self {
        return R.init(0b0110011, rd, 0, rs1, rs2, 32);
    }
    /// rd = rs1 - rs2
    fn subw(rd: Register, rs1: Register, rs2: Register) Self {
        return R.init(0b0111011, rd, 0, rs1, rs2, 32);
    }
    /// Bitwise xor.
    /// rd = rs1 ^ rs2
    fn xor(rd: Register, rs1: Register, rs2: Register) Self {
        return R.init(0b0110011, rd, 4, rs1, rs2, 0);
    }
    /// Bitwise or.
    /// rd = rs1 | rs2
    fn or_(rd: Register, rs1: Register, rs2: Register) Self {
        return R.init(0b0110011, rd, 6, rs1, rs2, 0);
    }
    /// Bitwise and.
    /// rd = rs1 & rs2
    fn and_(rd: Register, rs1: Register, rs2: Register) Self {
        return R.init(0b0110011, rd, 7, rs1, rs2, 0);
    }
    /// Shift left logical.
    /// rd = rs1 << rs2
    fn sll(rd: Register, rs1: Register, rs2: Register) Self {
        return R.init(0b0110011, rd, 1, rs1, rs2, 0);
    }
    /// Shift left logical word.
    /// rd = rs1 << rs2
    fn sllw(rd: Register, rs1: Register, rs2: Register) Self {
        return R.init(0b0111011, rd, 1, rs1, rs2, 0);
    }
    /// Shift right logical.
    /// rd = rs1 >> rs2
    fn srl(rd: Register, rs1: Register, rs2: Register) Self {
        return R.init(0b0110011, rd, 5, rs1, rs2, 0);
    }
    /// Shift right logical word.
    /// rd = rs1 >> rs2
    fn srlw(rd: Register, rs1: Register, rs2: Register) Self {
        return R.init(0b0111011, rd, 5, rs1, rs2, 0);
    }
    /// Shift right arithmetic (preserves sign bit).
    /// rd = (rs1 >> rs2) | (rs1 & (1 << (bits - 1)))
    fn sra(rd: Register, rs1: Register, rs2: Register) Self {
        return R.init(0b0110011, rd, 5, rs1, rs2, 32);
    }
    /// Shift right arithmetic (preserves sign bit) word.
    /// rd = (rs1 >> rs2) | (rs1 & (1 << (bits - 1)))
    fn sraw(rd: Register, rs1: Register, rs2: Register) Self {
        return R.init(0b0111011, rd, 5, rs1, rs2, 32);
    }
    /// Set less than, signed.
    /// rd = rs1 s< rs2 ? 1 : 0
    fn slt(rd: Register, rs1: Register, rs2: Register) Self {
        return R.init(0b0110011, rd, 2, rs1, rs2, 0);
    }
    /// Set less than, unsigned.
    /// rd = rs1 u< rs2 ? 1 : 0
    fn sltu(rd: Register, rs1: Register, rs2: Register) Self {
        return R.init(0b0110011, rd, 3, rs1, rs2, 0);
    }

    // start M extension

    /// Multiply.
    /// rd = rs1 * r2
    fn mul(rd: Register, rs1: Register, rs2: Register) Self {
        return R.init(0b0110011, rd, 0, rs1, rs2, 1);
    }
    /// Multiply, get high bits.
    /// rd = (rs1 * r2) >> 64
    fn mulh(rd: Register, rs1: Register, rs2: Register) Self {
        return R.init(0b0110011, rd, 1, rs1, rs2, 1);
    }
    /// Multiply signed `rs1` by unsigned `rs2`, get high bits.
    /// rd = (rs1 * r2) >> 64
    fn mulhsu(rd: Register, rs1: Register, rs2: Register) Self {
        return R.init(0b0110011, rd, 2, rs1, rs2, 1);
    }
    /// Multiply unsigned `rs1` by unsigned `rs2`, get high bits.
    /// rd = (rs1 * r2) >> 64
    fn mulhu(rd: Register, rs1: Register, rs2: Register) Self {
        return R.init(0b0110011, rd, 3, rs1, rs2, 1);
    }
    /// Divide, signed.
    /// rd = rs1 / r2
    fn div(rd: Register, rs1: Register, rs2: Register) Self {
        return R.init(0b0110011, rd, 4, rs1, rs2, 1);
    }
    /// Divide, unsigned.
    /// rd = rs1 / r2
    fn divu(rd: Register, rs1: Register, rs2: Register) Self {
        return R.init(0b0110011, rd, 5, rs1, rs2, 1);
    }
    /// Remainder, unsigned.
    /// rd = rs1 % r2
    fn rem(rd: Register, rs1: Register, rs2: Register) Self {
        return R.init(0b0110011, rd, 6, rs1, rs2, 1);
    }
    /// Remainder, signed (modulus).
    /// rd = rs1 % r2
    fn remu(rd: Register, rs1: Register, rs2: Register) Self {
        return R.init(0b0110011, rd, 7, rs1, rs2, 1);
    }

    // end M extension

    /// Multiply words.
    /// rd = rs1 * r2
    fn mulw(rd: Register, rs1: Register, rs2: Register) Self {
        return R.init(0b0111011, rd, 0, rs1, rs2, 1);
    }
    /// Divide words, signed.
    /// rd = rs1 / r2
    fn divw(rd: Register, rs1: Register, rs2: Register) Self {
        return R.init(0b0111011, rd, 4, rs1, rs2, 1);
    }
    /// Divide words, unsigned.
    /// rd = rs1 / r2
    fn divuw(rd: Register, rs1: Register, rs2: Register) Self {
        return R.init(0b0111011, rd, 5, rs1, rs2, 1);
    }
    /// Remainder of words, unsigned.
    /// rd = rs1 % r2
    fn remw(rd: Register, rs1: Register, rs2: Register) Self {
        return R.init(0b0111011, rd, 6, rs1, rs2, 1);
    }
    /// Remainder of words, signed (modulus).
    /// rd = rs1 % r2
    fn remuw(rd: Register, rs1: Register, rs2: Register) Self {
        return R.init(0b0111011, rd, 7, rs1, rs2, 1);
    }

    const I = packed struct(u32) {
        opcode: Opcode,
        rd: Register,
        funct3: u3,
        rs1: Register,
        imm: u12,

        fn init(opcode: Opcode, rd: Register, funct3: u3, rs1: Register, imm: i12) Self {
            return .{ .i = .{
                .opcode = opcode,
                .rd = rd,
                .funct3 = funct3,
                .rs1 = rs1,
                .imm = @bitCast(imm),
            } };
        }
    };
    /// Add immediate.
    /// rd = rs1 + imm
    fn addi(rd: Register, rs1: Register, imm: i12) Self {
        return I.init(0b0010011, rd, 0, rs1, imm);
    }
    /// Add immediate word.
    /// rd = rs1 + imm
    fn addiw(rd: Register, rs1: Register, imm: i12) Self {
        return I.init(0b0011011, rd, 0, rs1, imm);
    }
    /// Xor immediate.
    /// rd = rs1 ^ imm
    fn xori(rd: Register, rs1: Register, imm: i12) Self {
        return I.init(0b0010011, rd, 4, rs1, imm);
    }
    /// Or immediate.
    /// rd = rs1 | imm
    fn ori(rd: Register, rs1: Register, imm: i12) Self {
        return I.init(0b0010011, rd, 6, rs1, imm);
    }
    /// And immediate.
    /// rd = rs1 & imm
    fn andi(rd: Register, rs1: Register, imm: i12) Self {
        return I.init(0b0010011, rd, 7, rs1, imm);
    }
    /// Shift left logical immediate.
    /// rd = rs1 << by
    fn slli(rd: Register, rs1: Register, by: u6) Self {
        return I.init(0b0010011, rd, 1, rs1, @intCast(by));
    }
    /// Shift left logical word immediate.
    /// rd = rs1 << by
    fn slliw(rd: Register, rs1: Register, by: u5) Self {
        return I.init(0b0011011, rd, 1, rs1, @intCast(by));
    }
    /// Shift right logical immediate.
    /// rd = rs1 >> by
    fn srli(rd: Register, rs1: Register, by: u6) Self {
        return I.init(0b0010011, rd, 5, rs1, @intCast(by));
    }
    /// Shift right logical word immediate.
    /// rd = rs1 >> by
    fn srliw(rd: Register, rs1: Register, by: u6) Self {
        return I.init(0b0011011, rd, 5, rs1, @intCast(by));
    }
    /// Shift right arithmetic immediate (preserves sign bit).
    /// rd = ((((~0 * (rs1 >> (bits - 1))) << bits) | rs1) >> by)
    fn srai(rd: Register, rs1: Register, by: u6) Self {
        return I.init(0b0010011, rd, 5, rs1, 0x20 << 5 | @as(i12, @intCast(by)));
    }
    /// Shift right arithmetic word immediate (preserves sign bit).
    /// rd = ((((~0 * (rs1 >> (bits - 1))) << bits) | rs1) >> by)
    fn sraiw(rd: Register, rs1: Register, by: u6) Self {
        return I.init(0b0011011, rd, 5, rs1, 0x20 << 5 | @as(i12, @intCast(by)));
    }
    /// Set less than immediate, signed.
    /// rd = rs1 s< imm ? 1 : 0
    fn slti(rd: Register, rs1: Register, imm: i12) Self {
        return I.init(0b0010011, rd, 2, rs1, imm);
    }
    /// Set less than sign extended immediate, unsigned comparison.
    /// rd = rs1 u< imm ? 1 : 0
    fn sltiu(rd: Register, rs1: Register, imm: i12) Self {
        return I.init(0b0010011, rd, 3, rs1, imm);
    }
    /// Load byte and sign extend.
    /// rd = @as(*i8, @ptrFromInt(rs1 + imm)).*
    fn lb(rd: Register, rs1: Register, imm: i12) Self {
        return I.init(0b0000011, rd, 0, rs1, imm);
    }
    /// Load half and sign extend.
    /// rd = @as(*i16, @ptrFromInt(rs1 + imm)).*
    fn lh(rd: Register, rs1: Register, imm: i12) Self {
        return I.init(0b0000011, rd, 1, rs1, imm);
    }
    /// Load word and sign extend.
    /// rd = @as(*i32, @ptrFromInt(rs1 + imm)).*
    fn lw(rd: Register, rs1: Register, imm: i12) Self {
        return I.init(0b0000011, rd, 2, rs1, imm);
    }
    /// Load double.
    /// rd = @as(*u64, @ptrFromInt(rs1 + imm)).*
    fn ld(rd: Register, rs1: Register, imm: i12) Self {
        return I.init(0b0000011, rd, 3, rs1, imm);
    }
    /// Load byte and zero extend.
    /// rd = @as(*u8, @ptrFromInt(rs1 + imm)).*
    fn lbu(rd: Register, rs1: Register, imm: i12) Self {
        return I.init(0b0000011, rd, 4, rs1, imm);
    }
    /// Load half and zero extend.
    /// rd = @as(*u16, @ptrFromInt(rs1 + imm)).*
    fn lhu(rd: Register, rs1: Register, imm: i12) Self {
        return I.init(0b0000011, rd, 5, rs1, imm);
    }
    /// Load word and zero extend.
    /// rd = @as(*u32, @ptrFromInt(rs1 + imm)).*
    fn lwu(rd: Register, rs1: Register, imm: i12) Self {
        return I.init(0b0000011, rd, 6, rs1, imm);
    }
    /// Jump and link register.
    /// rd = pc + 4; pc = rs1 + imm
    fn jalr(rd: Register, rs1: Register, imm: i12) Self {
        return I.init(0b1100111, rd, 0, rs1, imm);
    }
    /// Environment call. Issue a syscall on linux
    fn ecall() Self {
        return I.init(0b1110011, .zero, 0, .zero, 0);
    }

    const S = packed struct(u32) {
        opcode: Opcode,
        imm4_0: u5,
        funct3: u3,
        rs1: Register,
        rs2: Register,
        imm11_5: u7,

        fn init(opcode: Opcode, funct3: u3, rs1: Register, rs2: Register, imm: i12) Self {
            const umm: u12 = @bitCast(imm);
            return .{ .s = .{
                .opcode = opcode,
                .imm4_0 = @truncate(umm),
                .funct3 = funct3,
                .rs1 = rs1,
                .rs2 = rs2,
                .imm11_5 = @truncate(umm >> 5),
            } };
        }
    };
    /// Store byte.
    /// @as(*u8, @ptrFromInt(rs1 + imm)).* = @truncate(rs2)
    fn sb(rs1: Register, imm: i12, rs2: Register) Self {
        return S.init(0b0100011, 0, rs1, rs2, imm);
    }
    /// Store half.
    /// @as(*u16, @ptrFromInt(rs1 + imm)).* = @truncate(rs2)
    fn sh(rs1: Register, imm: i12, rs2: Register) Self {
        return S.init(0b0100011, 1, rs1, rs2, imm);
    }
    /// Store word.
    /// @as(*u32, @ptrFromInt(rs1 + imm)).* = @truncate(rs2)
    fn sw(rs1: Register, imm: i12, rs2: Register) Self {
        return S.init(0b0100011, 2, rs1, rs2, imm);
    }
    /// Store double.
    /// @as(*u64, @ptrFromInt(rs1 + imm)).* = @truncate(rs2)
    fn sd(rs1: Register, imm: i12, rs2: Register) Self {
        return S.init(0b0100011, 3, rs1, rs2, imm);
    }

    /// The lowest bit of the immediate value, imm, is not stored as imm will always be even.
    const B = packed struct(u32) {
        opcode: Opcode,
        imm11: u1,
        imm4_1: u4,
        funct3: u3,
        rs1: Register,
        rs2: Register,
        imm10_5: u6,
        imm12: u1,

        fn init(opcode: Opcode, funct3: u3, rs1: Register, rs2: Register, imm: i13) Self {
            std.debug.assert(imm & 1 == 0);
            const umm: u13 = @bitCast(imm);
            return .{ .b = .{
                .opcode = opcode,
                .imm11 = @truncate(umm >> 11),
                .imm4_1 = @truncate(umm >> 1),
                .funct3 = funct3,
                .rs1 = rs1,
                .rs2 = rs2,
                .imm10_5 = @truncate(umm >> 5),
                .imm12 = @truncate(umm >> 12),
            } };
        }
    };
    /// Branch if equal.
    /// pc = if (rs1 == rs2) then pc + imm else pc + 4
    ///
    /// Advancing pc by 4 is done by all instructions, but shown here for clarity.
    /// Note that `imm` must be even since function addresses always are.
    fn beq(rs1: Register, rs2: Register, imm: i13) Self {
        return B.init(0b1100011, 0, rs1, rs2, imm);
    }
    /// Branch if not equal.
    /// pc = if (rs1 != rs2) then pc + imm else pc + 4
    ///
    /// Advancing pc by 4 is done by all instructions, but shown here for clarity.
    /// Note that `imm` must be even since function addresses always are.
    fn bne(rs1: Register, rs2: Register, imm: i13) Self {
        return B.init(0b1100011, 1, rs1, rs2, imm);
    }
    /// Branch if less than, signed.
    /// pc = if (rs1 s< rs2) then pc + imm else pc + 4
    ///
    /// Advancing pc by 4 is done by all instructions, but shown here for clarity.
    /// Note that `imm` must be even since function addresses always are.
    fn blt(rs1: Register, rs2: Register, imm: i13) Self {
        return B.init(0b1100011, 4, rs1, rs2, imm);
    }
    /// Branch if greater than or equal, signed.
    /// pc = if (rs1 s>= rs2) then pc + imm else pc + 4
    ///
    /// Advancing pc by 4 is done by all instructions, but shown here for clarity.
    /// Note that `imm` must be even since function addresses always are.
    fn bge(rs1: Register, rs2: Register, imm: i13) Self {
        return B.init(0b1100011, 5, rs1, rs2, imm);
    }
    /// Branch if less than, unsigned.
    /// pc = if (rs1 u< rs2) then pc + imm else pc + 4
    ///
    /// Advancing pc by 4 is done by all instructions, but shown here for clarity.
    /// Note that `imm` must be even since function addresses always are.
    fn bltu(rs1: Register, rs2: Register, imm: i13) Self {
        return B.init(0b1100011, 6, rs1, rs2, imm);
    }
    /// Branch if greater than or equal, unsigned.
    /// pc = if (rs1 u>= rs2) then pc + imm else pc + 4
    ///
    /// Advancing pc by 4 is done by all instructions, but shown here for clarity.
    /// Note that `imm` must be even since function addresses always are.
    fn bgeu(rs1: Register, rs2: Register, imm: i13) Self {
        return B.init(0b1100011, 7, rs1, rs2, imm);
    }

    const U = packed struct(u32) {
        opcode: Opcode,
        rd: Register,
        imm12_31: u20,

        fn init(opcode: Opcode, rd: Register, imm: i20) Self {
            return .{ .u = .{
                .opcode = opcode,
                .rd = rd,
                .imm12_31 = @bitCast(imm),
            } };
        }
    };
    /// Load upper immediate.
    /// rd = imm << 12
    fn lui(rd: Register, imm: i20) Self {
        return U.init(0b0110111, rd, imm);
    }
    /// Add upper immediate to pc.
    /// rd = pc + (imm << 12)
    fn auipc(rd: Register, imm: i20) Self {
        return U.init(0b0010111, rd, imm);
    }

    const J = packed struct(u32) {
        opcode: Opcode,
        rd: Register,
        imm12_19: u8,
        imm11: u1,
        imm1_10: u10,
        imm20: u1,

        fn init(opcode: Opcode, rd: Register, imm: i21) Self {
            std.debug.assert(imm & 1 == 0);
            const umm: u21 = @bitCast(imm);
            return .{ .j = .{
                .opcode = opcode,
                .rd = rd,
                .imm12_19 = @truncate(umm >> 12),
                .imm11 = @truncate(umm >> 11),
                .imm1_10 = @truncate(umm >> 1),
                .imm20 = @intCast(umm >> 20),
            } };
        }
    };
    /// Jump and link.
    /// rd = pc + 4; pc = pc + imm
    fn jal(rd: Register, imm: i21) Self {
        return J.init(0b1101111, rd, imm);
    }

    const Self = @This();
};

const RegisterAllocator = struct {
    allocated: std.AutoHashMap(compile.VReg, Register),
    available: std.ArrayList(Register),

    fn init(allocator: Allocator) !RegisterAllocator {
        var available: std.ArrayList(Register) = .init(allocator);
        for ([_]Register{ .s11, .s10, .s9, .s8, .s7, .s6, .s5, .s4, .s3, .s2, .s1, .s0 }) |reg| {
            try available.append(reg);
        }
        var allocated: std.AutoHashMap(compile.VReg, Register) = .init(allocator);
        try allocated.ensureTotalCapacity(@intCast(available.items.len));
        return .{
            .allocated = allocated,
            .available = available,
        };
    }

    fn deinit(self: *RegisterAllocator) void {
        self.allocated.deinit();
        self.available.deinit();
    }

    fn get(self: *const RegisterAllocator, vreg: compile.VReg) Register {
        return self.allocated.get(vreg).?;
    }

    fn allocate(self: *RegisterAllocator, vreg: compile.VReg) !Register {
        const reg = self.available.pop() orelse return error.OutOfRegisters;
        self.allocated.putAssumeCapacityNoClobber(vreg, reg);
        return reg;
    }

    fn free(self: *RegisterAllocator, vreg: compile.VReg) void {
        const ent = self.allocated.fetchRemove(vreg).?;
        const reg = ent.value;

        std.debug.assert(std.mem.indexOfScalar(Register, self.available.items, reg) == null);
        return self.available.appendAssumeCapacity(reg);
    }

    fn allocateAux(self: *RegisterAllocator) !Register {
        const reg = self.available.pop() orelse return error.OutOfRegisters;
        return reg;
    }

    fn freeAux(self: *RegisterAllocator, reg: Register) void {
        var it = self.allocated.valueIterator();
        while (it.next()) |r| std.debug.assert(reg != r.*);
        std.debug.assert(std.mem.indexOfScalar(Register, self.available.items, reg) == null);
        return self.available.appendAssumeCapacity(reg);
    }
};

const Relocation = struct {
    instr: usize,
    target: compile.BlockRef,
};

const Context = struct {
    register_allocator: RegisterAllocator,
    instructions: std.ArrayList(Instruction),
    relocations: std.ArrayList(Relocation),
    block_starts: std.ArrayList(usize),

    // Current stuff that changes often, basically here to avoid prop drilling.
    block: ?*const compile.BasicBlock = null,
    current_instruction_index: ?usize = null,

    fn deinit(self: *Context) void {
        self.register_allocator.deinit();
        self.instructions.deinit();
    }

    fn addRelocation(self: *Context, target: compile.BlockRef) !void {
        try self.relocations.append(.{
            .instr = self.instructions.items.len,
            .target = target,
        });
    }

    fn emit(self: *Context, inst: Instruction) !void {
        try self.instructions.append(inst);
    }

    /// Frees all virtual registers who's last use is the current `compile.Instr` or earlier. This
    /// must be called after the current compile.Instr's sources have been retrieved, since this
    /// will deallocate them, and after allocating auxiliary registers since otherwise they may
    /// collide with the sources. Should be called before allocating results to allow for more
    /// register re-use.
    fn freeUnusedVRegs(self: *Context) !void {
        // TODO: make this do stuff again.
        _ = self;
        // var it = self.register_allocator.allocated.keyIterator();
        // while (it.next()) |vreg| {
        //     if (self.block.?.vreg_last_use.get(vreg.*).? <= self.current_instruction_index.?) {
        //         self.register_allocator.free(vreg.*);
        //     }
        // }
    }

    fn genConstantInner(self: *Context, reg: Register, value: u64) !void {
        if (value <= std.math.maxInt(i12)) {
            // If the highest bit is set, we will get sign extension from this, but it will be
            // cleared by the next addiw.
            try self.emit(Instruction.addi(reg, .zero, @intCast(value)));
        } else if (value <= std.math.maxInt(i32)) {
            const lower: u12 = @truncate(value);
            const upper: u20 = @truncate(if (lower >> 11 == 1) (value >> 12) + 1 else value >> 12);
            // If the highest bit in upper is set here, we will get an unwanted sign extension, but
            // that will be cleared in the `addiw` right afterwards. If the highest bit was set,
            // then that must be because (lower >> 11 == 1) happened, so lower is not zero, thus the
            // `addiw` is guaranteed to be run. The only other way the highest bit can be set would
            // be if value had `1 << 32 == 1`, but then it would not be smaller than
            // `std.math.maxInt(i32)`.
            try self.emit(.lui(reg, @bitCast(upper)));
            if (lower > 0) try self.emit(.addiw(reg, reg, @bitCast(lower)));
        } else {
            const thisVal: u12 = @truncate(value);
            const nextVal = if (thisVal >> 11 == 1) (value >> 12) + 1 else value >> 12;
            try self.genConstantInner(reg, nextVal);
            try self.emit(.slli(reg, reg, 12)); // TODO: sometimes these `slli`s can be collapsed
            if (thisVal > 0) try self.emit(.addi(reg, reg, @bitCast(thisVal)));
        }
    }

    fn genConstant(self: *Context, constant: compile.Instr.Constant) !void {
        try self.freeUnusedVRegs();
        const reg = try self.register_allocator.allocate(constant.dest);
        try self.genConstantInner(reg, constant.value);
    }

    fn genBinOp(self: *Context, bin_op: compile.Instr.BinOp) !void {
        const lhs = self.register_allocator.get(bin_op.lhs);
        const rhs = self.register_allocator.get(bin_op.rhs);
        try self.freeUnusedVRegs();
        const reg = try self.register_allocator.allocate(bin_op.dest);
        switch (bin_op.op) {
            .add => try self.emit(.add(reg, lhs, rhs)),
            .sub => try self.emit(.sub(reg, lhs, rhs)),
        }
    }

    fn genProcCall(self: *Context, call: compile.Instr.ProcCall) !void {
        switch (call.proc) {
            .print => {
                const arg = self.register_allocator.get(call.arg);

                try self.freeUnusedVRegs();

                const num = try self.register_allocator.allocateAux();
                defer self.register_allocator.freeAux(num);

                if (arg != num) try self.emit(.addi(num, arg, 0));

                const quot = try self.register_allocator.allocateAux();
                defer self.register_allocator.freeAux(quot);
                const digit = try self.register_allocator.allocateAux();
                defer self.register_allocator.freeAux(digit);
                const count = try self.register_allocator.allocate(call.dest);

                try self.emit(.addi(digit, .zero, '\n'));
                try self.emit(.addi(.sp, .sp, -1));
                try self.emit(.sb(.sp, 0, digit));
                try self.emit(.addi(count, .zero, 1));
                try self.emit(.addi(quot, .zero, 10));

                try self.emit(.addi(digit, num, 0));
                try self.emit(.divu(num, digit, quot));
                try self.emit(.remu(digit, digit, quot));
                try self.emit(.addi(digit, digit, '0'));
                try self.emit(.addi(.sp, .sp, -1));
                try self.emit(.sb(.sp, 0, digit));
                try self.emit(.addi(count, count, 1));
                try self.emit(.bne(num, .zero, -4 * 7));

                try self.emit(.addi(.a7, .zero, @intFromEnum(std.os.linux.syscalls.RiscV64.write)));
                try self.emit(.addi(.a0, .zero, 1)); // fd = stdout
                try self.emit(.addi(.a1, .sp, 0)); // buf = sp
                try self.emit(.addi(.a2, count, 0)); // count = count
                try self.emit(.ecall()); // syscall(no, fd, buf, count)
                try self.emit(.add(.sp, .sp, count));

                try self.emit(.addi(count, count, -1));
            },
            .read_int => {
                try self.freeUnusedVRegs();

                const result = try self.register_allocator.allocate(call.dest);
                const ptr = try self.register_allocator.allocateAux();
                defer self.register_allocator.freeAux(ptr);
                const char = try self.register_allocator.allocateAux();
                defer self.register_allocator.freeAux(char);
                const newline = try self.register_allocator.allocateAux();
                defer self.register_allocator.freeAux(newline);

                try self.emit(.addi(.sp, .sp, -21));
                try self.emit(.addi(newline, .zero, '\n'));
                try self.emit(.addi(result, .zero, 0));

                try self.emit(.addi(.a7, .zero, @intFromEnum(std.os.linux.syscalls.RiscV64.read)));
                try self.emit(.addi(.a0, .zero, 0)); // fd = stdin
                try self.emit(.addi(.a1, .sp, 0)); // buf = sp
                try self.emit(.addi(.a2, .zero, 21)); // count = count
                try self.emit(.ecall()); // syscall(no, fd, buf, count)

                try self.emit(.addi(ptr, .sp, 0));
                try self.emit(.add(.a0, .a0, .sp)); // a0 = end

                // loop start
                try self.emit(.bgeu(ptr, .a0, 4 * 8)); // done
                try self.emit(.lb(char, ptr, 0));
                try self.emit(.beq(char, newline, 4 * 6)); // done
                try self.emit(.mul(result, result, newline)); // '\n' happens to also be 10
                try self.emit(.addi(char, char, -'0'));
                // assert 0 <= char <= 9
                try self.emit(.add(result, result, char));

                try self.emit(.addi(ptr, ptr, 1));
                try self.emit(.jal(.zero, -4 * 7)); // -> loop start
                // done

                try self.emit(.addi(.sp, .sp, 21));
            },
        }
    }

    fn genBranch(self: *Context, branch: compile.Instr.Branch) !void {
        const cond = self.register_allocator.get(branch.cond);
        try self.freeUnusedVRegs();

        try self.addRelocation(branch.false);
        try self.emit(.beq(cond, .zero, 0));

        try self.addRelocation(branch.true);
        try self.emit(.jal(.zero, 0));
    }

    fn genJump(self: *Context, jump: compile.Instr.Jump) !void {
        try self.addRelocation(jump.to);
        try self.emit(.jal(.zero, 0));
    }

    fn genExit(self: *Context, _: compile.Instr.Exit) !void {
        try self.emit(.addi(.a0, .zero, 0));
        try self.emit(.addi(.a7, .zero, 93));
        try self.emit(.ecall());
    }

    fn codegenInstr(self: *Context, instr: compile.Instr) !void {
        switch (instr.type) {
            inline else => |ty| {
                const func = comptime blk: {
                    const typeName = @typeName(@TypeOf(ty));
                    var it = std.mem.splitBackwardsScalar(u8, typeName, '.');
                    const base = it.first();
                    if (!@hasDecl(Context, "gen" ++ base)) {
                        @compileError(std.fmt.comptimePrint(
                            "codegen.Context must have a member named 'gen{s}' " ++
                                "since compile.Instr.Type has a variant named {s}",
                            .{ base, typeName },
                        ));
                    }
                    break :blk @field(Context, "gen" ++ base);
                };
                try func(self, ty);
            },
        }
    }

    fn codegenBlock(self: *Context, block: compile.BasicBlock) !void {
        self.block = &block;
        defer self.block = null;
        for (block.instrs.items, 0..) |instr, i| {
            self.current_instruction_index = i;
            try self.codegenInstr(instr);
        }
    }

    fn codegenProc(self: *Context, proc: compile.Procedure) !void {
        for (proc.blocks) |block| {
            try self.block_starts.append(self.instructions.items.len);
            try self.codegenBlock(block);
        }
    }
};

pub fn create_elf(allocator: Allocator, proc: compile.Procedure) ![]u8 {
    var ctx: Context = .{
        .register_allocator = try .init(allocator),
        .instructions = .init(allocator),
        .relocations = .init(allocator),
        .block_starts = .init(allocator),
    };
    defer ctx.deinit();

    try ctx.codegenProc(proc);

    // TODO: make this less sheiße
    for (ctx.relocations.items) |relocation| {
        const instr = &ctx.instructions.items[relocation.instr];
        const opcode: Instruction.Opcode = @truncate(@as(u32, @bitCast(instr.*)));
        const target: isize = @intCast(ctx.block_starts.items[@intFromEnum(relocation.target)]);
        const from: isize = @intCast(relocation.instr);
        switch (opcode) {
            0b1101111 => {
                const jal: Instruction.J = instr.j;
                instr.* = .jal(jal.rd, @intCast((target - from) * 4));
            },
            0b1100011 => {
                const b: Instruction.B = instr.b;
                if (b.funct3 != 0)
                    std.debug.panic("Not yet implemented instruction with relocation\n", .{});
                instr.* = .beq(b.rs1, b.rs2, @intCast((target - from) * 4));
            },
            else => std.debug.panic("Not yet implemented instruction with relocation\n", .{}),
        }
    }

    std.debug.print("allocated regs: {}\n", .{root.fmtHashMap(ctx.register_allocator.allocated)});

    var output_buffer: std.ArrayList(u8) = .init(allocator);
    errdefer output_buffer.deinit();
    try output_buffer.appendNTimes(undefined, @sizeOf(elf.Elf64_Ehdr) + @sizeOf(elf.Elf64_Phdr));
    const output = output_buffer.writer();
    for (ctx.instructions.items) |instr| {
        try output.writeInt(u32, @bitCast(instr), .little);
    }

    const base_addr = 0x10000000;
    const elf_header: elf.Elf64_Ehdr = .{
        .e_ident = elf.MAGIC.* ++ [_]u8{
            elf.ELFCLASS64, // EI_CLASS
            elf.ELFDATA2LSB, // EI_DATA
            1, // EI_VERSION
            @intFromEnum(elf.OSABI.NONE), // EI_OSABI
            0, // EI_ABIVERSION
        } ++ [1]u8{0} ** 7, // EI_PAD
        .e_type = .EXEC,
        .e_machine = elf.EM.RISCV,
        .e_version = 1,
        .e_entry = base_addr + @sizeOf(elf.Elf64_Ehdr) + @sizeOf(elf.Elf64_Phdr), // we place the code directly after the headers
        .e_phoff = @sizeOf(elf.Elf64_Ehdr), // we place the program header(s) directly after the elf header
        .e_shoff = 0,
        .e_flags = 0, // ?
        .e_ehsize = @sizeOf(elf.Elf64_Ehdr),
        .e_phentsize = @sizeOf(elf.Elf64_Phdr),
        .e_phnum = 1,
        .e_shentsize = @sizeOf(elf.Elf64_Shdr),
        .e_shnum = 0,
        .e_shstrndx = 0,
    };
    const program_header: elf.Elf64_Phdr = .{
        .p_type = elf.PT_LOAD,
        .p_flags = elf.PF_R | elf.PF_X,
        .p_offset = 0,
        .p_vaddr = base_addr,
        .p_paddr = base_addr,
        .p_filesz = output_buffer.items.len,
        .p_memsz = output_buffer.items.len,
        .p_align = 0x1000,
    };

    @memcpy(output_buffer.items[0..@sizeOf(elf.Elf64_Ehdr)], std.mem.asBytes(&elf_header));
    @memcpy(output_buffer.items[@sizeOf(elf.Elf64_Ehdr)..][0..@sizeOf(elf.Elf64_Phdr)], std.mem.asBytes(&program_header));

    return output_buffer.toOwnedSlice();
}