remove intermediate enum / directly encode assembly

This commit is contained in:
2026-06-11 21:49:04 -04:00
parent 91f5db6950
commit b03f755252
11 changed files with 414 additions and 668 deletions
+197 -180
View File
@@ -1,232 +1,249 @@
use std::collections::HashMap;
use crate::backend::Symbol;
use super::*;
use crate::backend::{LibImport, LinkedProgram, SymImport, SymTable, Symbol};
use util::*;
pub struct Encoder<'a> {
pub data: Vec<u8>,
pub sym_tab: SymTable<u64>,
pub missing: Vec<(usize, Symbol)>,
pub sym_refs: HashMap<Symbol, Vec<usize>>,
pub program: &'a Program<X86_64>,
type ERes = Result<(), CompilerMsg>;
/// machine code
#[derive(Default)]
pub struct Code {
pub(super) bytes: Vec<u8>,
pub(super) missing: Vec<(usize, Symbol)>,
}
pub fn compile(p: &Program<X86_64>) -> Result<LinkedProgram<u64>, CompilerMsg> {
let mut encoder = Encoder::new(p);
p.encode_data(&mut encoder.data, &mut encoder.sym_tab);
for f in &p.funcs {
let addr = encoder.data.len();
encoder.sym_tab.insert(f.sym, addr as u64);
for instr in &f.instrs {
compile_instr(&mut encoder, instr)?;
}
}
for (pos, sym) in encoder.missing.drain(..) {
let addr = encoder
.sym_tab
.get(sym)
.ok_or(CompilerMsg::from(format!("missing symbol {sym:?}")))?;
encoder.data[pos..pos + 4].copy_from_slice(&addr_offset(pos, addr))
}
let imports = p
.external
.iter()
.map(|e| LibImport {
name: e.file.clone(),
syms: e
.syms
.iter()
.map(|&s| SymImport {
name: p.sym_info(s).name.clone(),
usages: encoder.sym_refs.entry(s).or_default().clone(),
})
.collect(),
})
.collect();
Ok(LinkedProgram {
code: encoder.data,
entry: p.entry.and_then(|e| encoder.sym_tab.get(e)),
imports,
})
#[derive(Clone, Copy)]
pub struct Mem {
pub reg: Reg,
pub disp: u32,
}
type BInstr = crate::backend::Instr<X86_64>;
fn compile_instr(encoder: &mut Encoder, instr: &BInstr) -> Result<(), CompilerMsg> {
match instr {
BInstr::Asm(asm) => {
for i in &asm.instrs {
encoder.asm(*i)?;
}
}
_ => todo!(),
}
Ok(())
#[derive(Clone, Copy)]
pub enum RegImmMem {
Reg(Reg),
Imm(u64),
Mem(Mem),
}
impl Encoder<'_> {
// assembly
#[derive(Clone, Copy)]
pub enum RegMem {
Reg(Reg),
Mem(Mem),
}
pub fn mov_rr(&mut self, dst: RegH, src: RegH, width: Width) {
if width == Width::B16 {
self.data.push(0x66);
}
if src.requires_rex(width) || dst.requires_rex(width) {
self.data.push(rex(width, src, 0, dst));
}
self.data.push(0x88 | width.gt8() as u8);
self.data.push(modrm_regs(src, dst));
}
pub fn mem(reg: Reg, disp: u32) -> Mem {
Mem { reg, disp }
}
pub fn mov_ri(&mut self, dst: RegWH, src: u64) -> Result<(), CompilerMsg> {
if dst.width == Width::B16 {
self.data.push(0x66);
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(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) => {
self.prefix16(dst);
if dst.requires_rex() {
self.bytes.push(rex(dst.width(), 0, 0, dst));
}
if src > dst.width().max() {
return Err("immediate cannot fit in register".into());
}
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) => {
if src > u32::MAX as u64 {
return Err("cannot move 64 bit immediate into memory".into());
}
self.bytes.extend([rex(1, dst.reg, 0, 0), 0xc7]);
self.modrm_regdisp(dst.reg, dst.disp);
self.bytes.extend(src.to_le_bytes());
}
RegImmMem::Mem(_) => return Err("cannot move memory to memory".into()),
},
}
if dst.requires_rex() {
self.data.push(rex(dst.width, 0, 0, dst));
}
if src > dst.width.max() {
return Err("immediate cannot fit in register".into());
}
let opcode = 0xb0 | ((dst.width.gt8() as u8) << 3);
self.data.push(opcode | dst.base());
self.data.extend(&src.to_le_bytes()[..dst.width.bytes()]);
Ok(())
}
pub fn mov_rm(&mut self, dst: RegWH, src: Mem) {}
pub fn mov_mr(&mut self, dst: Mem, src: RegWH) {}
pub fn mov_mi(&mut self, dst: Mem, src: u32) {
self.data.extend([rex(1, dst.reg, 0, 0), 0xc7]);
self.modrm_regdisp(dst.reg, dst.disp);
self.data.extend(src.to_le_bytes());
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 imm.try_into() {
Ok(imm) => {
const U8: u32 = 2 << 8;
if let 0..U8 = imm {
self.bytes.push(0x6a);
self.bytes.push(imm as u8);
} else {
self.bytes.push(0x68);
self.bytes.extend(imm.to_le_bytes());
}
}
Err(_) => return Err("immediate must be 32 bit".into()),
},
RegImmMem::Mem(mem) => todo!(),
}
Ok(())
}
pub fn lea(&mut self, dst: RegWH, sym: Symbol) {
self.data
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.data.extend([0xcd, code])
self.bytes.extend([0xcd, code])
}
pub fn syscall(&mut self) {
self.data.extend([0x0f, 0x05])
self.bytes.extend([0x0f, 0x05])
}
pub fn call_i(&mut self, sym: Symbol) {
self.data.push(0xe8);
pub fn call(&mut self, sym: Symbol) {
self.bytes.push(0xe8);
self.sym_offset4(sym);
}
pub fn call_m(&mut self, sym: Symbol) {
self.data.extend([0xff, 0x15]);
pub fn call_mem(&mut self, sym: Symbol) {
self.bytes.extend([0xff, 0x15]);
self.sym_offset4(sym);
}
pub fn ret(&mut self) {
self.data.push(0xc3);
self.bytes.push(0xc3);
}
pub fn push_r(&mut self, reg: Reg, width: Width64) {
if reg.gt8() {
self.data.push(0x41);
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);
}
self.data.push(0x50 | reg.base());
}
pub fn push_i(&mut self, imm: u32) {
const U8: u32 = 2 << 8;
if let 0..U8 = imm {
self.data.push(0x6a);
self.data.push(imm as u8);
fn modrm_regdisp(&mut self, reg: Reg, disp: u32) {
let disp8 = disp < u8::MAX as u32;
let mod_ = if disp8 { 0b01 } else { 0b10 };
self.bytes.push(modrm(mod_, 0, reg.base()));
if reg.val() == rsp.val() {
// SIB
self.bytes.push(0x24);
}
if disp8 {
self.bytes.push(disp as u8);
} else {
self.data.push(0x68);
self.data.extend(imm.to_le_bytes());
self.bytes.extend(disp.to_le_bytes());
}
}
pub fn pop(&mut self, reg: Reg) {
if reg.gt8() {
self.data.push(0x41);
}
self.data.push(0x58 | reg.base());
}
/// inserts a 32 bit offset from a symbol
pub fn sym_offset4(&mut self, sym: Symbol) {
let Some(addr) = self.sym_tab.get(sym) else {
let pos = self.data.len();
self.data.extend([0; 4]);
if self.program.sym_info(sym).external {
self.sym_refs.entry(sym).or_default().push(pos);
} else {
self.missing.push((pos, sym));
}
return;
};
self.data.extend(addr_offset(self.data.len(), addr));
fn sym_offset4(&mut self, sym: Symbol) {
let pos = self.bytes.len();
self.bytes.extend([0; 4]);
self.missing.push((pos, sym));
}
pub fn asm(&mut self, instr: Instr) -> Result<(), CompilerMsg> {
match instr {
Instr::Mov(v) => match v {
Mov::RR { dst, src, width } => self.mov_rr(dst, src, width),
Mov::RI { dst, src } => self.mov_ri(dst, src)?,
Mov::RM { dst, src } => self.mov_rm(dst, src),
Mov::MI { dst, src } => self.mov_mi(dst, src),
Mov::MR { dst, src } => self.mov_mr(dst, src),
},
Instr::Int(code) => self.int(code),
Instr::Syscall => self.syscall(),
Instr::Lea { dst, sym } => self.lea(dst, sym),
Instr::Call(sym) => self.call_i(sym),
Instr::CallM(sym) => self.call_m(sym),
Instr::Ret => self.ret(),
Instr::Push(v) => match v {
Push::Reg(reg, width) => self.push_r(reg, width),
Push::Imm(imm) => self.push_i(imm),
},
Instr::Pop(reg) => self.pop(reg),
Instr::Sub => self.data.extend([0x48, 0x83, 0xec, 0x28]),
}
Ok(())
}
pub fn modrm_regdisp(&mut self, reg: impl Into<Reg>, disp: u32) {
let reg = reg.into();
let disp8 = disp < u8::MAX as u32;
let mod_ = if disp8 { 0b01 } else { 0b10 };
self.data.push(modrm(mod_, 0, reg.base()));
if reg == rsp.reg {
// SIB
self.data.push(0x24);
}
if disp8 {
self.data.push(disp as u8);
} else {
self.data.extend(disp.to_le_bytes());
}
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)));
}
}
impl<'a> Encoder<'a> {
pub fn new(program: &'a Program<X86_64>) -> Self {
Self {
data: Default::default(),
sym_tab: SymTable::new(program.sym_count()),
missing: Default::default(),
sym_refs: Default::default(),
program,
}
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 u32 as u64)
}
}