From c29d4cd9ec8a41eb163f2da0e7b29da9ad943342 Mon Sep 17 00:00:00 2001
From: Antoine Viallon <antoine@lesviallon.fr>
Date: Fri, 5 Jan 2024 23:29:53 +0100
Subject: [PATCH] optimizations/register_allocation: init

---
 compiler/ir.py                                |  3 +
 compiler/optimizations/__init__.py            |  1 +
 compiler/optimizations/register_allocation.py | 58 +++++++++++++++++++
 3 files changed, 62 insertions(+)
 create mode 100644 compiler/optimizations/__init__.py
 create mode 100644 compiler/optimizations/register_allocation.py

diff --git a/compiler/ir.py b/compiler/ir.py
index 711e4e0..bf79b4d 100644
--- a/compiler/ir.py
+++ b/compiler/ir.py
@@ -152,6 +152,7 @@ class IRRegister(IRAssignable):
     def __init__(self, location: SourceLocation):
         super().__init__(location)
         self.id = IRRegister.register_id
+        self.real = None
         IRRegister.register_id += 1
         IRRegister.registers[self.id] = self
 
@@ -162,6 +163,8 @@ class IRRegister(IRAssignable):
         del IRRegister.registers[self.id]
 
     def codegen(self):
+        if self.real is not None:
+            return f"%r{self.real}"
         return f"%r{self.id}"
 
 
diff --git a/compiler/optimizations/__init__.py b/compiler/optimizations/__init__.py
new file mode 100644
index 0000000..fc5582d
--- /dev/null
+++ b/compiler/optimizations/__init__.py
@@ -0,0 +1 @@
+from .register_allocation import RegisterAllocation
diff --git a/compiler/optimizations/register_allocation.py b/compiler/optimizations/register_allocation.py
new file mode 100644
index 0000000..c07bba2
--- /dev/null
+++ b/compiler/optimizations/register_allocation.py
@@ -0,0 +1,58 @@
+from .. import ir
+from ..errors import CompilationError
+from ..logger import Logger
+
+logger = Logger(__name__)
+
+
+class RegisterAllocation:
+    """
+    Rules:
+    Read After Write: write (a) MUST be before read (b)
+    Read After Read: order does not matter
+    Write After Read: write (b) MUST be after read (a)
+    Write After Write: order should be preserved
+
+    This means we must allocate a new register on its first u
+    """
+
+    def __init__(self, intermediate_representation: list[ir.IRItem]):
+        self.intermediate_representation = intermediate_representation
+        self.used_registers: dict[int, bool] = {}
+
+        for register in ir.IRRegister.registers:
+            self.used_registers[register] = True
+
+    def alloc_register(self) -> int:
+        for register, free in self.used_registers.items():
+            if free:
+                self.used_registers[register] = False
+                return register
+
+    def free_register(self, identifier: int):
+        self.used_registers[identifier] = True
+
+    def analyze(self):
+        for i, item in enumerate(self.intermediate_representation):
+            if not isinstance(item, ir.IRAction):
+                raise CompilationError(item.location, f"Expected an action, got {item}")
+            action: ir.IRAction = item
+
+            for assignable in (action.reads + action.writes):
+                if not isinstance(assignable, ir.IRRegister):
+                    continue
+
+                register: ir.IRRegister = assignable
+
+                sorted_users = register.sorted_users
+
+                # This is the first use, allocate a new register
+                if sorted_users[0].ir_location.line == i:
+                    logger.debug(f"Try to allocate new register for virtual register '{register}'")
+                    new_register = self.alloc_register()
+                    register.real = new_register
+                    logger.debug(f"Allocated new register '{new_register}' for virtual register '{register}'")
+
+                if sorted_users[-1].ir_location.line == i:
+                    self.free_register(register.real)
+                    logger.debug(f"Freed register '{register.real}' (was last used by virt. register '{register}')")