do some registers

codegen/src/registers.rs

@@ -0,0 +1,122 @@
+//! # Register allocation
+//!
+//! Register allocation is not very smart, but also not too stupid. It tries to put SSA
+//! registers into machine registers as much as possible.
+//!
+//! ```text
+//! bb0:
+//!   %0 = 0
+//!   %1 = 1
+//!   %2 = add %0 %1
+//!   switch %2, then bb1, else bb2
+//!
+//! bb1:
+//!   %3 = add %1, 1
+//!   
+//! bb2:
+//!   %4 = add %2, 2
+//! ```
+//!
+//! For all SSA registers, we establish their "point of last use". This is the bb,stmt where their last usage occurs.
+//!
+//! First, we establish a list of possible registers to allocate.
+//! Since we immediately alloca all parameters, all the param registers are free real estate.
+//!
+//! ```text
+//! rax, rdi, rsi, rcx, rdx, r8, r9
+//! ```
+//!
+//! This forms our priority list of registers.
+//!
+//! Every time a statement has a return value, we try to assign that SSA register into a new machine register.
+//! For this, we iterate through the register list above and find the first register that's free. If we see a register
+//! that is not used anymore at the current location, we throw it out and use that new slot.
+//!
+//! When codegening an SSA register, we look into a lookup table from SSA register to machine register/stack spill and use that.
+//!
+//! When the list above is full, we spill the register to the stack. This should be rare. If the register doesn't fit into a machine
+//! register, it's also spilled.
+//!
+//! We do a first pass over the function to calculate all the offsets and registers
+//! we want to use.
+
+use std::cell::Cell;
+
+use analysis::ir::{self, Func, Location, Register};
+use rustc_hash::FxHashMap;
+
+/// A machine register from our register list described in the module documentation.
+#[derive(Debug, Clone, Copy)]
+pub struct MachineReg(pub usize);
+
+#[derive(Debug, Clone, Copy)]
+pub enum RegValue {
+    /// The SSA register contains an address on the stack.
+    /// The offset is the offset from the start of the function.
+    StackRelative { offset: u64 },
+    /// The SSA register resides on the stack as it has been spilled.
+    /// This should be rather rare in practice.
+    Spilled { offset: u64 },
+    /// The SSA register resides in a machine register.
+    MachineReg(MachineReg),
+}
+
+#[derive(Debug)]
+pub struct FunctionLayout {
+    /// Where a register comes from at a particular usage of a register.
+    register_uses: FxHashMap<(Location, Register), RegValue>,
+    total_stack_space: u64,
+}
+
+pub fn compute_layout(f: &Func) -> FunctionLayout {
+    let register_uses = FxHashMap::default();
+
+    FunctionLayout {
+        register_uses,
+        total_stack_space: 0,
+    }
+}
+
+pub struct LayoutPrinter<'a>(Cell<Option<&'a Func<'a>>>, &'a FunctionLayout);
+
+impl<'a> ir::pretty::Customizer<'a> for LayoutPrinter<'a> {
+    fn start_func(&self, func: &'a Func<'a>) {
+        self.0.set(Some(func));
+    }
+
+    fn fmt_reg(
+        &self,
+        reg: Register,
+        f: &mut std::fmt::Formatter<'_>,
+        loc: Location,
+    ) -> std::fmt::Result {
+        let layout = self.1.register_uses.get(&(loc, reg));
+        write!(f, "{{")?;
+
+        match self.0.get().unwrap().regs[reg.0 as usize].name {
+            None => write!(f, "%{}", reg.0)?,
+            Some(name) => write!(f, "%{name}")?,
+        }
+
+        write!(f, ", ")?;
+        match layout {
+            Some(RegValue::MachineReg(mach)) => write!(f, "reg-{}", mach.0)?,
+            Some(RegValue::Spilled { offset }) => write!(f, "spill-{offset}")?,
+            Some(RegValue::StackRelative { offset }) => {
+                write!(f, "i-forgot-what-this-meant-{offset}")?
+            }
+            None => write!(f, "<unknown>")?,
+        }
+
+        write!(f, "}}")
+    }
+}
+
+pub fn debug_layout(func: &Func, layout: &FunctionLayout) {
+    let custom = LayoutPrinter(Cell::default(), layout);
+
+    println!("----- code layout");
+    println!("{}", ir::pretty::func_to_string(func, &custom));
+
+    dbg!(layout);
+}