lang/src/arch/x86_64/encode.rs

use crate::backend::Symbol;

use super::*;
use util::*;

type ERes = Result<(), CompilerMsg>;

/// machine code
#[derive(Default)]
pub struct Code {
    pub(super) bytes: Vec<u8>,
    pub(super) missing: Vec<(usize, Symbol)>,
}

#[derive(Clone, Copy)]
pub struct Mem {
    pub reg: Reg,
    pub disp: i32,
    pub width: Width,
}

#[derive(Clone, Copy)]
pub enum RegImmMem {
    Reg(Reg),
    Imm(u64),
    Mem(Mem),
}

#[derive(Clone, Copy)]
pub enum RegMem {
    Reg(Reg),
    Mem(Mem),
}

pub fn mem(reg: Reg, disp: i32, width: Width) -> Mem {
    Mem { reg, disp, width }
}

impl Code {
    pub fn mov(&mut self, dst: impl Into<RegMem>, src: impl Into<RegImmMem>) -> ERes {
        let dst = dst.into();
        let src = src.into();
        match dst {
            RegMem::Reg(mut dst) => match src {
                RegImmMem::Reg(src) => {
                    if dst.width() != src.width() {
                        return Err("src and dst are not same width".into());
                    }
                    if dst.incompatible(&src) {
                        return Err("incompatible registers due to rex".into());
                    }
                    let width = dst.width();
                    self.prefix16(width);
                    if src.requires_rex() || dst.requires_rex() {
                        self.bytes.push(rex(width, src, 0, dst));
                    }
                    self.bytes.push(0x88 | width.gt8() as u8);
                    self.bytes.push(modrm_regs(src, dst));
                }
                RegImmMem::Imm(src) => {
                    let src_width = Width::fit(src);
                    if src_width > dst.width() {
                        return Err("immediate cannot fit in register".into());
                    }
                    self.prefix16(dst);
                    if src_width <= Width::B32 {
                        dst.lower64();
                    }
                    if dst.requires_rex() {
                        self.bytes.push(rex(dst.width(), 0, 0, dst));
                    }
                    let opcode = 0xb0 | ((dst.width().gt8() as u8) << 3);
                    self.bytes.push(opcode | dst.base());
                    self.bytes.extend(&src.to_le_bytes()[..dst.width().bytes()]);
                }
                RegImmMem::Mem(src) => todo!(),
            },
            RegMem::Mem(dst) => match src {
                RegImmMem::Reg(src) => todo!(),
                RegImmMem::Imm(src) => {
                    let src_width = Width::fit(src);
                    if src_width == Width::B64 {
                        return Err("cannot move 64 bit immediate into memory".into());
                    }
                    match dst.reg.width() {
                        Width::B8 | Width::B16 => return Err("invalid register width".into()),
                        Width::B32 => self.bytes.push(0x67),
                        Width::B64 => (),
                    }
                    self.prefix16(src_width);
                    if dst.reg.requires_mem_rex() {
                        self.bytes.push(rex(src_width, 0, 0, dst.reg));
                    }
                    self.bytes.push(0xc6 | (src_width != Width::B8) as u8);
                    self.modrm_regdisp(dst.reg, dst.disp);
                    self.bytes.extend(&src.to_le_bytes()[..src_width.bytes()]);
                }
                RegImmMem::Mem(_) => return Err("cannot move memory to memory".into()),
            },
        }
        Ok(())
    }

    pub fn push(&mut self, reg: impl Into<RegImmMem>) -> ERes {
        match reg.into() {
            RegImmMem::Reg(reg) => match reg.width() {
                Width::B64 => {
                    if reg.gt8() {
                        self.bytes.push(0x41);
                    }
                    self.bytes.push(0x50 | reg.base());
                }
                Width::B16 => {}
                _ => return Err("register must be 64 or 16 bit".into()),
            },
            RegImmMem::Imm(imm) => match Width::fit(imm) {
                Width::B8 => {
                    self.bytes.push(0x6a);
                    self.bytes.push(imm as u8);
                }
                Width::B16 | Width::B32 => {
                    self.bytes.push(0x68);
                    self.bytes.extend((imm as u32).to_le_bytes());
                }
                Width::B64 => return Err("immediate must be 32 bit or less".into()),
            },
            RegImmMem::Mem(mem) => todo!(),
        }
        Ok(())
    }

    pub fn pop(&mut self, reg: Reg) -> ERes {
        match reg.width() {
            Width::B64 | Width::B16 => (),
            _ => return Err("register must be 64 or 16 bit".into()),
        }
        self.prefix16(reg);
        if reg.gt8() {
            self.bytes.push(0x41);
        }
        self.bytes.push(0x58 | reg.base());
        Ok(())
    }

    pub fn lea(&mut self, dst: Reg, sym: Symbol) {
        self.bytes
            .extend([rex(1, dst, 0, 0), 0x8d, modrm_disp32(dst)]);
        self.sym_offset4(sym);
    }

    pub fn int(&mut self, code: u8) {
        self.bytes.extend([0xcd, code])
    }

    pub fn syscall(&mut self) {
        self.bytes.extend([0x0f, 0x05])
    }

    pub fn call(&mut self, sym: Symbol) {
        self.bytes.push(0xe8);
        self.sym_offset4(sym);
    }

    pub fn call_mem(&mut self, sym: Symbol) {
        self.bytes.extend([0xff, 0x15]);
        self.sym_offset4(sym);
    }

    pub fn ret(&mut self) {
        self.bytes.push(0xc3);
    }

    pub fn sub(&mut self) {
        // sub esp 40 iirc
        self.bytes.extend([0x48, 0x83, 0xec, 0x28]);
    }

    fn prefix16(&mut self, width: impl Into<Width>) {
        if width.into() == Width::B16 {
            self.bytes.push(0x66);
        }
    }

    fn modrm_regdisp(&mut self, reg: Reg, disp: i32) {
        const I8_MIN: i32 = i8::MIN as i32;
        const I8_MAX: i32 = i8::MAX as i32;
        let mod_ = match disp {
            0 => 0b00,
            I8_MIN..=I8_MAX => 0b01,
            _ => 0b10,
        };
        self.bytes.push(modrm(mod_, 0, reg.base()));
        if reg.val() == rsp.val() {
            // SIB
            self.bytes.push(0x24);
        }
        match mod_ {
            0b00 => (),
            0b01 => self.bytes.push(disp as u8),
            0b10 => self.bytes.extend(disp.to_le_bytes()),
            _ => unreachable!(),
        }
    }

    /// inserts a 32 bit offset from a symbol
    fn sym_offset4(&mut self, sym: Symbol) {
        let pos = self.bytes.len();
        self.bytes.extend([0; 4]);
        self.missing.push((pos, sym));
    }

    pub fn extend(&mut self, other: &Code) {
        let pos = self.bytes.len();
        self.bytes.extend(&other.bytes);
        self.missing
            .extend(other.missing.iter().map(|&(p, s)| (pos + p, s)));
    }
}

pub fn encode(f: impl FnOnce(&mut Code) -> Result<(), CompilerMsg>) -> Result<Code, CompilerMsg> {
    let mut code = Code::default();
    f(&mut code)?;
    Ok(code)
}

// fromrot
impl From<Reg> for RegImmMem {
    fn from(value: Reg) -> Self {
        Self::Reg(value)
    }
}

impl From<Reg> for RegMem {
    fn from(value: Reg) -> Self {
        Self::Reg(value)
    }
}

impl From<Mem> for RegImmMem {
    fn from(value: Mem) -> Self {
        Self::Mem(value)
    }
}

impl From<Mem> for RegMem {
    fn from(value: Mem) -> Self {
        Self::Mem(value)
    }
}

impl From<u64> for RegImmMem {
    fn from(value: u64) -> Self {
        Self::Imm(value)
    }
}

impl From<i64> for RegImmMem {
    fn from(value: i64) -> Self {
        Self::Imm(value as u64)
    }
}

impl From<i32> for RegImmMem {
    fn from(value: i32) -> Self {
        Self::Imm(value as u64)
    }
}