from-scratch/stages/as0/test.rs
2026-07-09 14:17:15 +02:00

416 lines
12 KiB
Rust

unsafe extern "C" {
// unsafe fn parse_gpr(src: *mut Source, prefix: u8) -> (u64, u64);
// unsafe fn parse_reg(src: *mut Source) -> (u64, u64);
unsafe fn try_parse_reg(_: *const u8) -> (usize, Register);
unsafe fn try_parse_mem(_: *const u8, _: *mut Operand) -> (usize, *const Operand);
#[link_name = "try_parse_operand"]
unsafe fn try_parse_operand_impl(_: *const u8, _: *mut Operand) -> (usize, *const Operand);
#[link_name = "try_parse_inst"]
unsafe fn try_parse_inst_impl(_: *const u8, _: *mut Instruction)
-> (usize, *const Instruction);
#[link_name = "encode_inst"]
unsafe fn encode_inst_impl(_: *mut Instruction) -> usize;
unsafe fn init_label_tables();
unsafe fn fasthash(_: *const u8, _: usize) -> u32;
#[link_name = "buf"]
static mut BUF: [u8; 0x100];
}
#[repr(C)]
#[derive(Debug, PartialEq, Eq)]
struct Register {
size: u32,
num: u32,
}
#[repr(u32)]
#[derive(Debug, PartialEq, Eq)]
enum OperandKind {
None = 0,
Reg = 1,
Mem = 2,
Imm = 4,
MemLabel = 2 | 8,
ImmLabel = 4 | 8,
}
#[repr(C, align(8))]
#[derive(Debug, PartialEq, Eq)]
struct Operand {
kind: OperandKind,
size: u32,
reg: Register,
index: Register,
scale: u32,
disp: i32,
}
impl Operand {
fn parse(bytes: &[u8]) -> Option<(usize, Self)> {
let mut op = core::mem::MaybeUninit::<Operand>::uninit();
let (n, ptr) = unsafe { try_parse_operand_impl(bytes.as_ptr(), op.as_mut_ptr()) };
if n == 0 {
None
} else {
Some((n, unsafe { op.assume_init() }))
}
}
fn label(name: &[u8]) -> Self {
let mut this = Self {
kind: OperandKind::ImmLabel,
size: 2,
reg: Register::invalid(),
index: Register::invalid(),
scale: 0,
disp: 0,
};
unsafe {
let hash = fasthash(name.as_ptr(), name.len());
(&raw mut this).byte_add(8).cast::<u32>().write(hash as u32);
}
this
}
fn imm(value: i64) -> Self {
let size = if value.abs() == 0 {
value as u32
} else {
unsafe {
(value.abs() as u64)
.checked_ilog2()
.unwrap_unchecked()
.saturating_sub(1)
.checked_ilog2()
.unwrap_or(0)
+ 1
}
};
let mut this = Self {
kind: OperandKind::Imm,
size,
reg: Register::invalid(),
index: Register::invalid(),
scale: 0,
disp: 0,
};
unsafe {
(&raw mut this).byte_add(8).cast::<i64>().write(value);
}
this
}
}
#[repr(C)]
#[derive(Debug, PartialEq, Eq)]
struct Instruction {
mnemonic_offset: i32,
num_operands: u32,
operands: [Operand; 2],
}
impl Instruction {
fn parse(bytes: &[u8]) -> Option<(usize, Self)> {
let mut inst = core::mem::MaybeUninit::<Instruction>::uninit();
let (n, ptr) = unsafe { try_parse_inst_impl(bytes.as_ptr(), inst.as_mut_ptr()) };
if n == 0 {
None
} else {
Some((n, unsafe { inst.assume_init() }))
}
}
}
impl Register {
fn new(num: u32, size: u32) -> Self {
Self {
num,
size: size.ilog2(),
}
}
fn invalid() -> Self {
Self { num: !0, size: !0 }
}
fn rbp() -> Self {
Self { num: 5, size: 0 }
}
}
static mut TABLES_INIT: std::cell::LazyCell<()> =
std::cell::LazyCell::new(|| unsafe { init_label_tables() });
#[cfg(test)]
mod tests {
use super::*;
use core::mem::MaybeUninit;
fn write_to_buf(b: &[u8]) {
assert!(b.len() <= 0x100, "Buffer overflow");
unsafe {
core::ptr::copy_nonoverlapping(b.as_ptr(), (&raw mut BUF).cast(), b.len());
}
}
#[test]
fn parse_inst() {
_ = unsafe { *TABLES_INIT };
let res = Instruction::parse(b"mov byte ['rax], 4\0");
// assert_eq!(
// res,
// Some((
// 12,
// Instruction {
// mnemonic_offset: 0,
// num_operands: 2,
// operands: [
// Operand {
// kind: OperandKind::Reg,
// size: 3,
// reg: Register::new(0, 8),
// index: Register::invalid(),
// scale: 0,
// disp: 0,
// },
// Operand {
// kind: OperandKind::Reg,
// size: 3,
// reg: Register::new(3, 8),
// index: Register::invalid(),
// scale: 0,
// disp: 0,
// },
// ],
// }
// ))
// );
unsafe {
let mut inst = res.unwrap().1;
encode_inst_impl((&raw mut inst).cast());
}
}
#[test]
fn parse_mem() {
let mut op = MaybeUninit::<Operand>::uninit();
let (n, _) =
unsafe { try_parse_mem(b"qword [rax + rbx * 2 + 0x10]\0".as_ptr(), op.as_mut_ptr()) };
assert_eq!(
(n, unsafe { op.assume_init_ref() }),
(
28,
&Operand {
kind: OperandKind::Mem,
size: 3,
reg: Register::new(0, 8),
index: Register::new(3, 8),
scale: 2,
disp: 0x10,
}
)
);
let (n, _) =
unsafe { try_parse_mem(b"byte [r12 + r13 * 4 - 0x20]\0".as_ptr(), op.as_mut_ptr()) };
assert_eq!(
(n, unsafe { op.assume_init_ref() }),
(
27,
&Operand {
kind: OperandKind::Mem,
size: 0,
reg: Register::new(12, 8),
index: Register::new(13, 8),
scale: 4,
disp: -0x20,
}
)
);
let (n, _) = unsafe { try_parse_mem(b"word [r8 * 8 + 5]\0".as_ptr(), op.as_mut_ptr()) };
assert_eq!(
(n, unsafe { op.assume_init_ref() }),
(
17,
&Operand {
kind: OperandKind::Mem,
size: 1,
reg: Register::rbp(),
index: Register::new(8, 8),
scale: 8,
disp: 5,
}
)
);
}
#[test]
fn parse_operand() {
let cases = [
(
&b"rax\0"[..],
Operand {
kind: OperandKind::Reg,
size: 3,
reg: Register::new(0, 8),
index: Register::invalid(),
scale: 0,
disp: 0,
},
),
(
&b"qword [rax + rbx * 2 + 0x10]\0"[..],
Operand {
kind: OperandKind::Mem,
size: 3,
reg: Register::new(0, 8),
index: Register::new(3, 8),
scale: 2,
disp: 0x10,
},
),
(
&b"byte [r12 + r13 * 4 - 0x20]\0"[..],
Operand {
kind: OperandKind::Mem,
size: 0,
reg: Register::new(12, 8),
index: Register::new(13, 8),
scale: 4,
disp: -0x20,
},
),
(
&b"word [r8 * 8 + 5]\0"[..],
Operand {
kind: OperandKind::Mem,
size: 1,
reg: Register::rbp(),
index: Register::new(8, 8),
scale: 8,
disp: 5,
},
),
(&b"0x12345678\0"[..], Operand::imm(0x12345678)),
(&b"-0x12345678\0"[..], Operand::imm(-0x12345678)),
(&b"0\0"[..], Operand::imm(0)),
(&b"-1\0"[..], Operand::imm(-1)),
(&b"1\0"[..], Operand::imm(1)),
(&b"255\0"[..], Operand::imm(255)),
(&b"-255\0"[..], Operand::imm(-255)),
(&b"'wazzaah\0"[..], Operand::label(b"wazzaah")),
(
&b"qword [rsp + 'hii]"[..],
Operand {
kind: OperandKind::MemLabel,
size: 3,
reg: Register::new(4, 8),
index: Register::invalid(),
scale: 0,
disp: unsafe { fasthash(b"hii".as_ptr(), 3) as i32 },
},
),
];
for (text, expected) in cases {
let (n, result) = Operand::parse(text).expect("Failed to parse operand");
assert_eq!(
result,
expected,
"Failed to parse '{}'",
std::str::from_utf8(text).unwrap()
);
}
}
#[test]
fn parse_reg() {
let cases = [
(&b"rax\0"[..], (0, 8)),
(&b"rcx\0"[..], (1, 8)),
(&b"rdx\0"[..], (2, 8)),
(&b"rbx\0"[..], (3, 8)),
(&b"rsp\0"[..], (4, 8)),
(&b"rbp\0"[..], (5, 8)),
(&b"rsi\0"[..], (6, 8)),
(&b"rdi\0"[..], (7, 8)),
(&b"r8\0"[..], (8, 8)),
(&b"r9\0"[..], (9, 8)),
(&b"r10\0"[..], (10, 8)),
(&b"r11\0"[..], (11, 8)),
(&b"r12\0"[..], (12, 8)),
(&b"r13\0"[..], (13, 8)),
(&b"r14\0"[..], (14, 8)),
(&b"r15\0"[..], (15, 8)),
(&b"rip\0"[..], (!0, 8)),
(&b"eax\0"[..], (0, 4)),
(&b"ecx\0"[..], (1, 4)),
(&b"edx\0"[..], (2, 4)),
(&b"ebx\0"[..], (3, 4)),
(&b"esp\0"[..], (4, 4)),
(&b"ebp\0"[..], (5, 4)),
(&b"esi\0"[..], (6, 4)),
(&b"edi\0"[..], (7, 4)),
(&b"r8d\0"[..], (8, 4)),
(&b"r9d\0"[..], (9, 4)),
(&b"r10d\0"[..], (10, 4)),
(&b"r11d\0"[..], (11, 4)),
(&b"r12d\0"[..], (12, 4)),
(&b"r13d\0"[..], (13, 4)),
(&b"r14d\0"[..], (14, 4)),
(&b"r15d\0"[..], (15, 4)),
(&b"r8w\0"[..], (8, 2)),
(&b"r9w\0"[..], (9, 2)),
(&b"r10w\0"[..], (10, 2)),
(&b"r11w\0"[..], (11, 2)),
(&b"r12w\0"[..], (12, 2)),
(&b"r13w\0"[..], (13, 2)),
(&b"r14w\0"[..], (14, 2)),
(&b"r15w\0"[..], (15, 2)),
(&b"r8b\0"[..], (8, 1)),
(&b"r9b\0"[..], (9, 1)),
(&b"r10b\0"[..], (10, 1)),
(&b"r11b\0"[..], (11, 1)),
(&b"r12b\0"[..], (12, 1)),
(&b"r13b\0"[..], (13, 1)),
(&b"r14b\0"[..], (14, 1)),
(&b"r15b\0"[..], (15, 1)),
(&b"ax\0"[..], (0, 2)),
(&b"cx\0"[..], (1, 2)),
(&b"dx\0"[..], (2, 2)),
(&b"bx\0"[..], (3, 2)),
(&b"sp\0"[..], (4, 2)),
(&b"bp\0"[..], (5, 2)),
(&b"si\0"[..], (6, 2)),
(&b"di\0"[..], (7, 2)),
(&b"al\0"[..], (0, 1)),
(&b"cl\0"[..], (1, 1)),
(&b"dl\0"[..], (2, 1)),
(&b"bl\0"[..], (3, 1)),
(&b"spl\0"[..], (4, 1)),
(&b"bpl\0"[..], (5, 1)),
(&b"sil\0"[..], (6, 1)),
(&b"dil\0"[..], (7, 1)),
];
for (text, (reg, size)) in cases {
let (n, result) = unsafe { try_parse_reg(text.as_ptr()) };
assert_eq!(
(n, result),
(text.len() - 1, Register::new(reg, size)),
"Failed to parse '{}' n={n}",
std::str::from_utf8(text).unwrap()
);
}
}
}