Split out the register allocator.

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

@@ -4,37 +4,12 @@
 #![allow(dead_code)]
 
 use super::{jit_ir, CompilationError};
+use reg_alloc::RegisterAllocator;
 
 mod abs_stack;
+mod reg_alloc;
 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 }
-    }
-
-    fn frame_off(&self) -> usize {
-        self.frame_off
-    }
-}
-
 /// A trait that defines access to JIT compiled code.
 pub(crate) trait CodeGenOutput {
     /// Disassemble the code-genned trace into a string.
@@ -45,7 +20,10 @@ pub(crate) trait CodeGenOutput {
 /// 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) -> Result<Self, CompilationError>
+    fn new(
+        jit_mod: &'a jit_ir::Module,
+        ra: &'a mut dyn RegisterAllocator,
+    ) -> Result<Self, CompilationError>
     where
         Self: Sized;
 

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 });
+    }
+}