From 259083d31b1595002916937b76dc331d1ef34899 Mon Sep 17 00:00:00 2001
From: Jacob Lifshay <programmerjake@gmail.com>
Date: Fri, 16 Dec 2022 00:22:18 -0800
Subject: [PATCH] attempt to speed up code

---
 .../register_allocator.py                     | 44 ++++++++++++++-----
 src/bigint_presentation_code/util.py          | 21 +++++++++
 2 files changed, 55 insertions(+), 10 deletions(-)

diff --git a/src/bigint_presentation_code/register_allocator.py b/src/bigint_presentation_code/register_allocator.py
index 959e1c3..6732953 100644
--- a/src/bigint_presentation_code/register_allocator.py
+++ b/src/bigint_presentation_code/register_allocator.py
@@ -303,21 +303,45 @@ class MergedSSAVal(Interned):
                             return lhs, rhs
         return None
 
-    @lru_cache(maxsize=None, typed=True)
     def copy_merged(self, lhs_loc, rhs, rhs_loc, copy_relation):
         # type: (Loc | None, MergedSSAVal, Loc | None, _CopyRelation) -> MergedSSAVal
+        retval = self.try_copy_merged(lhs_loc, rhs, rhs_loc, copy_relation)
+        if isinstance(retval, MergedSSAVal):
+            return retval
+        raise retval
+
+    def try_copy_merged(self, lhs_loc,  # type: Loc | None
+                        rhs,  # type: MergedSSAVal
+                        rhs_loc,  # type: Loc | None
+                        copy_relation,  # type: _CopyRelation
+                        ):
+        # type: (...) -> MergedSSAVal | BadMergedSSAVal
         cr_lhs, cr_rhs = copy_relation
         if cr_lhs.ssa_val not in self.ssa_vals:
             cr_lhs, cr_rhs = cr_rhs, cr_lhs
-        lhs_merged = self.with_offset_to_match(
-            cr_lhs.ssa_val, additional_offset=-cr_lhs.reg_idx)
-        if lhs_loc is not None:
-            lhs_merged = lhs_merged.with_loc(lhs_loc)
-        rhs_merged = rhs.with_offset_to_match(
-            cr_rhs.ssa_val, additional_offset=-cr_rhs.reg_idx)
-        if rhs_loc is not None:
-            rhs_merged = rhs_merged.with_loc(rhs_loc)
-        return lhs_merged.merged(rhs_merged).normalized()
+        return self.__try_copy_merged(lhs_loc=lhs_loc, cr_lhs=cr_lhs,
+                                      rhs=rhs, rhs_loc=rhs_loc, cr_rhs=cr_rhs)
+
+    @lru_cache(maxsize=None, typed=True)
+    def __try_copy_merged(self, lhs_loc,  # type: Loc | None
+                          cr_lhs,  # type: SSAValSubReg
+                          rhs,  # type: MergedSSAVal
+                          rhs_loc,  # type: Loc | None
+                          cr_rhs,  # type: SSAValSubReg
+                          ):
+        # type: (...) -> MergedSSAVal | BadMergedSSAVal
+        try:
+            lhs_merged = self.with_offset_to_match(
+                cr_lhs.ssa_val, additional_offset=-cr_lhs.reg_idx)
+            if lhs_loc is not None:
+                lhs_merged = lhs_merged.with_loc(lhs_loc)
+            rhs_merged = rhs.with_offset_to_match(
+                cr_rhs.ssa_val, additional_offset=-cr_rhs.reg_idx)
+            if rhs_loc is not None:
+                rhs_merged = rhs_merged.with_loc(rhs_loc)
+            return lhs_merged.merged(rhs_merged).normalized()
+        except BadMergedSSAVal as e:
+            return e
 
 
 @final
diff --git a/src/bigint_presentation_code/util.py b/src/bigint_presentation_code/util.py
index 9d2b7b1..9f1e5ab 100644
--- a/src/bigint_presentation_code/util.py
+++ b/src/bigint_presentation_code/util.py
@@ -7,6 +7,7 @@ from bigint_presentation_code.type_util import Self, final
 
 _T_co = TypeVar("_T_co", covariant=True)
 _T = TypeVar("_T")
+_T2 = TypeVar("_T2")
 
 __all__ = [
     "BaseBitSet",
@@ -141,6 +142,18 @@ class OSet(MutableSet[_T]):
         # type: (_T) -> None
         self.__items.pop(value, None)
 
+    def remove(self, value):
+        # type: (_T) -> None
+        del self.__items[value]
+
+    def pop(self):
+        # type: () -> _T
+        return self.__items.popitem()[0]
+
+    def clear(self):
+        # type: () -> None
+        self.__items.clear()
+
     def __repr__(self):
         # type: () -> str
         if len(self) == 0:
@@ -201,6 +214,14 @@ class FMap(Mapping[_T, _T_co], Interned):
         # type: () -> str
         return f"FMap({self.__items})"
 
+    def get(self, key, default=None):
+        # type: (_T, _T_co | _T2) -> _T_co | _T2
+        return self.__items.get(key, default)
+
+    def __contains__(self, key):
+        # type: (_T | object) -> bool
+        return key in self.__items
+
 
 def trailing_zero_count(v, default=-1):
     # type: (int, int) -> int
-- 
2.30.2