New codegen backend: Start emitting code.

This makes a start at emitting X86_64 code from the JIT IR.

Obviously this is non-functional at this point (it's currently not even
called from the JIT pipeline), but should serve as something we can
iterate upon and at least unit test in isolation.

Many things missing:
 - Trace input handling.
 - Correct allocation sizes.
 - Stackmaps
 - Debugger support.
 - Loads more testing.
 - Cross-arch testing support.

This PR is already large, so I've held off doing those for now.

ykrt/Cargo.toml

@@ -23,6 +23,8 @@ yktracec = { path = "../yktracec" }
 static_assertions = "1.1.0"
 typed-index-collections = "3.1.0"
 thiserror = "1.0.56"
+dynasmrt = "2.0.0"
+iced-x86 = { version = "1.21.0", features = ["decoder", "std"] }
 
 [dependencies.llvm-sys]
 # note: using a git version to get llvm linkage features in llvm-sys (not in a
@@ -46,4 +48,5 @@ yk_jitstate_debug = []
 yk_testing = []
 
 [dev-dependencies]
+fm = "0.2.2"
 num-traits = "0.2.16"

ykrt/src/compile/jitc_yk/codegen/mod.rs

@@ -0,0 +1,218 @@
+//! The JIT's Code Generator.
+
+// FIXME: eventually delete.
+#![allow(dead_code)]
+
+use super::{jit_ir, CompilationError};
+use dynasmrt::ExecutableBuffer;
+#[cfg(any(debug_assertions, test))]
+use dynasmrt::{AssemblyOffset, DynasmApi};
+#[cfg(any(debug_assertions, test))]
+use std::{cell::Cell, collections::HashMap, slice};
+use typed_index_collections::TiVec;
+
+/// FIXME: one arch is hard-compiled-in, but we should be able to "cross-unit-test" different
+/// architectures.
+#[cfg(target_arch = "x86_64")]
+mod x86_64;
+
+/// Describes storage for a local variable.
+///
+/// FIXME: This assumes everything is a stack slot. When we come to implement a register allocator,
+/// this should be an enum of stack/reg.
+#[derive(Debug, Clone, Copy)]
+struct LocalAlloc {
+    /// The offset from the base pointer of the allocation.
+    ///
+    /// This is independent of which direction the stack grows. In other words, for architectures
+    /// where the stack grows downwards, you'd subtract this from the base pointer to find the
+    /// address of the allocation.
+    ///
+    /// FIXME: when we have register allocation, consider addressing stack slots by the stack
+    /// pointer, thus freeing up the base pointer for general purpose use.
+    frame_off: usize,
+}
+
+impl LocalAlloc {
+    fn new(frame_off: usize) -> Self {
+        Self { frame_off }
+    }
+}
+
+impl LocalAlloc {
+    fn frame_off(&self) -> usize {
+        self.frame_off
+    }
+}
+
+/// The result of running the code generator.
+pub(crate) struct CodeGenOutput {
+    /// The executable code itself.
+    buf: ExecutableBuffer,
+    /// Comments to be shown when printing the compiled trace using `AsmPrinter`.
+    ///
+    /// Used for testing and debugging.
+    #[cfg(any(debug_assertions, test))]
+    comments: HashMap<AssemblyOffset, Vec<String>>,
+}
+
+impl CodeGenOutput {
+    #[cfg(any(debug_assertions, test))]
+    fn disassemble(&self) -> String {
+        let tp = AsmPrinter::new(&self.buf);
+        #[cfg(any(debug_assertions, test))]
+        let tp = tp.set_comments(&self.comments);
+        tp.to_string()
+    }
+}
+
+/// The code generator.
+pub(crate) struct CodeGen<'a> {
+    jit_mod: &'a jit_ir::Module,
+    #[cfg(target_arch = "x86_64")]
+    asm: dynasmrt::x64::Assembler,
+    /// Abstract stack pointer, as a relative offset from `RBP`. The higher this number, the larger
+    /// the JITted code's stack. That means that even on a host where the stack grows down, this
+    /// value grows up.
+    asp: usize,
+    /// Maps each JIT IR local variable to where the JIT will store it at the machine level.
+    ///
+    /// Conceptually this is a hashmap, but since the keys are integers and there can never be
+    /// "gaps" in the collection, we can store it as a vector.
+    local_allocs: TiVec<jit_ir::InstrIdx, LocalAlloc>,
+    /// Comments used by the trace printer for debugging and testing only.
+    ///
+    /// Each assembly offset can have zero or more comment lines.
+    #[cfg(any(debug_assertions, test))]
+    comments: Cell<HashMap<AssemblyOffset, Vec<String>>>,
+}
+
+impl<'a> CodeGen<'a> {
+    pub fn new(jit_mod: &'a jit_ir::Module) -> Result<Self, CompilationError> {
+        #[cfg(target_arch = "x86_64")]
+        let asm = dynasmrt::x64::Assembler::new()
+            .map_err(|e| CompilationError::Unrecoverable(e.to_string()))?;
+        Ok(Self {
+            jit_mod,
+            asm,
+            asp: 0,
+            local_allocs: TiVec::new(),
+            #[cfg(any(debug_assertions, test))]
+            comments: Cell::new(HashMap::new()),
+        })
+    }
+
+    /// Aligns the abstract stack pointer to the specified number of bytes.
+    fn align_asp(&mut self, to: usize) {
+        let rem = self.asp % to;
+        if rem != 0 {
+            self.asp += to - rem;
+        }
+    }
+
+    pub fn codegen(mut self) -> Result<CodeGenOutput, CompilationError> {
+        let alloc_off = self.emit_prologue();
+
+        // FIXME: we'd like to be able to assemble code backwards as this would simplify register
+        // allocation and side-step the need to patch up the prolog after the fact. dynasmrs
+        // doesn't support this, but it's on their roadmap:
+        // https://github.com/CensoredUsername/dynasm-rs/issues/48
+        for (idx, inst) in self.jit_mod.instrs().iter().enumerate() {
+            #[cfg(any(debug_assertions, test))]
+            self.comment(self.asm.offset(), inst.to_string());
+            self.codegen_inst(jit_ir::InstrIdx::new(idx)?, inst);
+        }
+
+        // Now we know the size of the stack frame (i.e. self.asp), patch the allocation with the
+        // correct amount.
+        self.patch_frame_allocation(alloc_off);
+
+        self.asm
+            .commit()
+            .map_err(|e| CompilationError::Unrecoverable(e.to_string()))?;
+
+        let buf = self
+            .asm
+            .finalize()
+            .map_err(|_| CompilationError::Unrecoverable("failed to finalize assembler".into()))?;
+
+        #[cfg(not(any(debug_assertions, test)))]
+        return Ok(CodeGenOutput { buf });
+        #[cfg(any(debug_assertions, test))]
+        {
+            let comments = self.comments.take();
+            return Ok(CodeGenOutput { buf, comments });
+        }
+    }
+
+    fn codegen_inst(&mut self, instr_idx: jit_ir::InstrIdx, inst: &jit_ir::Instruction) {
+        match inst {
+            jit_ir::Instruction::LoadArg(i) => self.codegen_loadarg_instr(instr_idx, &i),
+            jit_ir::Instruction::Load(i) => self.codegen_load_instr(instr_idx, &i),
+            _ => todo!(),
+        }
+    }
+
+    fn allocation(&'a self, idx: jit_ir::InstrIdx) -> &'a LocalAlloc {
+        &self.local_allocs[idx]
+    }
+
+    /// Add a comment to the trace.
+    #[cfg(any(debug_assertions, test))]
+    fn comment(&mut self, off: AssemblyOffset, line: String) {
+        self.comments.get_mut().entry(off).or_default().push(line);
+    }
+}
+
+/// Stringifies emitted code for testing and debugging purposes.
+#[cfg(any(debug_assertions, test))]
+struct AsmPrinter<'a> {
+    buf: &'a ExecutableBuffer,
+    comments: Option<&'a HashMap<AssemblyOffset, Vec<String>>>,
+}
+
+#[cfg(any(debug_assertions, test))]
+impl<'a> AsmPrinter<'a> {
+    fn new(buf: &'a ExecutableBuffer) -> Self {
+        Self {
+            buf,
+            comments: None,
+        }
+    }
+
+    /// Specifies the comments to print alongside the trace.
+    fn set_comments(mut self, comments: &'a HashMap<AssemblyOffset, Vec<String>>) -> Self {
+        self.comments = Some(comments);
+        self
+    }
+
+    /// Returns the disassembled trace.
+    fn to_string(&self) -> String {
+        let mut out = Vec::new();
+        out.push("--- Begin jit-asm ---".to_string());
+        let len = self.buf.len();
+        let bptr = self.buf.ptr(AssemblyOffset(0));
+        let code = unsafe { slice::from_raw_parts(bptr, len) };
+        self.target_print(&mut out, code, bptr as u64, len);
+        out.push("--- End jit-asm ---".into());
+        out.join("\n")
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::CodeGenOutput;
+    use fm::FMatcher;
+
+    /// Test helper to use `fm` to match a disassembled trace.
+    pub(crate) fn match_asm(cgo: &CodeGenOutput, pattern: &str) {
+        let dis = cgo.disassemble();
+        match FMatcher::new(pattern).unwrap().matches(&dis) {
+            Ok(()) => (),
+            Err(e) => panic!(
+                "\n!!! Emitted code didn't match !!!\n\n{}\nFull asm:\n{}\n",
+                e, dis
+            ),
+        }
+    }
+}