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.

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/abs_stack.rs

@@ -0,0 +1,81 @@
+//! The abstract stack.
+
+/// This data structure keeps track of an abstract stack pointer for a JIT frame during code
+/// generation. The abstract stack pointer is zero-based, so the stack pointer value also serves as
+/// the size of the stack.
+///
+/// The implementation is platform agnostic: as the abstract stack gets bigger, the abstract stack
+/// pointer grows upwards, even on architectures where the stack grows downwards.
+#[derive(Debug, Default)]
+pub(crate) struct AbstractStack(usize);
+
+impl AbstractStack {
+    /// Aligns the abstract stack pointer to the specified number of bytes.
+    ///
+    /// Returns the newly aligned stack pointer.
+    pub(crate) fn align(&mut self, to: usize) -> usize {
+        let rem = self.0 % to;
+        if rem != 0 {
+            self.0 += to - rem;
+        }
+        self.0
+    }
+
+    /// Makes the stack bigger by `nbytes` bytes.
+    ///
+    /// Returns the new stack pointer.
+    pub(crate) fn grow(&mut self, nbytes: usize) -> usize {
+        self.0 += nbytes;
+        self.0
+    }
+
+    /// Returns the stack pointer value.
+    pub(crate) fn size(&self) -> usize {
+        self.0
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::AbstractStack;
+
+    #[test]
+    fn grow() {
+        let mut s = AbstractStack::default();
+        assert_eq!(s.grow(8), 8);
+        assert_eq!(s.grow(8), 16);
+        assert_eq!(s.grow(1), 17);
+        assert_eq!(s.grow(0), 17);
+        assert_eq!(s.grow(1000), 1017);
+    }
+
+    #[test]
+    fn align() {
+        let mut s = AbstractStack::default();
+        for i in 1..100 {
+            assert_eq!(s.align(i), 0);
+            assert_eq!(s.align(i), 0);
+        }
+        for i in 1..100 {
+            s.grow(1);
+            assert_eq!(s.align(1), i);
+            assert_eq!(s.align(1), i);
+        }
+        assert_eq!(s.align(8), 104);
+        for i in 105..205 {
+            assert_eq!(s.align(i), i);
+            assert_eq!(s.align(i), i);
+        }
+        assert_eq!(s.align(12345678), 12345678);
+        assert_eq!(s.align(12345678), 12345678);
+    }
+
+    #[test]
+    fn size() {
+        let mut s = AbstractStack::default();
+        for i in 1..100 {
+            s.grow(1);
+            assert_eq!(s.size(), i);
+        }
+    }
+}

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

@@ -0,0 +1,50 @@
+//! The JIT's Code Generator.
+
+// FIXME: eventually delete.
+#![allow(dead_code)]
+
+use super::{jit_ir, CompilationError};
+use reg_alloc::RegisterAllocator;
+
+mod abs_stack;
+mod reg_alloc;
+mod x86_64;
+
+/// A trait that defines access to JIT compiled code.
+pub(crate) trait CodeGenOutput {
+    /// Disassemble the code-genned trace into a string.
+    #[cfg(any(debug_assertions, test))]
+    fn disassemble(&self) -> String;
+}
+
+/// All code generators conform to this contract.
+trait CodeGen<'a> {
+    /// Instantiate a code generator for the specified JIT module.
+    fn new(
+        jit_mod: &'a jit_ir::Module,
+        ra: &'a mut dyn RegisterAllocator,
+    ) -> Result<Self, CompilationError>
+    where
+        Self: Sized;
+
+    /// Perform code generation.
+    fn codegen(self) -> Result<Box<dyn CodeGenOutput>, CompilationError>;
+}
+
+#[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: Box<dyn 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
+            ),
+        }
+    }
+}

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

@@ -0,0 +1,68 @@
+//! Register allocation.
+//!
+//! This module:
+//!  - describes the generic interface to register allocators.
+//!  - contains concrete implementations of register allocators.
+
+use super::{super::jit_ir, abs_stack::AbstractStack};
+
+mod spill_alloc;
+#[cfg(test)]
+pub(crate) use spill_alloc::SpillAllocator;
+
+/// Describes a local variable allocation.
+#[derive(Debug, Clone, Copy, Eq, PartialEq)]
+pub(crate) enum LocalAlloc {
+    /// The local variable is on the stack.
+    Stack {
+        /// 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.
+        ///
+        /// OPT: consider addressing relative to the stack pointer, thus freeing up the base
+        /// pointer for general purpose use.
+        frame_off: usize,
+    },
+    /// The local variable is in a register.
+    ///
+    /// FIXME: unimplemented.
+    Register,
+}
+
+impl LocalAlloc {
+    /// Create a [Self::Stack] allocation.
+    pub(crate) fn new_stack(frame_off: usize) -> Self {
+        Self::Stack { frame_off }
+    }
+}
+
+/// Indicates the direction of stack growth.
+pub(crate) enum StackDirection {
+    GrowsUp,
+    GrowsDown,
+}
+
+/// The API to regsiter allocators.
+///
+/// Register allocators are responsible for assigning storage for local variables.
+pub(crate) trait RegisterAllocator {
+    /// Creates a register allocator for a stack growing in the specified direction.
+    fn new(stack_dir: StackDirection) -> Self
+    where
+        Self: Sized;
+
+    /// Allocates `size` bytes storage space for the local variable defined by the instruction with
+    /// index `local`.
+    fn allocate(
+        &mut self,
+        local: jit_ir::InstrIdx,
+        size: usize,
+        stack: &mut AbstractStack,
+    ) -> LocalAlloc;
+
+    /// Return the allocation for the value computed by the instruction at the specified
+    /// instruction index.
+    fn allocation<'a>(&'a self, idx: jit_ir::InstrIdx) -> &'a LocalAlloc;
+}

ykrt/src/compile/jitc_yk/codegen/reg_alloc/spill_alloc.rs

@@ -0,0 +1,142 @@
+//! The spill allocator.
+//!
+//! This is a register allocator that always allocates to the stack, so in fact it's not much of a
+//! register allocator at all.
+
+use super::{
+    super::{abs_stack::AbstractStack, jit_ir},
+    LocalAlloc, RegisterAllocator, StackDirection,
+};
+use typed_index_collections::TiVec;
+
+pub(crate) struct SpillAllocator {
+    allocs: TiVec<jit_ir::InstrIdx, LocalAlloc>,
+    stack_dir: StackDirection,
+}
+
+impl RegisterAllocator for SpillAllocator {
+    fn new(stack_dir: StackDirection) -> SpillAllocator {
+        Self {
+            allocs: Default::default(),
+            stack_dir,
+        }
+    }
+
+    fn allocate(
+        &mut self,
+        local: jit_ir::InstrIdx,
+        size: usize,
+        stack: &mut AbstractStack,
+    ) -> LocalAlloc {
+        // Under the current design, there can't be gaps in [self.allocs] and local variable
+        // allocations happen sequentially. So the local we are currently allocating should be the
+        // next unallocated index.
+        debug_assert!(jit_ir::InstrIdx::new(self.allocs.len()).unwrap() == local);
+
+        // Align the stack to the size of the allocation.
+        //
+        // FIXME: perhaps we should align to the largest alignment of the constituent fields?
+        // To do this we need to first finish proper type sizes.
+        let post_align_off = stack.align(size);
+
+        // Make space for the allocation.
+        let post_grow_off = stack.grow(size);
+
+        // If the stack grows up, then the allocation's offset is the stack height *before* we've
+        // made space on the stack, otherwise it's the stack height *after*.
+        let alloc_off = match self.stack_dir {
+            StackDirection::GrowsUp => post_align_off,
+            StackDirection::GrowsDown => post_grow_off,
+        };
+
+        let alloc = LocalAlloc::new_stack(alloc_off);
+        self.allocs.push(alloc);
+        alloc
+    }
+
+    fn allocation<'a>(&'a self, idx: jit_ir::InstrIdx) -> &'a LocalAlloc {
+        &self.allocs[idx]
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::compile::jitc_yk::{
+        codegen::{
+            abs_stack::AbstractStack,
+            reg_alloc::{LocalAlloc, RegisterAllocator, SpillAllocator, StackDirection},
+        },
+        jit_ir::InstrIdx,
+    };
+
+    #[test]
+    fn grow_down() {
+        let mut stack = AbstractStack::default();
+        let mut sa = SpillAllocator::new(StackDirection::GrowsDown);
+
+        let idx = InstrIdx::new(0).unwrap();
+        sa.allocate(idx, 8, &mut stack);
+        debug_assert_eq!(stack.size(), 8);
+        debug_assert_eq!(sa.allocation(idx), &LocalAlloc::Stack { frame_off: 8 });
+
+        let idx = InstrIdx::new(1).unwrap();
+        sa.allocate(idx, 1, &mut stack);
+        debug_assert_eq!(stack.size(), 9);
+        debug_assert_eq!(sa.allocation(idx), &LocalAlloc::Stack { frame_off: 9 });
+    }
+
+    #[test]
+    fn grow_up() {
+        let mut stack = AbstractStack::default();
+        let mut sa = SpillAllocator::new(StackDirection::GrowsUp);
+
+        let idx = InstrIdx::new(0).unwrap();
+        sa.allocate(idx, 8, &mut stack);
+        debug_assert_eq!(stack.size(), 8);
+        debug_assert_eq!(sa.allocation(idx), &LocalAlloc::Stack { frame_off: 0 });
+
+        let idx = InstrIdx::new(1).unwrap();
+        sa.allocate(idx, 1, &mut stack);
+        debug_assert_eq!(stack.size(), 9);
+        debug_assert_eq!(sa.allocation(idx), &LocalAlloc::Stack { frame_off: 8 });
+    }
+
+    #[cfg(debug_assertions)]
+    #[should_panic]
+    #[test]
+    fn allocate_out_of_order() {
+        let mut stack = AbstractStack::default();
+        let mut sa = SpillAllocator::new(StackDirection::GrowsUp);
+        // panics because the backing store for local allocations are a "unsparse vector" and Local
+        // 0 hasn't been allocated yet.
+        sa.allocate(InstrIdx::new(1).unwrap(), 1, &mut stack);
+    }
+
+    #[test]
+    fn compose_alloc_and_align_down() {
+        let mut stack = AbstractStack::default();
+        let mut sa = SpillAllocator::new(StackDirection::GrowsDown);
+
+        sa.allocate(InstrIdx::new(0).unwrap(), 8, &mut stack);
+        stack.align(32);
+
+        let idx = InstrIdx::new(1).unwrap();
+        sa.allocate(idx, 1, &mut stack);
+        debug_assert_eq!(stack.size(), 33);
+        debug_assert_eq!(sa.allocation(idx), &LocalAlloc::Stack { frame_off: 33 });
+    }
+
+    #[test]
+    fn compose_alloc_and_align_up() {
+        let mut stack = AbstractStack::default();
+        let mut sa = SpillAllocator::new(StackDirection::GrowsUp);
+
+        sa.allocate(InstrIdx::new(0).unwrap(), 8, &mut stack);
+        stack.align(32);
+
+        let idx = InstrIdx::new(1).unwrap();
+        sa.allocate(idx, 1, &mut stack);
+        debug_assert_eq!(stack.size(), 33);
+        debug_assert_eq!(sa.allocation(idx), &LocalAlloc::Stack { frame_off: 32 });
+    }
+}