use crate::ir::ir_variable::IrVariableId;
use crate::ir::register_allocation::RegisterAssignment;
use std::collections::HashMap;

pub type StackVariableOffset = isize;

pub enum VariableLocation {
    Register(RegisterAssignment),
    Stack(StackVariableOffset),
}

pub struct VariableLocations {
    register_variables: HashMap<IrVariableId, RegisterAssignment>,
    stack_variables: HashMap<IrVariableId, StackVariableOffset>,
}

impl VariableLocations {
    pub fn new(
        register_variables: HashMap<IrVariableId, RegisterAssignment>,
        stack_variables: HashMap<IrVariableId, StackVariableOffset>,
    ) -> Self {
        Self {
            register_variables,
            stack_variables,
        }
    }

    pub fn get_variable_location(&self, id: &IrVariableId) -> VariableLocation {
        match self.register_variables.get(id) {
            Some(register_assignment) => VariableLocation::Register(*register_assignment),
            None => VariableLocation::Stack(self.stack_variables[id]),
        }
    }

    pub fn register_variables_count(&self) -> usize {
        self.register_variables.len()
    }

    pub fn stack_size(&self) -> usize {
        self.stack_variables.len()
    }
}