Make rzxplay.py let EI block an interrupt when a short frame follows

skoolkid · Apr 1, 2024 · 1e0c91e · 1e0c91e
1 parent d1b8e3e
commit 1e0c91e
Show file tree

Hide file tree

Showing 2 changed files with 62 additions and 9 deletions.
diff --git a/skoolkit/rzxplay.py b/skoolkit/rzxplay.py
@@ -376,15 +376,19 @@ def process_block(block, options, context):
         registers[25] = 0
         fetch_counter = tracer.next_frame()
         if registers[26]:
-            if memory[pc] == 0x76:
+            opcode = memory[pc]
+            if opcode == 0x76:
                 # Advance PC if the CPU was halted
                 registers[24] = (registers[24] + 1) % 65536
-            elif memory[pc] == 0xED and memory[(pc + 1) % 65536] in (0x57, 0x5F):
+                simulator.accept_interrupt(registers, memory, 0)
+            elif opcode == 0xED and memory[(pc + 1) % 65536] in (0x57, 0x5F):
                 # Reset bit 2 of F if the last instruction was LD A,I/R
                 registers[1] &= 0b11111011
-            # Always accept an interrupt at a frame boundary, even if the
-            # instruction just executed would normally block it (e.g. EI)
-            simulator.accept_interrupt(registers, memory, 0)
+                simulator.accept_interrupt(registers, memory, 0)
+            elif opcode != 0xFB or fetch_counter > 2:
+                # Accept an interrupt at a frame boundary, unless the last
+                # instruction was EI and the next frame is short
+                simulator.accept_interrupt(registers, memory, 0)
         if show_progress:
             p = (context.frame_count / total_frames) * 100
             write(f'[{p:5.1f}%]\x08\x08\x08\x08\x08\x08\x08\x08')

diff --git a/tests/test_rzxplay.py b/tests/test_rzxplay.py
@@ -441,7 +441,7 @@ def test_ld_a_r_at_frame_boundary(self):
         """
         self._test_rzx(rzx, exp_output, '--quiet --no-screen', exp_trace)
 
-    def test_ei_does_not_block_interrupt(self):
+    def test_ei_blocks_interrupt_before_frame_with_fetch_count_1(self):
         ram = [0] * 0xC000
         pc = 0xC000
         code = (
@@ -452,13 +452,62 @@ def test_ei_does_not_block_interrupt(self):
         registers = {'PC': pc}
         z80data = self.write_z80_file(None, ram, registers=registers, ret_data=True)
         rzx = RZX()
-        frames = [(1, 0, []), (2, 0, [])]
+        frames = [(1, 0, []), (1, 0, []), (3, 0, [])]
         rzx.add_snapshot(z80data, 'z80', frames)
         exp_output = ''
         exp_trace = """
             F:0 C:00001 I:00000 $C000 EI
-            F:1 C:00002 I:00000 $0038 PUSH AF
-            F:1 C:00001 I:00000 $0039 PUSH HL
+            F:1 C:00001 I:00000 $C001 XOR A
+            F:2 C:00003 I:00000 $0038 PUSH AF
+            F:2 C:00002 I:00000 $0039 PUSH HL
+            F:2 C:00001 I:00000 $003A LD HL,($5C78)
+        """
+        self._test_rzx(rzx, exp_output, '--quiet --no-screen', exp_trace)
+
+    def test_ei_blocks_interrupt_before_frame_with_fetch_count_2(self):
+        ram = [0] * 0xC000
+        pc = 0xC000
+        code = (
+            0xFB, # EI
+            0xAF, # XOR A
+            0xA8, # XOR B
+        )
+        ram[pc - 0x4000:pc - 0x4000 + len(code)] = code
+        registers = {'PC': pc}
+        z80data = self.write_z80_file(None, ram, registers=registers, ret_data=True)
+        rzx = RZX()
+        frames = [(1, 0, []), (2, 0, []), (3, 0, [])]
+        rzx.add_snapshot(z80data, 'z80', frames)
+        exp_output = ''
+        exp_trace = """
+            F:0 C:00001 I:00000 $C000 EI
+            F:1 C:00002 I:00000 $C001 XOR A
+            F:1 C:00001 I:00000 $C002 XOR B
+            F:2 C:00003 I:00000 $0038 PUSH AF
+            F:2 C:00002 I:00000 $0039 PUSH HL
+            F:2 C:00001 I:00000 $003A LD HL,($5C78)
+        """
+        self._test_rzx(rzx, exp_output, '--quiet --no-screen', exp_trace)
+
+    def test_ei_does_not_block_interrupt_before_frame_with_fetch_count_3(self):
+        ram = [0] * 0xC000
+        pc = 0xC000
+        code = (
+            0xFB, # EI
+            0xAF, # XOR A
+        )
+        ram[pc - 0x4000:pc - 0x4000 + len(code)] = code
+        registers = {'PC': pc}
+        z80data = self.write_z80_file(None, ram, registers=registers, ret_data=True)
+        rzx = RZX()
+        frames = [(1, 0, []), (3, 0, [])]
+        rzx.add_snapshot(z80data, 'z80', frames)
+        exp_output = ''
+        exp_trace = """
+            F:0 C:00001 I:00000 $C000 EI
+            F:1 C:00003 I:00000 $0038 PUSH AF
+            F:1 C:00002 I:00000 $0039 PUSH HL
+            F:1 C:00001 I:00000 $003A LD HL,($5C78)
         """
         self._test_rzx(rzx, exp_output, '--quiet --no-screen', exp_trace)