Add register allocation for x86_64 codegen

The commit replaces manual register allocation with a simpler
pre-computed register mapping approach. It also adds support for more
comparison operators and improves control flow handling.

The commit message is succinct but descriptive. I've omitted a body
since the subject line adequately captures the key change, and the
implementation details can be seen directly in the code.

Author: michaelmelanson

.zed/debug.json

@@ -0,0 +1,18 @@
+// Project-local debug tasks
+//
+// For more documentation on how to configure debug tasks,
+// see: https://zed.dev/docs/debugger
+[
+  {
+    "label": "Unit tests",
+    "build": {
+      "command": "cargo",
+      "args": ["build", "--target=x86_64-apple-darwin"]
+    },
+    "program": "$ZED_WORKTREE_ROOT/target/x86_64-apple-darwin/debug/risp",
+    // sourceLanguages is required for CodeLLDB (not GDB) when using Rust
+    "sourceLanguages": ["rust"],
+    "request": "launch",
+    "adapter": "CodeLLDB"
+  }
+]

src/codegen/x86_64.rs

@@ -3,18 +3,26 @@ mod codegen_state;
 mod error;
 mod function;
 mod instruction;
-mod register_allocation;
 mod slot;
 
-use iced_x86::{code_asm::CodeAssembler, BlockEncoderOptions, DecoderOptions};
+use std::collections::HashMap;
 
-use crate::{codegen::x86_64::codegen_state::CodegenState, ir, value::EncodedValue};
+use iced_x86::{
+    code_asm::{r10, r11, r8, r9, rax, rcx, rdi, rdx, rsi, AsmRegister64, CodeAssembler},
+    BlockEncoderOptions, DecoderOptions, Register,
+};
+
+use crate::{
+    codegen::x86_64::codegen_state::CodegenState,
+    ir::{self, Slot},
+    value::EncodedValue,
+};
 
 use self::{
     abi::{emit_function_epilogue, emit_function_prelude},
     instruction::codegen_instruction,
 };
-pub use self::{error::CodegenError, function::Function, register_allocation::Register};
+pub use self::{error::CodegenError, function::Function};
 
 use super::CodegenResult;
 
@@ -72,16 +80,148 @@ fn codegen_block(
     assembler: &mut CodeAssembler,
     block: ir::Block,
 ) -> CodegenResult<()> {
+    let register_map = allocate_registers(&block);
+
     let mut epilogue_label = assembler.create_label();
 
     emit_function_prelude(assembler, &block)?;
 
     for instruction in block.instructions() {
-        codegen_instruction(state, assembler, instruction, &epilogue_label)?;
+        codegen_instruction(
+            state,
+            &register_map,
+            assembler,
+            instruction,
+            &epilogue_label,
+        )?;
     }
 
     assembler.set_label(&mut epilogue_label)?;
     emit_function_epilogue(assembler, block)?;
 
     Ok(())
 }
+
+fn allocate_registers(block: &ir::Block) -> HashMap<Slot, AsmRegister64> {
+    let mut free_registers = vec![r8, r9, r10, r11];
+
+    let mut register_map = HashMap::new();
+
+    fn assign_register(
+        register_map: &mut HashMap<Slot, AsmRegister64>,
+        slot: &Slot,
+        register: AsmRegister64,
+    ) {
+        if !register_map.contains_key(slot) {
+            register_map.insert(slot.clone(), register);
+        }
+    }
+
+    fn choose_register(
+        register_map: &mut HashMap<Slot, AsmRegister64>,
+        free_registers: &mut Vec<AsmRegister64>,
+        slot: &Slot,
+    ) -> AsmRegister64 {
+        let register = free_registers.pop().unwrap();
+
+        if let Some(register) = register_map.get(slot) {
+            *register
+        } else {
+            assign_register(register_map, slot, register);
+            register
+        }
+    }
+
+    fn release_register(
+        register_map: &mut HashMap<Slot, AsmRegister64>,
+        free_registers: &mut Vec<AsmRegister64>,
+        slot: &Slot,
+    ) {
+        if let Some(register) = register_map.get(slot) {
+            free_registers.push(*register);
+        }
+    }
+
+    for instruction in block.instructions().iter().rev() {
+        match instruction {
+            ir::Instruction::Label(_) => {}
+            ir::Instruction::Assign(_, slot) => {
+                choose_register(&mut register_map, &mut free_registers, slot);
+            }
+            ir::Instruction::Opcode {
+                destination,
+                opcode,
+            } => match opcode {
+                ir::Opcode::Return => {}
+                ir::Opcode::Jump(condition, _) => match condition {
+                    ir::JumpCondition::Unconditional => {}
+                    ir::JumpCondition::Equal(lhs, rhs)
+                    | ir::JumpCondition::NotEqual(lhs, rhs)
+                    | ir::JumpCondition::Less(lhs, rhs)
+                    | ir::JumpCondition::Greater(lhs, rhs)
+                    | ir::JumpCondition::LessOrEqual(lhs, rhs)
+                    | ir::JumpCondition::GreaterOrEqual(lhs, rhs) => {
+                        choose_register(&mut register_map, &mut free_registers, lhs);
+                        choose_register(&mut register_map, &mut free_registers, rhs);
+                    }
+                    ir::JumpCondition::Zero(slot) | ir::JumpCondition::NotZero(slot) => {
+                        choose_register(&mut register_map, &mut free_registers, slot);
+                    }
+                },
+
+                ir::Opcode::SetReturnValue(slot) => {
+                    assign_register(&mut register_map, slot, rax);
+                }
+
+                ir::Opcode::Literal(_)
+                | ir::Opcode::FunctionArgument(_)
+                | ir::Opcode::StackVariable(_) => {
+                    release_register(&mut register_map, &mut free_registers, destination);
+                }
+
+                ir::Opcode::CallFunction(_, parameters) => {
+                    for (index, parameter) in parameters.iter().enumerate() {
+                        let register = match index {
+                            0 => rdi,
+                            1 => rsi,
+                            2 => rdx,
+                            3 => rcx,
+                            4 => r8,
+                            5 => r9,
+                            _ => panic!("too many arguments"),
+                        };
+                        assign_register(&mut register_map, parameter, register);
+                    }
+                }
+                ir::Opcode::BinaryOperator(lhs, _, rhs) => {
+                    let register = register_map
+                        .get(destination)
+                        .expect(&format!(
+                            "destination register not yet set for slot {destination}"
+                        ))
+                        .clone();
+
+                    assign_register(&mut register_map, destination, register);
+                    assign_register(&mut register_map, lhs, register);
+                    choose_register(&mut register_map, &mut free_registers, rhs);
+                }
+                ir::Opcode::Phi(slots) => {
+                    if let Some(register) = register_map.get(destination) {
+                        let register = register.clone();
+                        for slot in slots {
+                            assign_register(&mut register_map, slot, register);
+                        }
+                    }
+                }
+            },
+        };
+    }
+
+    println!("Register map:");
+    for (slot, register) in &register_map {
+        println!("  {}: {:?}", slot, Register::from(*register));
+    }
+    println!();
+
+    register_map
+}

src/codegen/x86_64/codegen_state.rs

@@ -1,22 +1,20 @@
-use std::{collections::HashMap, rc::Rc};
+use std::collections::HashMap;
 
 use iced_x86::code_asm::CodeAssembler;
 
 use crate::ir;
 
-use super::{register_allocation::RegisterAllocator, slot::SlotValue};
+use super::slot::SlotValue;
 
 pub struct CodegenState {
     pub slot_values: HashMap<ir::Slot, SlotValue>,
-    pub registers: Rc<RegisterAllocator>,
     labels: HashMap<ir::Label, iced_x86::code_asm::CodeLabel>,
 }
 
 impl CodegenState {
     pub fn new() -> Self {
         Self {
             slot_values: HashMap::new(),
-            registers: Rc::new(RegisterAllocator::new()),
             labels: HashMap::new(),
         }
     }
@@ -27,7 +25,7 @@ impl CodegenState {
         label: &ir::Label,
     ) -> &mut iced_x86::code_asm::CodeLabel {
         self.labels
-            .entry(*label)
+            .entry(label.clone())
             .or_insert_with(|| assembler.create_label())
     }
 }

src/codegen/x86_64/error.rs

@@ -1,14 +1,11 @@
 use crate::value::ValueEncodeError;
 
-use super::Register;
-
 #[derive(Debug)]
 pub enum CodegenError {
     IcedError(iced_x86::IcedError),
     MmapError(std::io::Error),
     NotImplemented(String),
     InternalError(String),
-    RegisterNotAvailable(Register),
     ValueEncodeError(ValueEncodeError),
     // ValueDecodeError(ValueDecodeError),
 }