register_allocator2.py works!

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

diff --git a/src/bigint_presentation_code/register_allocator2.py b/src/bigint_presentation_code/register_allocator2.py
index d3ca3983c9dfb114fb0a7454484e9a6c3688a200..aefa2914f2cc58ad6a55994f04fbd113cf6758d8 100644 (file)
--- a/src/bigint_presentation_code/register_allocator2.py
+++ b/src/bigint_presentation_code/register_allocator2.py
@@ -58,7 +58,7 @@ class BadMergedSSAVal(ValueError):
     pass
 
 
-@plain_data(frozen=True)
+@plain_data(frozen=True, repr=False)
 @final
 class MergedSSAVal:
     """a set of `SSAVal`s along with their offsets, all register allocated as
@@ -236,6 +236,10 @@ class MergedSSAVal:
             stop = max(stop, live_range.stop)
         return ProgramRange(start=start, stop=stop)
 
+    def __repr__(self):
+        return (f"MergedSSAVal({self.fn_analysis}, "
+                f"ssa_val_offsets={self.ssa_val_offsets})")
+
 
 @final
 class SSAValToMergedSSAValMap(Mapping[SSAVal, MergedSSAVal]):
@@ -354,13 +358,15 @@ class MergedSSAValToIGNodeMap(Mapping[MergedSSAVal, "IGNode"]):
             self.__merged_ssa_val_map[ssa_val] = final_merged_ssa_val
         return retval
 
-    def __repr__(self):
-        # type: () -> str
-        s = ",\n".join(repr(v) for v in self.__map.values())
+    def __repr__(self, repr_state=None):
+        # type: (None | IGNodeReprState) -> str
+        if repr_state is None:
+            repr_state = IGNodeReprState()
+        s = ",\n".join(v.__repr__(repr_state) for v in self.__map.values())
         return f"MergedSSAValToIGNodeMap({{{s}}})"
 
 
-@plain_data(frozen=True)
+@plain_data(frozen=True, repr=False)
 @final
 class InterferenceGraph:
     __slots__ = "fn_analysis", "merged_ssa_val_map", "nodes"
@@ -400,6 +406,23 @@ class InterferenceGraph:
                     retval.merge(out.tied_input.ssa_val, out)
         return retval
 
+    def __repr__(self, repr_state=None):
+        # type: (None | IGNodeReprState) -> str
+        if repr_state is None:
+            repr_state = IGNodeReprState()
+        s = self.nodes.__repr__(repr_state)
+        return f"InterferenceGraph(nodes={s}, <...>)"
+
+
+@plain_data(repr=False)
+class IGNodeReprState:
+    __slots__ = "node_ids", "did_full_repr"
+
+    def __init__(self):
+        super().__init__()
+        self.node_ids = {}  # type: dict[IGNode, int]
+        self.did_full_repr = OSet()  # type: OSet[IGNode]
+
 
 @final
 class IGNode:
@@ -427,17 +450,20 @@ class IGNode:
         # type: () -> int
         return hash(self.merged_ssa_val)
 
-    def __repr__(self, nodes=None):
-        # type: (None | dict[IGNode, int]) -> str
-        if nodes is None:
-            nodes = {}
-        if self in nodes:
-            return f"<IGNode #{nodes[self]}>"
-        nodes[self] = len(nodes)
-        edges = "{" + ", ".join(i.__repr__(nodes) for i in self.edges) + "}"
-        return (f"IGNode(#{nodes[self]}, "
+    def __repr__(self, repr_state=None, short=False):
+        # type: (None | IGNodeReprState, bool) -> str
+        if repr_state is None:
+            repr_state = IGNodeReprState()
+        node_id = repr_state.node_ids.get(self, None)
+        if node_id is None:
+            repr_state.node_ids[self] = node_id = len(repr_state.node_ids)
+        if short or self in repr_state.did_full_repr:
+            return f"<IGNode #{node_id}>"
+        repr_state.did_full_repr.add(self)
+        edges = ", ".join(i.__repr__(repr_state, True) for i in self.edges)
+        return (f"IGNode(#{node_id}, "
                 f"merged_ssa_val={self.merged_ssa_val}, "
-                f"edges={edges}, "
+                f"edges={{{edges}}}, "
                 f"loc={self.loc})")
 
     @property
@@ -460,6 +486,19 @@ class AllocationFailedError(Exception):
         self.node = node
         self.interference_graph = interference_graph
 
+    def __repr__(self, repr_state=None):
+        # type: (None | IGNodeReprState) -> str
+        if repr_state is None:
+            repr_state = IGNodeReprState()
+        return (f"{__class__.__name__}({self.args[0]!r}, "
+                f"node={self.node.__repr__(repr_state, True)}, "
+                f"interference_graph="
+                f"{self.interference_graph.__repr__(repr_state)})")
+
+    def __str__(self):
+        # type: () -> str
+        return self.__repr__()
+
 
 def allocate_registers(fn):
     # type: (Fn) -> dict[SSAVal, Loc]