feat(backend): Add backend

src/backend/register_allocator.rs

@@ -0,0 +1,202 @@
+use std::{cell::RefCell, collections::BTreeMap, fmt::Display, rc::{Rc, Weak}};
+
+use crate::ir::types::Variable;
+
+#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
+pub struct Register {
+    name: &'static str,
+}
+impl Display for Register {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.name)
+    }
+}
+macro_rules! register_declare {
+    ($($reg:ident => $name:expr),*) => {
+        $(
+            pub const $reg: Register = Register { name: $name };
+        )*
+    };
+}
+register_declare! {
+    REG_R0 => "r0",
+    REG_R1 => "r1",
+    REG_R2 => "r2",
+    REG_R3 => "r3",
+    REG_R4 => "r4",
+    REG_R5 => "r5",
+    REG_R6 => "r6",
+    REG_R7 => "r7",
+    REG_R8 => "r8",
+    REG_R9 => "r9",
+    REG_R10 => "r10",
+    REG_R11 => "r11",
+    REG_R12 => "r12",
+    REG_SP => "sp",
+    REG_LR => "lr",
+    REG_PC => "pc",
+    REG_FP => "fp"
+}
+
+pub const REGISTERS: &[Register] = &[
+    REG_R0, REG_R1, REG_R2, REG_R3, REG_R4, REG_R5, REG_R6, REG_R7, REG_R8, REG_R9, REG_R10, REG_R11, REG_R12, REG_SP, REG_LR, REG_PC, REG_FP
+];
+
+pub const REGISTERS_CAN_ALLOC: &[Register] = &[
+    REG_R0, REG_R1, REG_R2, REG_R3, REG_R4, REG_R5, REG_R6, REG_R7, REG_R8, REG_R9, REG_R10, REG_R11, REG_R12
+];
+pub struct RegisterAlloc {
+    allocator: Weak<RefCell<RegisterAllocatorInner>>,
+    pub reg: Register,
+    pub is_reused: bool,
+}
+impl Drop for RegisterAlloc {
+    fn drop(&mut self) {
+        if let Some(allocator) = self.allocator.upgrade() {
+            let mut allocator = allocator.borrow_mut();
+            allocator.mark_unused(self.reg);
+        }
+    }
+}
+pub enum RegisterUseKind {
+    Designated,
+    UsedByVariable(Variable),
+    AllocatedToVariable(Variable),
+    Free,
+}
+struct RegisterAllocatorInner {
+    register_map: BTreeMap<Register, RegisterUseKind>,
+    variable_to_register: BTreeMap<Variable, Register>,
+}
+pub struct RegisterAllocator {
+    // register_map: BTreeMap<Register, RegisterUseKind>,
+    // variable_to_register: BTreeMap<Variable, Register>,
+    inner: Rc<RefCell<RegisterAllocatorInner>>,
+}
+impl RegisterAllocatorInner {
+    fn mark_unused(&mut self, reg: Register) {
+        if let Some(use_kind) = self.register_map.get_mut(&reg) {
+            match use_kind {
+                RegisterUseKind::Designated => {
+                    *use_kind = RegisterUseKind::Free;
+                },
+                RegisterUseKind::UsedByVariable(var) => {
+                    *use_kind = RegisterUseKind::AllocatedToVariable(*var);
+                },
+                _ => panic!("Trying to mark a register as unused that is not in use"),
+            }
+        }
+    }
+
+}
+impl RegisterAllocator {
+    pub fn new() -> Self {
+        let mut register_map = BTreeMap::new();
+        for &reg in REGISTERS_CAN_ALLOC {
+            register_map.insert(reg, RegisterUseKind::Free);
+        }
+        Self {
+            inner: Rc::new(RefCell::new(RegisterAllocatorInner {
+                register_map,
+                variable_to_register: BTreeMap::new(),
+            })),
+        }
+    }
+
+    pub fn alloc(&mut self, var: Variable) -> Option<RegisterAlloc> {
+        let mut inner = self.inner.borrow_mut();
+        if let Some(&reg) = inner.variable_to_register.get(&var) {
+            // Variable already has a register allocated
+            let use_kind = inner.register_map.get_mut(&reg).expect("Inconsistent state: variable has a register but it's not in the register map");
+            assert!(matches!(use_kind, RegisterUseKind::UsedByVariable(v) | RegisterUseKind::AllocatedToVariable(v) if *v == var));
+            *use_kind = RegisterUseKind::UsedByVariable(var);
+            return Some(RegisterAlloc {
+                allocator: Rc::downgrade(&self.inner),
+                reg,
+                is_reused: true,
+            });
+        }
+        // Find a free register
+        for (&reg, use_kind) in inner.register_map.iter_mut() {
+            // Find free register first
+            if let RegisterUseKind::Free = use_kind {
+                *use_kind = RegisterUseKind::UsedByVariable(var);
+                inner.variable_to_register.insert(var, reg);
+                return Some(RegisterAlloc {
+                    allocator: Rc::downgrade(&self.inner),
+                    reg,
+                    is_reused: false,
+                });
+            }
+        }
+        let RegisterAllocatorInner{register_map, variable_to_register} = &mut *inner;
+        for (&reg, use_kind) in register_map.iter_mut() {
+            // Find allocated register then
+            if let RegisterUseKind::AllocatedToVariable(ori_var) = use_kind {
+                assert!(variable_to_register.remove(&ori_var).is_some());
+                *use_kind = RegisterUseKind::UsedByVariable(var);
+                variable_to_register.insert(var, reg);
+                return Some(RegisterAlloc {
+                    allocator: Rc::downgrade(&self.inner),
+                    reg,
+                    is_reused: false,
+                });
+            }
+        }
+        // No free register available
+        None
+    }
+
+    pub fn alloc_reg(&mut self, reg: Register) -> Option<RegisterAlloc> {
+        let mut inner = self.inner.borrow_mut();
+        let RegisterAllocatorInner{register_map, variable_to_register} = &mut *inner;
+        let use_kind = register_map.get_mut(&reg).expect("Trying to allocate a register that is not in the register map");
+        match use_kind {
+            RegisterUseKind::Free => {
+                *use_kind = RegisterUseKind::Designated;
+                return Some(RegisterAlloc {
+                    allocator: Rc::downgrade(&self.inner),
+                    reg,
+                    is_reused: false,
+                });
+            },
+            RegisterUseKind::UsedByVariable(_var) => {
+                return None;
+            },
+            RegisterUseKind::AllocatedToVariable(var) => {
+                variable_to_register.remove(var);
+                *use_kind = RegisterUseKind::Designated;
+                return Some(RegisterAlloc {
+                    allocator: Rc::downgrade(&self.inner),
+                    reg,
+                    is_reused: false,
+                });
+            },
+            RegisterUseKind::Designated => {
+                return None;
+            },
+        }
+
+    }
+
+    pub fn alloc_any(&mut self) -> Option<Register> {
+        let mut inner = self.inner.borrow_mut();
+        
+        for (&reg, use_kind) in inner.register_map.iter_mut() {
+            if let RegisterUseKind::Free = use_kind {
+                *use_kind = RegisterUseKind::Designated;
+                return Some(reg);
+            }
+        }
+        let RegisterAllocatorInner{register_map, variable_to_register} = &mut *inner;
+        for (&reg, use_kind) in register_map.iter_mut() {
+            if let RegisterUseKind::AllocatedToVariable(ori_var) = use_kind {
+                variable_to_register.remove(&ori_var);
+                *use_kind = RegisterUseKind::Designated;
+                return Some(reg);
+            }
+        }
+
+        None
+    }
+}