rusty-minic/src/backend/register_allocator.rs

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_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
    }
}