diff --git a/bemani/format/afp/decompile.py b/bemani/format/afp/decompile.py index 54d25fc..9a79b73 100644 --- a/bemani/format/afp/decompile.py +++ b/bemani/format/afp/decompile.py @@ -1060,7 +1060,7 @@ class ByteCodeDecompiler(VerboseOutput): current_action = i next_action = i + 1 - if action.opcode in [AP2Action.THROW, AP2Action.RETURN]: + if action.opcode in [AP2Action.THROW, AP2Action.RETURN, AP2Action.END]: # This should end execution, so we should cap off the current execution # and send it to the end. current_action_flow = find(current_action) @@ -1290,7 +1290,7 @@ class ByteCodeDecompiler(VerboseOutput): # We haven't done any fixing up, we're guaranteed this is an AP2Action. last_action = cast(AP2Action, chunk.actions[-1]) - if last_action.opcode in [AP2Action.THROW, AP2Action.RETURN, AP2Action.JUMP] and len(chunk.next_chunks) != 1: + if last_action.opcode in [AP2Action.THROW, AP2Action.RETURN, AP2Action.JUMP, AP2Action.END] and len(chunk.next_chunks) != 1: raise Exception(f"Chunk ID {chunk.id} has control flow action expecting one next chunk but has {len(chunk.next_chunks)}!") if len(chunk.next_chunks) == 2 and last_action.opcode != AP2Action.IF: raise Exception(f"Chunk ID {chunk.id} has two next chunks but control flow action is not an if statement!") @@ -1487,7 +1487,7 @@ class ByteCodeDecompiler(VerboseOutput): if isinstance(last_action, AP2Action): if last_action.opcode == AP2Action.IF and len(chunk.next_chunks) != 2: raise Exception(f"Somehow messed up the next pointers on if statement in chunk ID {chunk.id}!") - if last_action.opcode in [AP2Action.JUMP, AP2Action.RETURN, AP2Action.THROW] and len(chunk.next_chunks) != 1: + if last_action.opcode in [AP2Action.JUMP, AP2Action.RETURN, AP2Action.THROW, AP2Action.END] and len(chunk.next_chunks) != 1: raise Exception(f"Somehow messed up the next pointers on control flow statement in chunk ID {chunk.id}!") else: if len(chunk.next_chunks) > 1: @@ -1628,7 +1628,7 @@ class ByteCodeDecompiler(VerboseOutput): # Examine the last instruction. last_action = chunk.actions[-1] if isinstance(last_action, AP2Action): - if last_action.opcode in [AP2Action.THROW, AP2Action.RETURN]: + if last_action.opcode in [AP2Action.THROW, AP2Action.RETURN, AP2Action.END]: # The last action already dictates what we should do here. Break # the chain at this point. self.vprint(f"Breaking chain on {chunk.id} because it is a {last_action}.") diff --git a/bemani/tests/test_afp_decompile.py b/bemani/tests/test_afp_decompile.py index eaa7430..8d0c195 100644 --- a/bemani/tests/test_afp_decompile.py +++ b/bemani/tests/test_afp_decompile.py @@ -235,6 +235,25 @@ class TestAFPControlGraph(ExtendedTestCase): self.assertEqual(self.__equiv(chunks_by_id[0]), ["100: STOP", "101: RETURN"]) self.assertEqual(self.__equiv(chunks_by_id[1]), []) + def test_dead_code_elimination_end(self) -> None: + # Return case + bytecode = self.__make_bytecode([ + AP2Action(100, AP2Action.STOP), + AP2Action(101, AP2Action.END), + AP2Action(102, AP2Action.END), + ]) + chunks_by_id, offset_map = self.__call_graph(bytecode) + self.assertEqual(offset_map, {100: 0, 103: 1}) + self.assertItemsEqual(chunks_by_id.keys(), {0, 1}) + self.assertItemsEqual(chunks_by_id[0].previous_chunks, []) + self.assertItemsEqual(chunks_by_id[0].next_chunks, [1]) + self.assertItemsEqual(chunks_by_id[1].previous_chunks, [0]) + self.assertItemsEqual(chunks_by_id[1].next_chunks, []) + + # Also verify the code + self.assertEqual(self.__equiv(chunks_by_id[0]), ["100: STOP", "101: END"]) + self.assertEqual(self.__equiv(chunks_by_id[1]), []) + def test_dead_code_elimination_throw(self) -> None: # Throw case bytecode = self.__make_bytecode([ @@ -467,3 +486,34 @@ class TestAFPControlGraph(ExtendedTestCase): self.assertEqual(self.__equiv(chunks_by_id[6]), ["112: PUSH\n 'd'\nEND_PUSH"]) self.assertEqual(self.__equiv(chunks_by_id[7]), ["113: END"]) self.assertEqual(self.__equiv(chunks_by_id[8]), []) + + def test_if_handling_diamond_end_both_sides(self) -> None: + # If true-false diamond case but the cases never converge. + bytecode = self.__make_bytecode([ + # Beginning of the if statement. + PushAction(100, [True]), + IfAction(101, IfAction.IS_TRUE, 104), + # False case (fall through from if). + PushAction(102, ['b']), + AP2Action(103, AP2Action.END), + # True case. + PushAction(104, ['a']), + AP2Action(105, AP2Action.END), + ]) + chunks_by_id, offset_map = self.__call_graph(bytecode) + self.assertEqual(offset_map, {100: 0, 102: 1, 104: 2, 106: 3}) + self.assertItemsEqual(chunks_by_id.keys(), {0, 1, 2, 3}) + self.assertItemsEqual(chunks_by_id[0].previous_chunks, []) + self.assertItemsEqual(chunks_by_id[0].next_chunks, [1, 2]) + self.assertItemsEqual(chunks_by_id[1].previous_chunks, [0]) + self.assertItemsEqual(chunks_by_id[1].next_chunks, [3]) + self.assertItemsEqual(chunks_by_id[2].previous_chunks, [0]) + self.assertItemsEqual(chunks_by_id[2].next_chunks, [3]) + self.assertItemsEqual(chunks_by_id[3].previous_chunks, [1, 2]) + self.assertItemsEqual(chunks_by_id[3].next_chunks, []) + + # Also verify the code + self.assertEqual(self.__equiv(chunks_by_id[0]), ["100: PUSH\n True\nEND_PUSH", "101: IF, Comparison: IS TRUE, Offset To Jump To If True: 104"]) + self.assertEqual(self.__equiv(chunks_by_id[1]), ["102: PUSH\n 'b'\nEND_PUSH", "103: END"]) + self.assertEqual(self.__equiv(chunks_by_id[2]), ["104: PUSH\n 'a'\nEND_PUSH", "105: END"]) + self.assertEqual(self.__equiv(chunks_by_id[3]), [])