fix(register): use enum

src/codegen/codegen.rs

@@ -8,26 +8,19 @@ use crate::{
 use indoc::{formatdoc, writedoc};
 use std::{fmt::Write, fs::File};
 
-pub struct CodeGen<'a> {
+pub struct CodeGen {
     symtable: SymbolTable,
-    registers: RegisterAllocator<'a>,
+    registers: RegisterAllocator,
     data_section: String,
     text_section: String,
     bss_section: String,
 }
 
-impl<'a> CodeGen<'a> {
+impl CodeGen {
     pub fn new(symtable: SymbolTable) -> Self {
         Self {
             symtable,
-            registers: RegisterAllocator::new(vec![
-                Register::new("r8b", "r8"),
-                Register::new("r9b", "r9"),
-                Register::new("dil", "rdi"),
-                Register::new("sil", "rsi"),
-                Register::new("dl", "rdx"),
-                Register::new("cl", "rcx"),
-            ]),
+            registers: RegisterAllocator::new(),
             bss_section: "section .bss\n".to_string(),
             data_section: "section .data\n".to_string(),
             text_section: formatdoc!(
@@ -54,7 +47,7 @@ impl<'a> CodeGen<'a> {
         .unwrap();
     }
 
-    fn expr(&mut self, expr: &Expr) -> &Register {
+    fn expr(&mut self, expr: &Expr) -> Register {
         match expr {
             Expr::Binary(bin_expr) => self.bin_expr(bin_expr),
             Expr::Lit(lit) => match lit {
@@ -65,7 +58,7 @@ impl<'a> CodeGen<'a> {
         }
     }
 
-    fn bin_expr(&mut self, expr: &ExprBinary) -> &Register {
+    fn bin_expr(&mut self, expr: &ExprBinary) -> Register {
         match &expr.op {
             BinOp::Assign => {
                 let left = expr.left.as_ref();
@@ -75,7 +68,7 @@ impl<'a> CodeGen<'a> {
 
                     let right = self.expr(expr.right.as_ref());
 
-                    self.mov(name, right);
+                    self.mov(name, &right);
 
                     todo!();
                 } else {
@@ -137,7 +130,7 @@ impl<'a> CodeGen<'a> {
         .unwrap();
     }
 
-    fn load(&mut self, int_lit: &IntLitRepr) -> &Register {
+    fn load(&mut self, int_lit: &IntLitRepr) -> Register {
         let r = self.registers.alloc().unwrap();
 
         writedoc!(
@@ -153,11 +146,11 @@ impl<'a> CodeGen<'a> {
         r
     }
 
-    fn unary_expr(&mut self, unary_expr: &ExprUnary) -> &Register {
+    fn unary_expr(&mut self, unary_expr: &ExprUnary) -> Register {
         match unary_expr.op {
             UnOp::Negative => {
                 let r = self.expr(unary_expr.expr.as_ref());
-                self.negate(r);
+                self.negate(&r);
 
                 r
             }
@@ -176,7 +169,7 @@ impl<'a> CodeGen<'a> {
         .unwrap();
     }
 
-    fn add(&mut self, r1: &Register, r2: &Register) -> &Register {
+    fn add(&mut self, r1: Register, r2: Register) -> Register {
         writedoc!(
             self.text_section,
             "
@@ -192,7 +185,7 @@ impl<'a> CodeGen<'a> {
         r1
     }
 
-    fn sub(&mut self, r1: &Register, r2: &Register) -> &Register {
+    fn sub(&mut self, r1: Register, r2: Register) -> Register {
         writedoc!(
             self.text_section,
             "
@@ -208,7 +201,7 @@ impl<'a> CodeGen<'a> {
         r1
     }
 
-    fn mul(&mut self, r1: &Register, r2: &Register) -> &Register {
+    fn mul(&mut self, r1: Register, r2: Register) -> Register {
         writedoc!(
             self.text_section,
             "
@@ -225,7 +218,7 @@ impl<'a> CodeGen<'a> {
     }
 
     //NOTE: if mafs doesn't works, prolly because of this
-    fn div(&mut self, r1: &Register, r2: &Register) -> &Register {
+    fn div(&mut self, r1: Register, r2: Register) -> Register {
         writedoc!(
             self.text_section,
             "

src/codegen/register.rs

@@ -1,71 +1,84 @@
-#[derive(PartialEq)]
-pub struct Register<'a> {
-    byte: &'a str,
-    dword: &'a str,
+#[derive(Clone)]
+pub enum Register {
+    R8,
+    R9,
+    Rdi,
+    Rsi,
+    Rdx,
+    Rcx,
 }
 
-impl<'a> Register<'a> {
-    pub fn new(byte: &'a str, dword: &'a str) -> Self {
-        return Self { byte, dword };
-    }
+const REGISTERS_COUNT: u8 = 6;
 
+impl Register {
     pub fn byte(&self) -> &str {
-        self.byte
+        match self {
+            Self::R8 => "r8b",
+            Self::R9 => "r9b",
+            Self::Rdi => "dil",
+            Self::Rsi => "sil",
+            Self::Rdx => "dl",
+            Self::Rcx => "cl",
+        }
     }
 
     pub fn dword(&self) -> &str {
-        self.dword
+        match self {
+            Self::R8 => "r8",
+            Self::R9 => "r9",
+            Self::Rdi => "rdi",
+            Self::Rsi => "rsi",
+            Self::Rdx => "rdx",
+            Self::Rcx => "rcx",
+        }
     }
 }
 
-#[derive(Clone)]
-enum RegisterStatus {
-    Free,
-    InUse,
+impl From<u8> for Register {
+    fn from(value: u8) -> Self {
+        match value {
+            0 => Self::R8,
+            1 => Self::R9,
+            2 => Self::Rdi,
+            3 => Self::Rsi,
+            4 => Self::Rdx,
+            5 => Self::Rcx,
+            _ => unreachable!(),
+        }
+    }
 }
 
 #[derive(Debug)]
-enum AllocatorError {
+pub enum AllocatorError {
     DoubleFree,
     RanOutOfRegisters,
 }
 
-pub struct RegisterAllocator<'a> {
-    registers: Vec<Register<'a>>,
-    status: Vec<RegisterStatus>,
+pub struct RegisterAllocator {
+    used: Vec<u8>,
 }
 
-impl<'a> RegisterAllocator<'a> {
-    pub fn new(registers: Vec<Register<'a>>) -> Self {
-        return Self {
-            status: vec![RegisterStatus::Free; registers.len()],
-            registers,
-        };
+impl RegisterAllocator {
+    pub fn new() -> Self {
+        return Self { used: Vec::new() };
     }
 
-    pub fn alloc(&mut self) -> Result<&Register, AllocatorError> {
-        for (i, status) in self.status.iter_mut().enumerate() {
-            if let RegisterStatus::Free = status {
-                *status = RegisterStatus::InUse;
-
-                return Ok(&self.registers[i]);
+    pub fn alloc(&mut self) -> Result<Register, AllocatorError> {
+        for i in 0..REGISTERS_COUNT {
+            if !self.used.contains(&i) {
+                self.used.push(i);
+                return Ok(Register::from(i));
             }
         }
 
         Err(AllocatorError::RanOutOfRegisters)
     }
 
-    pub fn free(&mut self, r: &Register) -> Result<(), AllocatorError> {
-        for (i, register) in self.registers.iter().enumerate() {
-            if r == register {
-                if let RegisterStatus::Free = self.status[i] {
-                    return Err(AllocatorError::DoubleFree);
-                } else {
-                    self.status[i] = RegisterStatus::Free;
-
-                    return Ok(());
-                }
-            }
+    pub fn free(&mut self, r: Register) -> Result<(), AllocatorError> {
+        if self.used.contains(&(r.clone() as u8)) {
+            self.used.remove(r as usize);
+        } else {
+            return Err(AllocatorError::DoubleFree);
         }
 
         Ok(())