damn_simple_architecture/compiler/src/registers.rs

use std::collections::HashMap;

use crate::parser::CompilerError;

/// Register allocator for DSA assembly generation
/// Manages general-purpose registers (rg0-rgf) and handles stack spilling
pub struct RegisterAllocator {
    /// Available general-purpose registers
    available_registers: Vec<String>,

    /// Maps variable names to their current location (register or stack offset)
    variable_locations: HashMap<String, Location>,

    /// Maps registers to the variables they currently hold
    register_contents: HashMap<String, String>,

    /// Current stack offset for local variables (relative to bpr)
    /// Starts at -4 (going downward from base pointer)
    stack_offset: i32,

    /// Track which registers are currently in use
    in_use: HashMap<String, bool>,
}

#[derive(Debug, Clone)]
pub enum Location {
    Register(String),
    Stack(i32), // offset from bpr
}

impl RegisterAllocator {
    pub fn new() -> Self {
        // Initialize with available GP registers (rg0-rgf = 16 registers)
        let registers = vec![
            "rg0", "rg1", "rg2", "rg3", "rg4", "rg5", "rg6", "rg7", "rg8", "rg9", "rga",
            "rgb", "rgc", "rgd", "rge", "rgf",
        ]
        .into_iter()
        .map(String::from)
        .collect();

        RegisterAllocator {
            available_registers: registers,
            variable_locations: HashMap::new(),
            register_contents: HashMap::new(),
            stack_offset: -4, // Start at -4 (first local below saved bpr)
            in_use: HashMap::new(),
        }
    }

    /// Allocate a temporary register for expression evaluation
    /// Returns the register name and optionally assembly code to save it
    pub fn alloc_temp(&mut self) -> Result<(String, Vec<String>), CompilerError> {
        let mut code = Vec::new();

        // Try to find an unused register
        for reg in &self.available_registers {
            if !self.in_use.get(reg).unwrap_or(&false) {
                self.in_use.insert(reg.clone(), true);
                return Ok((reg.clone(), code));
            }
        }

        // All registers in use - need to spill one
        // Choose the first register with a variable we can spill
        // Find a register to spill
        let reg_to_spill = self
            .available_registers
            .iter()
            .find(|reg| self.register_contents.contains_key(*reg))
            .cloned();

        if let Some(reg) = reg_to_spill {
            // Spill this variable to stack
            let spill_code = self.spill_register(&reg)?;
            code.extend(spill_code);

            self.in_use.insert(reg.clone(), true);
            return Ok((reg, code));
        }

        Err(CompilerError::Generic(
            "All registers are used up yet there are no variables to spill to the stack"
                .to_string(),
        ))
    }

    /// Free a temporary register after use
    /// NOTE: This will NOT free registers that contain variables!
    /// Variables persist throughout their scope and must not be freed
    pub fn free_temp(&mut self, reg: &str) {
        // Check if this register contains a variable
        if self.register_contents.contains_key(reg) {
            // This register holds a variable - don't free it!
            // Variables are only freed when they go out of scope via free_var()
            return;
        }

        // This is a true temporary - safe to free
        self.in_use.insert(reg.to_string(), false);
    }

    /// Allocate a register for a named variable
    /// Returns the register and any necessary assembly code
    pub fn alloc_var(
        &mut self,
        var_name: &str,
    ) -> Result<(String, Vec<String>), CompilerError> {
        if let Some(location) = self.variable_locations.get(var_name).cloned() {
            match location {
                Location::Register(reg) => {
                    return Ok((reg.clone(), Vec::new()));
                }
                Location::Stack(offset) => {
                    // Variable was pushed, need to calculate actual position
                    let (reg, mut code) = self.alloc_temp()?;

                    // Load from bpr + offset (offset is negative)
                    code.push(format!("\tsubi bpr {} {}", -(offset + 4), reg));
                    code.push(format!("\tldw {}, {}", reg, reg));

                    // Update location to register
                    self.variable_locations
                        .insert(var_name.to_string(), Location::Register(reg.clone()));
                    self.register_contents
                        .insert(reg.clone(), var_name.to_string());

                    return Ok((reg, code));
                }
            }
        }

        // Variable doesn't have a location yet, allocate a new register
        let (reg, code) = self.alloc_temp()?;
        self.variable_locations
            .insert(var_name.to_string(), Location::Register(reg.clone()));
        self.register_contents
            .insert(reg.clone(), var_name.to_string());

        Ok((reg, code))
    }

    /// Get the current location of a variable
    pub fn get_var_location(&self, var_name: &str) -> Option<&Location> {
        self.variable_locations.get(var_name)
    }

    /// Load a variable into a register (allocating if necessary)
    /// Returns the register and assembly code to load it
    pub fn load_var(
        &mut self,
        var_name: &str,
    ) -> Result<(String, Vec<String>), CompilerError> {
        self.alloc_var(var_name)
    }

    /// Store a value from a register into a variable
    /// Updates tracking and returns any necessary assembly code
    pub fn store_var(&mut self, var_name: &str, source_reg: &str) -> Vec<String> {
        let mut code = Vec::new();

        // Check if variable already has a location
        if let Some(location) = self.variable_locations.get(var_name) {
            match location {
                Location::Register(dest_reg) => {
                    if dest_reg != source_reg {
                        code.push(format!("\tmov {}, {}", source_reg, dest_reg));
                    }
                }
                Location::Stack(offset) => {
                    code.push(format!("\tstw {}, bpr, {}", source_reg, offset));
                }
            }
        } else {
            // Variable doesn't exist yet, we can just use the same reg.

            self.variable_locations.insert(
                var_name.to_string(),
                Location::Register(source_reg.to_string()),
            );
            self.register_contents
                .insert(source_reg.to_string(), var_name.to_string());
            self.in_use.insert(source_reg.to_string(), true);

            // this is not needed for now as if we're storing a var we already have a temp
            // register allocated.
            // if let Some(free_reg) = self.find_free_register() {
            //     if &free_reg != source_reg {
            //         code.push(format!("\tmov {}, {}", source_reg, free_reg));
            //     }
            //     self.variable_locations
            //         .insert(var_name.to_string(),
            // Location::Register(free_reg.clone()));
            //     self.register_contents
            //         .insert(free_reg.clone(), var_name.to_string());
            //     self.in_use.insert(free_reg, true);
            // } else {
            //     // No free registers - allocate on stack
            //     code.push(format!("\tstw {}, bpr, {}", source_reg, self.stack_offset));
            //     self.variable_locations
            //         .insert(var_name.to_string(), Location::Stack(self.stack_offset));
            //     self.stack_offset -= 4; // Move to next stack slot
            // }
        }

        code
    }

    /// Spill a register to the stack
    /// Returns assembly code to perform the spill
    pub fn spill_register(&mut self, reg: &str) -> Result<Vec<String>, CompilerError> {
        let mut code = Vec::new();

        if let Some(var_name) = self.register_contents.get(reg).cloned() {
            // PUSH register to stack (spr decrements automatically)
            code.push(format!("\tpush {}", reg));

            // Track that we pushed one word
            self.stack_offset -= 4;

            // Update variable location - it's now at current spr
            // Note: We track offset from bpr for consistency
            self.variable_locations
                .insert(var_name.clone(), Location::Stack(self.stack_offset));

            // Remove from register tracking
            self.register_contents.remove(reg);
        }

        Ok(code)
    }

    /// Find a free register (not currently in use)
    fn find_free_register(&self) -> Option<String> {
        for reg in &self.available_registers {
            if !self.in_use.get(reg).unwrap_or(&false) {
                return Some(reg.clone());
            }
        }
        None
    }

    /// Spill all registers to stack (useful before function calls)
    pub fn spill_all(&mut self) -> Vec<String> {
        let mut code = Vec::new();

        let regs_to_spill: Vec<String> = self.register_contents.keys().cloned().collect();

        for reg in regs_to_spill {
            if let Ok(spill_code) = self.spill_register(&reg) {
                code.extend(spill_code);
            }
        }

        code
    }

    /// Get the total stack offset
    pub fn get_stack_offset(&self) -> i32 {
        self.stack_offset
    }

    /// Get the total stack space needed for local variables
    pub fn get_stack_size(&self) -> i32 {
        -self.stack_offset // Convert negative offset to positive size
    }

    /// Reset allocator for a new function
    pub fn reset(&mut self) {
        self.variable_locations.clear();
        self.register_contents.clear();
        self.stack_offset = -4;
        self.in_use.clear();
    }

    /// Mark a variable as dead (no longer needed)
    /// Frees its register if it's in one
    pub fn free_var(&mut self, var_name: &str) {
        if let Some(Location::Register(reg)) = self.variable_locations.get(var_name) {
            let reg = reg.clone();
            self.register_contents.remove(&reg);
            self.in_use.insert(reg, false);
        }
        self.variable_locations.remove(var_name);
    }

    /// Get list of registers that contain variables and are in use
    /// These need to be saved before function calls
    pub fn get_caller_saved_registers(&self) -> Vec<String> {
        self.register_contents
            .iter()
            .filter(|(reg, _)| *self.in_use.get(*reg).unwrap_or(&false))
            .map(|(reg, _)| reg.clone())
            .collect()
    }

    /// Save caller-saved registers before a function call
    /// Returns assembly code to save them
    pub fn save_caller_saved(&mut self) -> Vec<String> {
        let mut code = Vec::new();

        // For simplicity, save all currently used registers
        // In a more sophisticated compiler, you'd only save registers that are live
        for (reg, var_name) in self.register_contents.clone() {
            if *self.in_use.get(&reg).unwrap_or(&false) {
                code.push(format!("\tpush {}", reg));
            }
        }

        code
    }

    /// Restore caller-saved registers after a function call
    /// Returns assembly code to restore them
    pub fn restore_caller_saved(&mut self, saved_regs: &[String]) -> Vec<String> {
        let mut code = Vec::new();

        // Restore in reverse order (LIFO)
        for reg in saved_regs.iter().rev() {
            code.push(format!("\tpop {}", reg));
        }

        code
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_basic_allocation() {
        let mut allocator = RegisterAllocator::new();

        let (reg1, code1) = allocator.alloc_temp().unwrap();
        assert_eq!(code1.len(), 0); // No spill needed
        assert_eq!(reg1, "rg0");

        let (reg2, code2) = allocator.alloc_temp().unwrap();
        assert_eq!(code2.len(), 0);
        assert_eq!(reg2, "rg1");

        allocator.free_temp(&reg1);

        let (reg3, code3) = allocator.alloc_temp().unwrap();
        assert_eq!(code3.len(), 0);
        assert_eq!(reg3, "rg0"); // Reuses freed register
    }

    #[test]
    fn test_variable_allocation() {
        let mut allocator = RegisterAllocator::new();

        let (reg, _) = allocator.alloc_var("x").unwrap();
        assert_eq!(reg, "rg0");

        // Requesting same variable again should return same register
        let (reg2, _) = allocator.alloc_var("x").unwrap();
        assert_eq!(reg2, "rg0");
    }

    #[test]
    fn test_stack_allocation() {
        let mut allocator = RegisterAllocator::new();

        // Allocate all 16 registers
        for i in 0..16 {
            allocator.alloc_var(&format!("var{}", i)).unwrap();
        }

        // Next allocation should spill to stack
        let (reg, code) = allocator.alloc_var("var16").unwrap();
        assert!(code.len() > 0); // Should have spill code
    }
}