Optimize jumps on jumps

This optimization is done at the basic block levels, not using opt templates.
When a call/jr/jp jumps to another jr/jp with no condition or with the same
condition as its original caller, the labels is updated. So:

jp X
...
X:
jp Y

is converted to
jp Y

Also clean up some iteration over dictionaries.

Author: boriel

arch/zx48k/optimizer/__init__.py

@@ -185,7 +185,6 @@ def optimize(initial_memory):
     if OPTIONS.optimization.value <= 2:  # if -O2 or lower, do nothing and return
         return '\n'.join(x for x in initial_memory if not RE_PRAGMA.match(x))
 
-    init()
     bb = basicblock.BasicBlock(initial_memory)
     cleanup_local_labels(bb)
     initialize_memory(bb)
@@ -210,6 +209,26 @@ def optimize(initial_memory):
     basic_blocks[0].prev = LABELS['*START*'].basic_block
     LABELS[END_PROGRAM_LABEL].basic_block.add_goes_to(LABELS['*__END_PROGRAM*'].basic_block)
 
+    # In O3 we simplify the graph by reducing jumps over jumps
+    for label in JUMP_LABELS:
+        block = LABELS[label].basic_block
+        if isinstance(block, DummyBasicBlock):
+            continue
+
+        # The instruction that starts this block must be one of jr / jp
+        first = block.get_first_non_label_instruction()
+        if first is None or first.inst not in ('jp', 'jr'):
+            continue
+
+        for blk in list(LABELS[label].used_by):
+            if not first.condition_flag or blk[-1].condition_flag == first.condition_flag:
+                new_label = first.opers[0]
+                blk[-1].asm = blk[-1].code.replace(label, new_label)
+                block.delete_comes_from(blk)
+                LABELS[label].used_by.remove(blk)
+                LABELS[new_label].used_by.add(blk)
+                blk.add_goes_to(LABELS[new_label].basic_block)
+
     for x in basic_blocks:
         x.compute_cpu_state()
 

arch/zx48k/optimizer/basicblock.py

@@ -302,10 +302,11 @@ def update_goes_and_comes(self):
             return  # subroutine returns are updated from CALLer blocks
 
         if oper and oper[0]:
-            if oper[0] not in LABELS.keys():
+            if oper[0] not in LABELS:
                 __DEBUG__("INFO: %s is not defined. No optimization is done." % oper[0], 1)
                 LABELS[oper[0]] = LabelInfo(oper[0], 0, DummyBasicBlock(ALL_REGS, ALL_REGS))
 
+            LABELS[oper[0]].used_by.add(self)
             self.add_goes_to(LABELS[oper[0]].basic_block)
 
         if inst in {'djnz', 'jp', 'jr'}:
@@ -651,7 +652,7 @@ def get_basic_blocks(block):
                     block.delete_goes_to(new_block)
 
                 for l in mem.opers:
-                    if l in LABELS.keys():
+                    if l in LABELS:
                         JUMP_LABELS.add(l)
                         block.label_goes.append(l)
                 break

tests/functional/opt3_gotogosub.asm

@@ -0,0 +1,53 @@
+	org 32768
+__START_PROGRAM:
+	di
+	push ix
+	push iy
+	exx
+	push hl
+	exx
+	ld hl, 0
+	add hl, sp
+	ld (__CALL_BACK__), hl
+	ei
+__LABEL__10:
+	call __LABEL__50
+__LABEL__20:
+	ld a, 1
+	ld (_a), a
+__LABEL__25:
+	ld bc, 0
+__END_PROGRAM:
+	di
+	ld hl, (__CALL_BACK__)
+	ld sp, hl
+	exx
+	pop hl
+	pop iy
+	pop ix
+	exx
+	ei
+	ret
+__CALL_BACK__:
+	DEFW 0
+__LABEL__30:
+	jp __LABEL__50
+__LABEL__40:
+	ld a, 2
+	ld (_a), a
+__LABEL__50:
+	ld hl, _a
+	inc (hl)
+__LABEL__60:
+	ret
+	ld bc, 0
+	jp __END_PROGRAM
+
+ZXBASIC_USER_DATA:
+_a:
+	DEFB 00
+	; Defines DATA END --> HEAP size is 0
+ZXBASIC_USER_DATA_END EQU ZXBASIC_MEM_HEAP
+	; Defines USER DATA Length in bytes
+ZXBASIC_USER_DATA_LEN EQU ZXBASIC_USER_DATA_END - ZXBASIC_USER_DATA
+	END

tests/functional/opt3_gotogosub.bas

@@ -0,0 +1,9 @@
+
+10 GOSUB 30
+20 LET a = 1
+25 END
+30 GOTO 50
+40 LET a = 2
+50 LET a = a + 1
+60 RETURN
+

tests/functional/opt3_gotogoto.asm

@@ -0,0 +1,48 @@
+	org 32768
+__START_PROGRAM:
+	di
+	push ix
+	push iy
+	exx
+	push hl
+	exx
+	ld hl, 0
+	add hl, sp
+	ld (__CALL_BACK__), hl
+	ei
+__LABEL__10:
+	jp __LABEL__50
+__LABEL__20:
+	ld a, 1
+	ld (_a), a
+__LABEL__30:
+	jp __LABEL__50
+__LABEL__40:
+	ld a, 2
+	ld (_a), a
+__LABEL__50:
+	ld hl, _a
+	inc (hl)
+	ld bc, 0
+__END_PROGRAM:
+	di
+	ld hl, (__CALL_BACK__)
+	ld sp, hl
+	exx
+	pop hl
+	pop iy
+	pop ix
+	exx
+	ei
+	ret
+__CALL_BACK__:
+	DEFW 0
+
+ZXBASIC_USER_DATA:
+_a:
+	DEFB 00
+	; Defines DATA END --> HEAP size is 0
+ZXBASIC_USER_DATA_END EQU ZXBASIC_MEM_HEAP
+	; Defines USER DATA Length in bytes
+ZXBASIC_USER_DATA_LEN EQU ZXBASIC_USER_DATA_END - ZXBASIC_USER_DATA
+	END

tests/functional/opt3_gotogoto.bas

@@ -0,0 +1,7 @@
+
+10 GOTO 30
+20 LET a = 1
+30 GOTO 50
+40 LET a = 2
+50 LET a = a + 1
+