deimos-lang/dmc-lib/src/ir/ir_variable.rs

use crate::type_info::TypeInfo;
use dvm_lib::instruction::Location;
use std::fmt::{Debug, Display, Formatter};
use std::hash::Hash;
use std::rc::Rc;

pub struct IrVariable {
    descriptor: IrVariableDescriptor,
    type_info: TypeInfo,
}

impl IrVariable {
    pub fn new_vr(name: Rc<str>, block_id: usize, type_info: &TypeInfo) -> Self {
        Self {
            descriptor: IrVariableDescriptor::VirtualRegister(IrVrVariableDescriptor::new(
                name, block_id,
            )),
            type_info: type_info.clone(),
        }
    }

    pub fn new_stack(name: Rc<str>, block_id: usize, type_info: TypeInfo) -> Self {
        Self {
            descriptor: IrVariableDescriptor::Stack(IrStackVariableDescriptor::new(name, block_id)),
            type_info: type_info.clone(),
        }
    }

    pub fn type_info(&self) -> &TypeInfo {
        &self.type_info
    }

    pub fn descriptor(&self) -> &IrVariableDescriptor {
        &self.descriptor
    }

    pub fn descriptor_mut(&mut self) -> &mut IrVariableDescriptor {
        &mut self.descriptor
    }

    pub fn set_descriptor(&mut self, descriptor: IrVariableDescriptor) {
        self.descriptor = descriptor;
    }
}

impl Display for IrVariable {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.descriptor)
    }
}

pub enum IrVariableDescriptor {
    VirtualRegister(IrVrVariableDescriptor),
    Stack(IrStackVariableDescriptor),
}

impl Display for IrVariableDescriptor {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            IrVariableDescriptor::VirtualRegister(vr_variable) => {
                write!(f, "{}", vr_variable)
            }
            IrVariableDescriptor::Stack(stack_variable) => {
                write!(f, "{}", stack_variable)
            }
        }
    }
}

impl IrVariableDescriptor {
    pub fn name(&self) -> &str {
        match self {
            IrVariableDescriptor::VirtualRegister(vr_variable) => vr_variable.name(),
            IrVariableDescriptor::Stack(stack_variable) => stack_variable.name(),
        }
    }

    pub fn name_owned(&self) -> Rc<str> {
        match self {
            IrVariableDescriptor::VirtualRegister(vr_variable) => vr_variable.name_owned(),
            IrVariableDescriptor::Stack(stack_variable) => stack_variable.name_owned(),
        }
    }

    pub fn as_location(&self) -> Location {
        match self {
            IrVariableDescriptor::VirtualRegister(register_variable) => {
                register_variable.as_location()
            }
            IrVariableDescriptor::Stack(stack_variable) => stack_variable.as_location(),
        }
    }
}

#[derive(Clone, Hash, PartialEq, Eq)]
pub struct IrVrVariableDescriptor {
    name: Rc<str>,
    block_id: usize,
    assigned_register: Option<usize>,
}

impl IrVrVariableDescriptor {
    pub fn new(name: Rc<str>, block_id: usize) -> Self {
        Self {
            name,
            block_id,
            assigned_register: None,
        }
    }

    pub fn name(&self) -> &str {
        &self.name
    }

    pub fn name_owned(&self) -> Rc<str> {
        self.name.clone()
    }

    pub fn block_id(&self) -> usize {
        self.block_id
    }

    pub fn set_assigned_register(&mut self, register: usize) {
        self.assigned_register = Some(register);
    }

    pub fn assigned_register(&self) -> usize {
        self.assigned_register.unwrap()
    }

    pub fn as_location(&self) -> Location {
        Location::Register(self.assigned_register.unwrap())
    }
}

impl Display for IrVrVariableDescriptor {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.name)
    }
}

impl Debug for IrVrVariableDescriptor {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.name)
    }
}

pub struct IrStackVariableDescriptor {
    name: Rc<str>,
    block_id: usize,
    offset: Option<isize>,
}

impl IrStackVariableDescriptor {
    pub fn new(name: Rc<str>, block_id: usize) -> Self {
        Self {
            name,
            block_id,
            offset: None,
        }
    }

    pub fn name(&self) -> &str {
        &self.name
    }

    pub fn name_owned(&self) -> Rc<str> {
        self.name.clone()
    }

    pub fn set_offset(&mut self, offset: isize) {
        self.offset = Some(offset);
    }

    pub fn offset(&self) -> isize {
        self.offset.unwrap()
    }

    pub fn as_location(&self) -> Location {
        Location::StackFrameOffset(self.offset.unwrap())
    }
}

impl Display for IrStackVariableDescriptor {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}_{}", self.name, self.block_id)
    }
}