Normalize mul and add further

nickdrozd · Aug 20, 2023 · 1b6764e · 1b6764e
1 parent 6d38bab
commit 1b6764e
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 7 deletions.
diff --git a/test/prog_data.py b/test/prog_data.py
@@ -1466,39 +1466,39 @@
 ALGEBRA_STRINGS: dict[str, tuple[int, str]] = {
     "1RB 1R_  0LC 0LD  1LD 1LC  1RE 1LB  1RF 1RD  0LD 0RA": (
         522,
-        "((1 + (24 * (4 ** 1073741817))) + (8168 * (4 ** 1073741817)))",
+        "(1 + (8192 * (4 ** 1073741817)))",
     ),
     "1RB 0LA  1LC 1LF  0LD 0LC  0LE 0LB  1RE 0RA  1R_ 1LD": (
         2424,
-        "((52202371 + (2266788192 * (64 ** 11524))) + (6496 * ((-64 + (64 ** 11524)) // 63)))",
+        "(52202371 + ((6496 * ((-64 + (64 ** 11524)) // 63)) + (2266788192 * (64 ** 11524))))",
     ),
     "1RB 1R_  1RC 1RA  1RD 0RB  1LE 0RC  0LF 0LD  0LB 1LA": (
         291,
         "((3 + (152 * (2 ** 327672))) + (-8 + (8 * (2 ** 327672))))",
     ),
     "1RB 1RC  1LC 0RF  1RA 0LD  0LC 0LE  1LD 0RA  1RE 1R_": (
         1698,
-        "((0 + (6132 * (4 ** 65530))) + (12 * (4 ** 65530)))",
+        "(0 + (6144 * (4 ** 65530)))",
     ),
     "1RB 1LE  1RC 1RF  1LD 0RB  1RE 0LC  1LA 0RD  1R_ 1RC": (
         5723,
-        "((2126281 + (2125760 * ((-4096 + (4096 ** 5052)) // 4095))) + ((~10^17) * (4096 ** 5052)))",
+        "(2126281 + ((2125760 * ((-4096 + (4096 ** 5052)) // 4095)) + ((~10^17) * (4096 ** 5052))))",
     ),
     "1RB 0LD  1RC 0RF  1LC 1LA  0LE 1R_  1LA 0RB  0RC 0RE": (
         413,
-        "((107 + (9 * (3 ** 22143))) + (63 * ((-3 + (3 ** 22143)) // 2)))",
+        "(107 + ((63 * ((-3 + (3 ** 22143)) // 2)) + (9 * (3 ** 22143))))",
     ),
     "1RB 0LB  0RC 1LB  1RD 0LA  1LE 1LF  1LA 0LD  1R_ 1LE": (
         3067,
         "(~10^463)",
     ),
     "1RB 2LA 1R_ 5LB 5LA 4LB  1LA 4RB 3RB 5LB 1LB 4RA": (
         731,
-        "((12 * (4 ** 8188)) + (4084 * (4 ** 8188)))",
+        "(0 + (4096 * (4 ** 8188)))",
     ),
     "1RB 2LB 4RB 1LA 1RB 1R_  1LA 3RA 5RA 4LB 0RA 4LA": (
         937,
-        "((111 + (80 * ((-4 + (4 ** 1358)) // 3))) + (10896 * (4 ** 1358)))",
+        "(111 + ((80 * ((-4 + (4 ** 1358)) // 3)) + (10896 * (4 ** 1358))))",
     ),
     "1RB 1RA 2LB 3LA  2LA 0LB 1LC 1LB  3RB 3RC 1R_ 1LC": (
         347,

diff --git a/tm/num.py b/tm/num.py
@@ -121,6 +121,11 @@ def __init__(self, l: Count, r: Count):
 
         assert isinstance(r, Num)
 
+        if (isinstance(l, Add)
+                and not isinstance(r, Add)
+                and isinstance(ll := l.l, int)):
+            l, r = ll, l.r + r
+
         super().__init__(l, r)
 
     def __mod__(self, other: int) -> int:
@@ -209,6 +214,12 @@ def __rmul__(self, other: Count) -> Count:
 
         return (self.l * other) * self.rcopy()
 
+    def __add__(self, other: Count) -> Count:
+        if isinstance(other, Mul) and self.r == other.r:
+            return (self.l + other.l) * self.r
+
+        return super().__add__(other)
+
 
 class Div(Num):
     join = '//'