feat: const-fold associative operations

Author: boriel

src/api/optimize.py

@@ -220,6 +220,46 @@ def visit_ADDRESS(self, node):
         yield node
 
     def visit_BINARY(self, node: symbols.BINARY):
+        if self.O_LEVEL > 1 and node.operator in ("PLUS", "MUL"):
+            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):
+                    left = node.left.left
+                    ll = node.left.right
+                elif chk.is_number(node.left.left):
+                    left = node.left.right
+                    ll = node.left.left
+
+                if left is not None:
+                    right = yield symbols.BINARY.make_node(
+                        operator=node.operator,
+                        left=ll,
+                        right=node.right,
+                        lineno=node.lineno,
+                        func=node.func,
+                    )
+                    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
         # Retry folding
         yield symbols.BINARY.make_node(node.operator, node.left, node.right, node.lineno, node.func, node.type_)

src/symbols/binary.py

@@ -144,4 +144,4 @@ def make_node(cls, operator, left, right, lineno, func=None, type_=None):
             else:
                 type_ = c_type
 
-        return cls(operator, a, b, type_=type_, lineno=lineno)
+        return cls(operator, a, b, type_=type_, lineno=lineno, func=func)

tests/functional/zx48k/opt2_assoc.asm

@@ -0,0 +1,108 @@
+	org 32768
+.core.__START_PROGRAM:
+	di
+	push ix
+	push iy
+	exx
+	push hl
+	exx
+	ld hl, 0
+	add hl, sp
+	ld (.core.__CALL_BACK__), hl
+	ei
+	jp .core.__MAIN_PROGRAM__
+.core.__CALL_BACK__:
+	DEFW 0
+.core.ZXBASIC_USER_DATA:
+	; Defines USER DATA Length in bytes
+.core.ZXBASIC_USER_DATA_LEN EQU .core.ZXBASIC_USER_DATA_END - .core.ZXBASIC_USER_DATA
+	.core.__LABEL__.ZXBASIC_USER_DATA_LEN EQU .core.ZXBASIC_USER_DATA_LEN
+	.core.__LABEL__.ZXBASIC_USER_DATA EQU .core.ZXBASIC_USER_DATA
+_a:
+	DEFB 00
+.core.ZXBASIC_USER_DATA_END:
+.core.__MAIN_PROGRAM__:
+	ld a, (_a)
+	add a, 3
+	ld (_a), a
+	ld h, 6
+	call .core.__MUL8_FAST
+	ld (_a), a
+	add a, 3
+	ld (_a), a
+	ld h, 6
+	call .core.__MUL8_FAST
+	ld (_a), a
+	add a, 3
+	ld (_a), a
+	ld h, 6
+	call .core.__MUL8_FAST
+	ld (_a), a
+	add a, 3
+	ld (_a), a
+	ld h, 6
+	call .core.__MUL8_FAST
+	ld (_a), a
+	ld hl, 0
+	ld b, h
+	ld c, l
+.core.__END_PROGRAM:
+	di
+	ld hl, (.core.__CALL_BACK__)
+	ld sp, hl
+	exx
+	pop hl
+	pop iy
+	pop ix
+	exx
+	ei
+	ret
+	;; --- end of user code ---
+#line 1 "/zxbasic/src/arch/zx48k/library-asm/mul8.asm"
+	    push namespace core
+__MUL8:		; Performs 8bit x 8bit multiplication
+	    PROC
+	    ;LOCAL __MUL8A
+	    LOCAL __MUL8LOOP
+	    LOCAL __MUL8B
+	    ; 1st operand (byte) in A, 2nd operand into the stack (AF)
+	    pop hl	; return address
+	    ex (sp), hl ; CALLE convention
+;;__MUL8_FAST: ; __FASTCALL__ entry
+	;;	ld e, a
+	;;	ld d, 0
+	;;	ld l, d
+	;;
+	;;	sla h
+	;;	jr nc, __MUL8A
+	;;	ld l, e
+	;;
+;;__MUL8A:
+	;;
+	;;	ld b, 7
+;;__MUL8LOOP:
+	;;	add hl, hl
+	;;	jr nc, __MUL8B
+	;;
+	;;	add hl, de
+	;;
+;;__MUL8B:
+	;;	djnz __MUL8LOOP
+	;;
+	;;	ld a, l ; result = A and HL  (Truncate to lower 8 bits)
+__MUL8_FAST: ; __FASTCALL__ entry, a = a * h (8 bit mul) and Carry
+	    ld b, 8
+	    ld l, a
+	    xor a
+__MUL8LOOP:
+	    add a, a ; a *= 2
+	    sla l
+	    jp nc, __MUL8B
+	    add a, h
+__MUL8B:
+	    djnz __MUL8LOOP
+	    ret		; result = HL
+	    ENDP
+	    pop namespace
+#line 38 "zx48k/opt2_assoc.bas"
+	END

tests/functional/zx48k/opt2_assoc.bas

@@ -0,0 +1,14 @@
+DIM a as Ubyte
+
+a = (a + 1) + 2
+a = (a * 2) * 3
+
+a = (1 + a) + 2
+a = (2 * a) * 3
+
+
+a = 2 + (a + 1)
+a = 3 * (a * 2)
+
+a = 2 + (1 + a)
+a = 3 * (2 * a)