Merge pull request #597 from boriel/feature/constant_fold_assoc_oper

Feature/constant fold assoc oper

Author: boriel

src/api/optimize.py

@@ -220,7 +220,13 @@ def visit_ADDRESS(self, node):
         yield node
 
     def visit_BINARY(self, node: symbols.BINARY):
+        node = yield self.generic_visit(node)  # This might convert consts to numbers if possible
+
         if self.O_LEVEL > 1 and node.operator in ("PLUS", "MUL"):
+            if chk.is_number(node.left) and not chk.is_number(node.right):
+                node.left, node.right = node.right, node.left
+                node = yield self.generic_visit(node)
+
             if node.left.token == "BINARY" and node.left.operator == node.operator and chk.is_number(node.right):
                 left = ll = None
                 if chk.is_number(node.left.right):
@@ -240,27 +246,31 @@ def visit_BINARY(self, node: symbols.BINARY):
                     )
                     node.left = left
                     node.right = right
-            elif node.right.token == "BINARY" and node.right.operator == node.operator and chk.is_number(node.left):
-                right = rr = None
-                if chk.is_number(node.right.left):
-                    right = node.right.right
-                    rr = node.right.left
-                elif chk.is_number(node.right.right):
-                    right = node.right.left
-                    rr = node.right.right
-
-                if right is not None:
-                    left = yield symbols.BINARY.make_node(
-                        operator=node.operator,
-                        left=node.left,
-                        right=rr,
-                        lineno=node.lineno,
-                        func=node.func,
-                    )
-                    node.left = left
-                    node.right = right
 
-        node = yield self.generic_visit(node)  # This might convert consts to numbers if possible
+            if (
+                node.left.token == node.right.token == "BINARY"
+                and node.operator == node.left.operator == node.right.operator
+                and chk.is_number(node.left.right, node.right.right)
+            ):
+                left = yield symbols.BINARY.make_node(
+                    operator=node.operator,
+                    left=node.left.left,
+                    right=node.right.left,
+                    func=node.left.func,
+                    lineno=node.left.lineno,
+                )
+                right = yield symbols.BINARY.make_node(
+                    operator=node.operator,
+                    left=node.left.right,
+                    right=node.right.right,
+                    func=node.right.func,
+                    lineno=node.right.lineno,
+                )
+
+                node = yield symbols.BINARY.make_node(
+                    operator=node.operator, left=left, right=right, func=node.func, lineno=node.lineno
+                )
+
         # Retry folding
         yield symbols.BINARY.make_node(node.operator, node.left, node.right, node.lineno, node.func, node.type_)
 

tests/functional/zx48k/opt2_assoc.asm

@@ -43,6 +43,11 @@ _a:
 	ld h, 6
 	call .core.__MUL8_FAST
 	ld (_a), a
+	ld hl, (_a - 1)
+	ld a, (_a)
+	add a, h
+	add a, 18
+	ld (_a), a
 	ld hl, 0
 	ld b, h
 	ld c, l
@@ -104,5 +109,5 @@ __MUL8B:
 	    ret		; result = HL
 	    ENDP
 	    pop namespace
-#line 38 "zx48k/opt2_assoc.bas"
+#line 43 "zx48k/opt2_assoc.bas"
 	END

tests/functional/zx48k/opt2_assoc.bas

@@ -12,3 +12,5 @@ a = 3 * (a * 2)
 
 a = 2 + (1 + a)
 a = 3 * (2 * a)
+
+a = 3 + (a + 4) + (5 + a) + 6