feature: enhance assoc const folding

Author: boriel

src/api/optimize.py

@@ -220,6 +220,8 @@ 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
@@ -245,7 +247,30 @@ def visit_BINARY(self, node: symbols.BINARY):
                     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