working on rewriting compiler ir to fix reg alloc issues

[bigint-presentation-code.git] / src / bigint_presentation_code / register_allocator.py

diff --git a/src/bigint_presentation_code/register_allocator.py b/src/bigint_presentation_code/register_allocator.py
index b44c32dd76023c948ff9be447d34eb16e02a77bf..cc794e9b926efc906ae8bd6906c1de9bdb7034ce 100644 (file)
--- a/src/bigint_presentation_code/register_allocator.py
+++ b/src/bigint_presentation_code/register_allocator.py
@@ -6,20 +6,14 @@ this uses an algorithm based on:
 """
 
 from itertools import combinations
-from typing import TYPE_CHECKING, Generic, Iterable, Mapping, TypeVar
+from typing import Generic, Iterable, Mapping, TypeVar
 
 from nmutil.plain_data import plain_data
 
 from bigint_presentation_code.compiler_ir import (GPRRangeType, Op, RegClass,
                                                   RegLoc, RegType, SSAVal)
-from bigint_presentation_code.ordered_set import OFSet, OSet
-
-if TYPE_CHECKING:
-    from typing_extensions import final
-else:
-    def final(v):
-        return v
-
+from bigint_presentation_code.type_util import final
+from bigint_presentation_code.util import OFSet, OSet
 
 _RegType = TypeVar("_RegType", bound=RegType)
 
@@ -342,6 +336,13 @@ class AllocationFailed:
         self.interference_graph = interference_graph
 
 
+class AllocationFailedError(Exception):
+    def __init__(self, msg, allocation_failed):
+        # type: (str, AllocationFailed) -> None
+        super().__init__(msg, allocation_failed)
+        self.allocation_failed = allocation_failed
+
+
 def try_allocate_registers_without_spilling(ops):
     # type: (list[Op]) -> dict[SSAVal, RegLoc] | AllocationFailed
 
@@ -422,5 +423,10 @@ def try_allocate_registers_without_spilling(ops):
 
 
 def allocate_registers(ops):
-    # type: (list[Op]) -> None
-    raise NotImplementedError
+    # type: (list[Op]) -> dict[SSAVal, RegLoc]
+    retval = try_allocate_registers_without_spilling(ops)
+    if isinstance(retval, AllocationFailed):
+        # TODO: implement spilling
+        raise AllocationFailedError(
+            "spilling required but not yet implemented", retval)
+    return retval