refactor(register): reorganize stuff

src/codegen/codegen.rs

@@ -1,3 +1,4 @@
+use super::{Register, RegisterAllocator};
 use crate::{
     parser::{
         BinOp, Expr, ExprBinary, ExprLit, ExprUnary, IntLitRepr, Stmt, StmtVarDecl, Type, UnOp,
@@ -7,40 +8,26 @@ use crate::{
 use indoc::{formatdoc, writedoc};
 use std::{fmt::Write, fs::File};
 
-struct Register<'a> {
-    name: &'a str,
-    in_use: bool,
-}
-
-impl<'a> Register<'a> {
-    pub fn new(name: &'a str) -> Self {
-        Self {
-            name,
-            in_use: false,
-        }
-    }
-}
-
-pub struct CodeGen<'a> {
+pub struct CodeGen {
     symtable: SymbolTable,
-    registers: [Register<'a>; 6],
+    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: [
-                Register::new("r8"),
-                Register::new("r9"),
-                Register::new("rdi"),
-                Register::new("rsi"),
-                Register::new("rdx"),
-                Register::new("rcx"),
-            ],
+            registers: RegisterAllocator::new(vec![
+                Register::new("", "r8"),
+                Register::new("", "r9"),
+                Register::new("", "rdi"),
+                Register::new("", "rsi"),
+                Register::new("", "rdx"),
+                Register::new("", "rcx"),
+            ]),
             bss_section: "section .bss\n".to_string(),
             data_section: "section .data\n".to_string(),
             text_section: formatdoc!(
@@ -54,22 +41,6 @@ impl<'a> CodeGen<'a> {
         }
     }
 
-    fn alloc(&mut self) -> usize {
-        for (i, r) in self.registers.iter_mut().enumerate() {
-            if !r.in_use {
-                r.in_use = true;
-
-                return i;
-            }
-        }
-
-        panic!("All registers were in use");
-    }
-
-    fn free(&mut self, r: usize) {
-        self.registers[r].in_use = false;
-    }
-
     fn declare(&mut self, var: &StmtVarDecl) {
         let size = self.size(&var.type_);
 
@@ -83,7 +54,7 @@ impl<'a> CodeGen<'a> {
         .unwrap();
     }
 
-    fn expr(&mut self, expr: &Expr) -> usize {
+    fn expr(&mut self, expr: &Expr) -> Register {
         match expr {
             Expr::Binary(bin_expr) => self.bin_expr(bin_expr),
             Expr::Lit(lit) => match lit {
@@ -94,7 +65,7 @@ impl<'a> CodeGen<'a> {
         }
     }
 
-    fn bin_expr(&mut self, expr: &ExprBinary) -> usize {
+    fn bin_expr(&mut self, expr: &ExprBinary) -> Register {
         match &expr.op {
             BinOp::Assign => {
                 let left = expr.left.as_ref();
@@ -103,10 +74,9 @@ impl<'a> CodeGen<'a> {
                     assert!(self.symtable.exists(name));
 
                     let right = self.expr(expr.right.as_ref());
+                    self.mov(name, &right);
 
-                    self.mov(name, right);
-
-                    69
+                    right
                 } else {
                     panic!("Cant assign to non ident");
                 }
@@ -115,25 +85,33 @@ impl<'a> CodeGen<'a> {
                 let left = self.expr(expr.left.as_ref());
                 let right = self.expr(expr.right.as_ref());
 
-                self.add(left, right)
+                self.add(&left, right);
+
+                left
             }
             BinOp::Sub => {
                 let left = self.expr(expr.left.as_ref());
                 let right = self.expr(expr.right.as_ref());
 
-                self.sub(left, right)
+                self.sub(&left, right);
+
+                left
             }
             BinOp::Mul => {
                 let left = self.expr(expr.left.as_ref());
                 let right = self.expr(expr.right.as_ref());
 
-                self.mul(left, right)
+                self.mul(&left, right);
+
+                left
             }
             BinOp::Div => {
                 let left = self.expr(expr.left.as_ref());
                 let right = self.expr(expr.right.as_ref());
 
-                self.div(left, right)
+                self.div(&left, right);
+
+                left
             }
             _ => panic!("lasjdf"),
         }
@@ -154,107 +132,101 @@ impl<'a> CodeGen<'a> {
         }
     }
 
-    fn mov(&mut self, label: &str, r: usize) {
+    fn mov(&mut self, label: &str, r: &Register) {
         writedoc!(
             self.text_section,
             "
             \tmov [{}], {}
             ",
             &label,
-            &self.registers[r].name
+            r.dword(),
         )
         .unwrap();
     }
 
-    fn load(&mut self, int_lit: &IntLitRepr) -> usize {
-        let r = self.alloc();
+    fn load(&mut self, int_lit: &IntLitRepr) -> Register {
+        let r = self.registers.alloc().unwrap();
 
         writedoc!(
             self.text_section,
             "
             \tmov {}, {}
             ",
-            &self.registers[r].name,
+            r.dword(),
             int_lit.to_string(),
         )
         .unwrap();
 
         r
     }
 
-    fn unary_expr(&mut self, unary_expr: &ExprUnary) -> usize {
+    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
             }
             _ => panic!("not unary operator is not supported yet"),
         }
     }
 
-    fn negate(&mut self, r: usize) {
+    fn negate(&mut self, r: &Register) {
         writedoc!(
             self.text_section,
             "
             \tneg {}
             ",
-            &self.registers[r].name,
+            r.dword(),
         )
         .unwrap();
     }
 
-    fn add(&mut self, r1: usize, r2: usize) -> usize {
+    fn add(&mut self, r1: &Register, r2: Register) {
         writedoc!(
             self.text_section,
             "
             \tadd {}, {}
             ",
-            &self.registers[r1].name,
-            &self.registers[r2].name,
+            r1.dword(),
+            r2.dword(),
         )
         .unwrap();
 
-        self.free(r2);
-
-        r1
+        self.registers.free(r2).unwrap();
     }
 
-    fn sub(&mut self, r1: usize, r2: usize) -> usize {
+    fn sub(&mut self, r1: &Register, r2: Register) {
         writedoc!(
             self.text_section,
             "
             \tsub {}, {}
             ",
-            &self.registers[r1].name,
-            &self.registers[r2].name,
+            r1.dword(),
+            r2.dword(),
         )
         .unwrap();
 
-        self.free(r2);
-
-        r1
+        self.registers.free(r2).unwrap();
     }
 
-    fn mul(&mut self, r1: usize, r2: usize) -> usize {
+    fn mul(&mut self, r1: &Register, r2: Register) {
         writedoc!(
             self.text_section,
             "
             \timul {}, {}
             ",
-            &self.registers[r1].name,
-            &self.registers[r2].name,
+            r1.dword(),
+            r2.dword(),
         )
         .unwrap();
 
-        self.free(r2);
-
-        r1
+        self.registers.free(r2).unwrap();
     }
 
     //NOTE: if mafs doesn't works, prolly because of this
-    fn div(&mut self, r1: usize, r2: usize) -> usize {
+    fn div(&mut self, r1: &Register, r2: Register) {
         writedoc!(
             self.text_section,
             "
@@ -263,15 +235,13 @@ impl<'a> CodeGen<'a> {
             \tidiv {}
             \tmov {}, rax
             ",
-            &self.registers[r1].name,
-            &self.registers[r2].name,
-            &self.registers[r1].name,
+            r1.dword(),
+            r2.dword(),
+            r1.dword(),
         )
         .unwrap();
 
-        self.free(r2);
-
-        r1
+        self.registers.free(r2).unwrap();
     }
 
     pub fn compile(&mut self, program: Vec<Stmt>, path: &str) {

src/codegen/mod.rs

@@ -1,3 +1,5 @@
 mod codegen;
+mod register;
 
 pub use codegen::CodeGen;
+pub use register::{Register, RegisterAllocator};

src/codegen/register.rs

@@ -0,0 +1,67 @@
+#[derive(Clone, PartialEq)]
+pub struct Register {
+    byte: &'static str,
+    dword: &'static str,
+}
+
+impl Register {
+    pub fn new(byte: &'static str, dword: &'static str) -> Self {
+        return Self { byte, dword };
+    }
+
+    pub fn byte(&self) -> &'static str {
+        self.byte
+    }
+
+    pub fn dword(&self) -> &'static str {
+        self.dword
+    }
+}
+
+#[derive(Debug)]
+pub enum AllocatorError {
+    DoubleFree,
+    RanOutOfRegisters,
+}
+
+pub struct RegisterAllocator {
+    registers: Vec<Register>,
+    used: Vec<u8>,
+}
+
+impl RegisterAllocator {
+    pub fn new(registers: Vec<Register>) -> Self {
+        return Self {
+            used: Vec::with_capacity(registers.len()),
+            registers,
+        };
+    }
+
+    pub fn alloc(&mut self) -> Result<Register, AllocatorError> {
+        for (i, reg) in self.registers.iter().enumerate() {
+            if !self.used.contains(&i.try_into().unwrap()) {
+                self.used.push(i.try_into().unwrap());
+
+                return Ok(reg.clone());
+            }
+        }
+
+        Err(AllocatorError::RanOutOfRegisters)
+    }
+
+    pub fn free(&mut self, r: Register) -> Result<(), AllocatorError> {
+        for (i, register) in self.registers.iter().enumerate() {
+            if r == *register {
+                if self.used.contains(&i.try_into().unwrap()) {
+                    self.used.remove(i);
+
+                    return Ok(());
+                } else {
+                    return Err(AllocatorError::DoubleFree);
+                }
+            }
+        }
+
+        Ok(())
+    }
+}