damn_simple_architecture/compiler/src/backend/dsa/scope.rs

use std::{cell::RefCell, collections::HashMap, ops::Deref, rc::Rc};

use crate::{
    backend::dsa::{
        instruction::{InsBlock, Instruction},
        registers::{Register, RegisterAllocator},
    },
    error,
    model::{CompilerError, Name, TypeId},
};

/// scope object
pub struct Scope<'a> {
    /// outer scope, for a function this will be the global scope.
    parent: Option<&'a mut Scope<'a>>,
    alloc: Rc<RefCell<Allocator>>,

    /// is the scope a function body or just a loop?
    /// depending on the type, ending a scope will have different behaviour
    r#type: ScopeType,

    /// variables
    variables: HashMap<String, Variable>,

    entry_stack_offset: i32,
}

impl<'a> Scope<'a> {
    pub fn new() -> Scope<'a> {
        let alloc = Rc::new(RefCell::new(Allocator::new()));
        let entry_stack_offset = alloc.borrow().get_stack_offset();

        Self {
            alloc,
            entry_stack_offset,
            parent: None,
            r#type: ScopeType::Function,
            variables: HashMap::new(),
        }
    }

    pub fn new_from(parent: &'a mut Scope<'a>, r#type: ScopeType) -> Scope<'a> {
        let alloc = Rc::clone(&parent.alloc);
        let entry_stack_offset = alloc.borrow().get_stack_offset();

        Self {
            alloc,
            entry_stack_offset,
            parent: Some(parent),
            r#type,
            variables: HashMap::new(),
        }
    }

    pub fn create_var(
        &mut self,
        name: String,
        r#type: TypeId,
    ) -> Result<(), CompilerError> {
        let mut var = Variable::new(name, r#type.clone());

        if r#type.size() > 4 {
            let slot = self.alloc.borrow_mut().allocate_stack_slot(r#type.size());
            var.stack_slot = Some(slot);
        } else {
            let reg = self.alloc.borrow_mut().allocate_var()?;
            var.register = Some(reg);
        }

        self.variables.insert(var.name.clone(), var);
        Ok(())
    }

    pub fn alloc_temp(&mut self) -> Result<TempReg, CompilerError> {
        self.alloc.borrow_mut().allocate_temp()
    }

    pub fn free_temp(&mut self, temp: &TempReg) {
        self.alloc.borrow_mut().free_temp(temp)
    }

    pub fn free_var(&mut self, reg: &AssignedReg) {
        self.alloc.borrow_mut().free_var(reg);
    }

    pub fn close(&mut self) {
        // tell the allocator that this scope is closing
        // this reverts the stack offset to what it was before this scope was created.
        self.alloc.clone().borrow_mut().destroy_scope(self);

        for var in self.variables.clone().values() {
            if let Some(reg) = var.register {
                self.free_var(&reg);
            }
        }
    }

    fn get_var(&mut self, var: &str) -> Option<&mut Variable> {
        self.variables.get_mut(var)
    }

    pub fn offset_read(
        &mut self,
        var: &str,
        offset: i32,
    ) -> Result<(TempReg, Instruction), CompilerError> {
        if let Some(var) = self.get_var(var) {
            let slot = var.stack_slot.ok_or_else(|| {
                error("Attempt to read from a var without a stack slot!")
            })?;

            return self.alloc.borrow_mut().offset_read(&slot, offset);
        }

        Err(CompilerError::Undefined(Name::new(var, None)))
    }

    pub fn offset_write(
        &mut self,
        reg: &TempReg,
        var: &str,
        offset: i32,
    ) -> Result<Instruction, CompilerError> {
        if let Some(var) = self.get_var(var) {
            let slot = var.stack_slot.ok_or_else(|| {
                error("Attempt to write to a var without a stack slot!")
            })?;

            return Ok(self.alloc.borrow_mut().offset_write(reg, &slot, offset));
        }

        Err(CompilerError::Undefined(Name::new(var, None)))
    }

    pub fn load_var(
        &mut self,
        var: &str,
    ) -> Result<(AssignedReg, Instruction), CompilerError> {
        if let Some(v) = self.get_var(var).cloned()
            && let Some(slot) = v.stack_slot
        {
            let res = self.alloc.borrow_mut().load_var(&slot)?;
            self.get_var(var).unwrap().register = Some(res.0);
            return Ok(res);
        }

        panic!("e")
    }

    pub fn spill_var(&mut self, var: &str) -> Result<Instruction, CompilerError> {
        if let Some(v) = self.get_var(var).cloned()
            && let Some(rg) = v.register
        {
            let mut slot = v.stack_slot;
            let res = self.alloc.borrow_mut().spill_var(&rg, &mut slot);
            self.get_var(var).unwrap().stack_slot = slot;
            return res;
        }

        Err(CompilerError::Undefined(Name::new(var, None)))
    }
}

impl Drop for Scope<'_> {
    fn drop(&mut self) {
        self.close()
    }
}

#[derive(PartialEq, Copy, Clone, Debug)]
pub enum ScopeType {
    Function,
    IfBlock,
    LoopBlock,
}

#[derive(Clone)]
pub struct Variable {
    pub name: String,

    /// the type of the variable.
    r#type: TypeId,

    /// size taken up in bytes.
    /// if size > 4, value must be stored on the stack.
    pub size: usize,

    pub stack_slot: Option<StackSlot>,
    pub register: Option<AssignedReg>,
}

impl Variable {
    pub fn new(name: String, r#type: TypeId) -> Self {
        Self {
            name,
            size: r#type.size(),
            r#type,
            stack_slot: None,
            register: None,
        }
    }
}

pub struct Allocator {
    stack_offset: i32,
    in_use: [(Register, bool); 16],
}

impl Allocator {
    pub fn new() -> Self {
        let mut in_use = [(Register::Null, false); 16];
        in_use.copy_from_slice(&Register::get_gp().map(|r| (r, false))[0..16]);

        Self {
            stack_offset: 0,
            in_use,
        }
    }

    pub fn get_stack_offset(&self) -> i32 {
        self.stack_offset
    }

    pub fn destroy_scope(&mut self, scope: &mut Scope) {
        self.stack_offset = scope.entry_stack_offset;

        for var in scope.variables.drain() {
            if let Some(assigned) = var.1.register {
                self.free_var(&assigned);
            }
        }
    }

    // what we need:

    // - create var in register from temporary register. free temp and use it.
    //
    // - create var on stack from struct/array literal. return stack offset to write to.
    //
    // - spill var from register to stack. return stack offset to write to.
    //
    // - read/write var from stack+offset into register to use while preserving the stack
    //   slot.
    //
    // - read / write bytes from the stack+offset in a larger variable into a register.

    pub fn offset_read(
        &mut self,
        slot: &StackSlot,
        offset: i32,
    ) -> Result<(TempReg, Instruction), CompilerError> {
        let register = self.allocate_temp()?;

        // instruction: reg = *(&var + offset)
        Ok((
            register.clone(),
            Instruction::ldw_reg_offset(
                Register::Spr,
                *register,
                (**slot + offset) - self.stack_offset,
            ),
        ))
    }

    pub fn offset_write(
        &mut self,
        reg: &TempReg,
        slot: &StackSlot,
        offset: i32,
    ) -> Instruction {
        // instruction: *(&var + offset) = reg
        Instruction::stw_reg_offset(
            **reg,
            Register::Spr,
            (**slot + offset) - self.stack_offset,
        )
    }

    pub fn load_var(
        &mut self,
        slot: &StackSlot,
    ) -> Result<(AssignedReg, Instruction), CompilerError> {
        let reg = self.allocate_var()?;

        Ok((
            reg.clone(),
            Instruction::ldw_reg_offset(Register::Spr, *reg, **slot - self.stack_offset),
        ))
    }

    pub fn spill_var(
        &mut self,
        reg: &AssignedReg,
        slot: &mut Option<StackSlot>,
        // var: &mut Variable,
    ) -> Result<Instruction, CompilerError> {
        if let Some(slot) = &slot {
            let block = Instruction::stw_reg_offset(
                **reg,
                Register::Spr,
                **slot - self.stack_offset,
            );

            self.free_var(reg);
            return Ok(block);
        }

        // var doesn't have a stack slot so we need to create one
        let new_slot = self.allocate_stack_slot(4); // alloc 4 bytes for reg value.
        let block = Instruction::push(**reg);

        self.free_var(reg);
        *slot = Some(new_slot);
        Ok(block)
    }

    pub fn allocate_stack_slot(&mut self, size: usize) -> StackSlot {
        self.stack_offset -= size as i32;
        let offset = self.stack_offset;
        StackSlot(offset)
    }

    pub fn allocate_var(&mut self) -> Result<AssignedReg, CompilerError> {
        if let Some(reg) = self.find_free_register() {
            Ok(AssignedReg(reg))
        } else {
            Err(CompilerError::Generic(
                "No free registers available".to_string(),
            ))
        }
    }

    pub fn allocate_temp(&mut self) -> Result<TempReg, CompilerError> {
        // allocates a temporary register
        if let Some(reg) = self.find_free_register() {
            Ok(TempReg(reg))
        } else {
            todo!("an efficient stack spilling algorithm. needs scope awareness.");
        }
    }

    pub fn free_temp(&mut self, temp: &TempReg) {
        // frees a temporary register.
        self.in_use[**temp as usize].1 = false;
    }

    pub fn free_var(&mut self, reg: &AssignedReg) {
        // frees a register.
        self.in_use[**reg as usize].1 = false;
    }

    // if we have register(s) free, return the first one.
    fn find_free_register(&mut self) -> Option<Register> {
        self.in_use.iter_mut().find_map(|(reg, used)| {
            if !*used {
                *used = true;
                Some(*reg)
            } else {
                None
            }
        })
    }
}

#[derive(Clone, Copy, Debug)]
pub struct TempReg(Register);
#[derive(Clone, Copy, Debug)]
pub struct AssignedReg(Register);
#[derive(Clone, Copy, Debug)]
pub struct StackSlot(i32);

impl Deref for TempReg {
    type Target = Register;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl Deref for AssignedReg {
    type Target = Register;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl Deref for StackSlot {
    type Target = i32;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}