lang/src/arch/x86_64/test/asm/setup.rs

use crate::arch::x86_64::*;
use std::{collections::HashMap, fs::OpenOptions, io::Write, process::Command};

const DISPS: &[i32] = &[
    0x0,
    i8::MIN as i32,
    i8::MAX as i32,
    i16::MIN as i32,
    i16::MAX as i32,
    i32::MIN,
    i32::MAX,
];

const IMMS: &[i128] = &[
    0x0,
    i8::MIN as i128,
    i8::MAX as i128,
    i16::MIN as i128,
    i16::MAX as i128,
    i32::MIN as i128,
    i32::MAX as i128,
    i64::MIN as i128,
    i64::MAX as i128,
    u8::MAX as i128,
    u8::MAX as i128 + 1,
    u16::MAX as i128,
    u16::MAX as i128 + 1,
    u32::MAX as i128,
    u32::MAX as i128 + 1,
    i64::MAX as i128,
];

const WIDTHS: &[Width] = &[Width::B8, Width::B16, Width::B32, Width::B64];

pub fn imms() -> impl Iterator<Item = i128> {
    IMMS.iter().cloned()
}

pub fn regs() -> impl Iterator<Item = RegW> {
    RegW::IMPORTANT.iter().cloned()
}

pub fn mems() -> impl Iterator<Item = Mem> {
    gen move {
        for &reg in RegW::IMPORTANT {
            for &disp in DISPS {
                for &width in WIDTHS {
                    yield mem(reg, disp, width);
                }
            }
        }
    }
}

pub struct TestCtx {
    path: String,
    code: Code,
    cache: HashMap<String, Result<Vec<u8>, String>>,
    changed: bool,
}

#[track_caller]
pub fn eq(
    ctx: &mut TestCtx,
    asm: impl AsRef<str>,
    instr: impl Fn(&mut Code) -> Result<(), CompilerMsg>,
) {
    let asm = asm.as_ref();
    let (mut res, cache) = eq_(ctx, asm, &instr);
    if res.is_err() && cache {
        ctx.cache.remove(asm);
        res = eq_(ctx, asm, &instr).0;
    }
    if let Err(err) = res {
        panic!("{err}");
    }
}

#[track_caller]
pub fn eq_(
    ctx: &mut TestCtx,
    asm: &str,
    instr: impl FnOnce(&mut Code) -> Result<(), CompilerMsg>,
) -> (Result<(), String>, bool) {
    let (expected, cache) = if let Some(val) = ctx.cache.get(asm) {
        (val, true)
    } else {
        ctx.changed = true;
        let res = nasm(asm);
        ctx.cache.insert(asm.to_string(), res);
        (ctx.cache.get(asm).unwrap(), false)
    };
    let code = &mut ctx.code;
    let res = instr(code);
    let res = match (expected, res) {
        (Ok(expected), Err(e)) => Err(format!(
            "{asm}: failed to compile: {}\nexpected: {expected:x?}",
            e.msg
        )),
        (Err(e), Ok(_)) => {
            let res = &code.bytes[..];
            Err(format!(
                "{asm}: should not have compiled:\n{e}\ngot: {res:x?}"
            ))
        }
        (Err(_), Err(_)) => Ok(()),
        (Ok(expected), Ok(_)) => {
            let res = &code.bytes[..];
            if expected != res {
                Err(format!("{asm}: expected {expected:x?}, got {res:x?}"))
            } else {
                Ok(())
            }
        }
    };
    ctx.code.bytes.clear();
    (res, cache)
}

fn nasm(input: &str) -> Result<Vec<u8>, String> {
    let fin = "/tmp/69420nasm_in.asm";
    let fout = "/tmp/69420nasm_out.o";
    let input = "result:".to_string() + input;
    write(fin, input.as_bytes());
    run(["nasm", "-w+error", "-felf64", fin, &format!("-o{fout}")])?;
    let output = run(["objdump", "--no-addresses", "-dw", "-Mintel", fout])?;
    let mut iter = output.lines().skip_while(|l| !l.contains("result")).skip(1);
    let res_line = iter.next().unwrap().trim();
    let end = res_line.find("\t").unwrap();
    let res_line = &res_line[..end];
    let bytes = res_line
        .trim()
        .split(" ")
        .map(|s| u8::from_str_radix(s, 16).unwrap())
        .collect();
    Ok(bytes)
}

fn run<const N: usize>(input: [&str; N]) -> Result<String, String> {
    let path = input[0];
    let mut cmd = Command::new(path);
    cmd.args(&input[1..]);
    let output = cmd.output().expect("failed to run");
    if output.status.code().unwrap() != 0 {
        return Err(output.stderr.try_into().unwrap());
    }
    Ok(output.stdout.try_into().unwrap())
}

fn write(path: &str, binary: &[u8]) {
    let mut file = OpenOptions::new()
        .create(true)
        .write(true)
        .truncate(true)
        .open(path)
        .expect("Failed to create file");
    file.write_all(binary).expect("Failed to write to file");
    file.sync_all().expect("Failed to sync file");
}

const CACHE_PATH: &str = "test/nasm_cache";

impl TestCtx {
    pub fn new(name: &str) -> Self {
        let path = CACHE_PATH.to_string() + "/" + name;
        let cache = match std::fs::read(&path) {
            Ok(bytes) => bitcode::decode(&bytes).unwrap_or_default(),
            Err(_) => Default::default(),
        };
        Self {
            path,
            code: Default::default(),
            cache,
            changed: Default::default(),
        }
    }
}

impl Drop for TestCtx {
    fn drop(&mut self) {
        if self.changed {
            write(&self.path, &bitcode::encode(&self.cache));
        }
    }
}