# vim: set fileencoding=utf-8 import unittest from typing import Dict, List, Sequence, Tuple, Union from bemani.tests.helpers import ExtendedTestCase from bemani.format.afp.types.ap2 import AP2Action, IfAction, JumpAction, PushAction, Register from bemani.format.afp.decompile import BitVector, ByteCode, ByteCodeChunk, ControlFlow, ByteCodeDecompiler, Statement class TestAFPBitVector(unittest.TestCase): def test_simple(self) -> None: bv = BitVector(5) self.assertEqual(len(bv), 5) self.assertEqual(bv.bitsSet, set()) bv.setBit(2) self.assertEqual(len(bv), 5) self.assertEqual(bv.bitsSet, {2}) bv.setBit(2) bv.setBit(3) self.assertEqual(len(bv), 5) self.assertEqual(bv.bitsSet, {2, 3}) bv.clearBit(2) bv.clearBit(1) self.assertEqual(len(bv), 5) self.assertEqual(bv.bitsSet, {3}) bv.setAllBitsTo(True) self.assertEqual(len(bv), 5) self.assertEqual(bv.bitsSet, {0, 1, 2, 3, 4}) bv.setAllBitsTo(False) self.assertEqual(len(bv), 5) self.assertEqual(bv.bitsSet, set()) def test_equality(self) -> None: bv1 = BitVector(5, init=True) bv2 = BitVector(5, init=False) self.assertFalse(bv1 == bv2) self.assertTrue(bv1 != bv2) bv2.setAllBitsTo(True) self.assertTrue(bv1 == bv2) self.assertFalse(bv1 != bv2) def test_clone(self) -> None: bv = BitVector(5) bv.setBit(2) bvclone = bv.clone() self.assertTrue(bv == bvclone) bv.setBit(3) bvclone.setBit(4) self.assertEqual(bv.bitsSet, {2, 3}) self.assertEqual(bvclone.bitsSet, {2, 4}) def test_boolean_logic(self) -> None: bv1 = BitVector(5).setBit(2).setBit(3) bv2 = BitVector(5).setBit(1).setBit(2) clone = bv1.clone().orVector(bv2) self.assertEqual(clone.bitsSet, {1, 2, 3}) clone = bv1.clone().andVector(bv2) self.assertEqual(clone.bitsSet, {2}) class TestAFPControlGraph(ExtendedTestCase): # Note that the offsets made up in these test functions are not realistic. Jump/If instructions # take up more than one opcode, and the end offset might be more than one byte past the last # action if that action takes up more than one byte. However, from the perspective of the # decompiler, it doesn't care about accurate sizes, only that the offsets are correct. def test_control_flow(self) -> None: cf = ControlFlow(1, 10, [20]) self.assertTrue(cf.contains(1)) self.assertFalse(cf.contains(10)) self.assertTrue(cf.contains(5)) self.assertFalse(cf.contains(20)) self.assertTrue(cf.is_first(1)) self.assertFalse(cf.is_first(10)) self.assertFalse(cf.is_first(5)) self.assertFalse(cf.is_first(20)) self.assertFalse(cf.is_last(1)) self.assertFalse(cf.is_last(10)) self.assertFalse(cf.is_last(5)) self.assertFalse(cf.is_last(20)) self.assertTrue(cf.is_last(9)) cf1, cf2 = cf.split(5, link=False) self.assertEqual(cf1.beginning, 1) self.assertEqual(cf1.end, 5) self.assertEqual(cf1.next_flow, []) self.assertEqual(cf2.beginning, 5) self.assertEqual(cf2.end, 10) self.assertEqual(cf2.next_flow, [20]) cf3, cf4 = cf.split(5, link=True) self.assertEqual(cf3.beginning, 1) self.assertEqual(cf3.end, 5) self.assertEqual(cf3.next_flow, [5]) self.assertEqual(cf4.beginning, 5) self.assertEqual(cf4.end, 10) self.assertEqual(cf4.next_flow, [20]) def __make_bytecode(self, actions: Sequence[AP2Action]) -> ByteCode: return ByteCode( actions, actions[-1].offset + 1, ) def __call_graph(self, bytecode: ByteCode) -> Tuple[Dict[int, ByteCodeChunk], Dict[int, int]]: # Just create a dummy compiler so we can access the internal method for testing. bcd = ByteCodeDecompiler(bytecode) # Call it, return the data in an easier to test fashion. chunks, offset_map = bcd._ByteCodeDecompiler__graph_control_flow(bytecode) return {chunk.id: chunk for chunk in chunks}, offset_map def __equiv(self, bytecode: Union[ByteCode, ByteCodeChunk, List[AP2Action]]) -> List[str]: if isinstance(bytecode, (ByteCode, ByteCodeChunk)): return [str(x) for x in bytecode.actions] else: return [str(x) for x in bytecode] def test_simple_bytecode(self) -> None: bytecode = self.__make_bytecode([ AP2Action(100, AP2Action.STOP), ]) chunks_by_id, offset_map = self.__call_graph(bytecode) self.assertEqual(offset_map, {100: 0, 101: 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"]) self.assertEqual(self.__equiv(chunks_by_id[1]), []) def test_jump_handling(self) -> None: bytecode = self.__make_bytecode([ JumpAction(100, 102), JumpAction(101, 104), JumpAction(102, 101), JumpAction(103, 106), JumpAction(104, 103), JumpAction(105, 107), JumpAction(106, 105), AP2Action(107, AP2Action.STOP), ]) chunks_by_id, offset_map = self.__call_graph(bytecode) self.assertEqual(offset_map, {100: 0, 101: 1, 102: 2, 103: 3, 104: 4, 105: 5, 106: 6, 107: 7, 108: 8}) self.assertItemsEqual(chunks_by_id.keys(), {0, 1, 2, 3, 4, 5, 6, 7, 8}) self.assertItemsEqual(chunks_by_id[0].previous_chunks, []) self.assertItemsEqual(chunks_by_id[0].next_chunks, [2]) self.assertItemsEqual(chunks_by_id[1].previous_chunks, [2]) self.assertItemsEqual(chunks_by_id[1].next_chunks, [4]) self.assertItemsEqual(chunks_by_id[2].previous_chunks, [0]) self.assertItemsEqual(chunks_by_id[2].next_chunks, [1]) self.assertItemsEqual(chunks_by_id[3].previous_chunks, [4]) self.assertItemsEqual(chunks_by_id[3].next_chunks, [6]) self.assertItemsEqual(chunks_by_id[4].previous_chunks, [1]) self.assertItemsEqual(chunks_by_id[4].next_chunks, [3]) self.assertItemsEqual(chunks_by_id[5].previous_chunks, [6]) self.assertItemsEqual(chunks_by_id[5].next_chunks, [7]) self.assertItemsEqual(chunks_by_id[6].previous_chunks, [3]) self.assertItemsEqual(chunks_by_id[6].next_chunks, [5]) self.assertItemsEqual(chunks_by_id[7].previous_chunks, [5]) self.assertItemsEqual(chunks_by_id[7].next_chunks, [8]) self.assertItemsEqual(chunks_by_id[8].previous_chunks, [7]) self.assertItemsEqual(chunks_by_id[8].next_chunks, []) # Also verify the code self.assertEqual(self.__equiv(chunks_by_id[0]), ["100: JUMP, Offset To Jump To: 102"]) self.assertEqual(self.__equiv(chunks_by_id[1]), ["101: JUMP, Offset To Jump To: 104"]) self.assertEqual(self.__equiv(chunks_by_id[2]), ["102: JUMP, Offset To Jump To: 101"]) self.assertEqual(self.__equiv(chunks_by_id[3]), ["103: JUMP, Offset To Jump To: 106"]) self.assertEqual(self.__equiv(chunks_by_id[4]), ["104: JUMP, Offset To Jump To: 103"]) self.assertEqual(self.__equiv(chunks_by_id[5]), ["105: JUMP, Offset To Jump To: 107"]) self.assertEqual(self.__equiv(chunks_by_id[6]), ["106: JUMP, Offset To Jump To: 105"]) self.assertEqual(self.__equiv(chunks_by_id[7]), ["107: STOP"]) self.assertEqual(self.__equiv(chunks_by_id[8]), []) def test_dead_code_elimination_jump(self) -> None: # Jump case bytecode = self.__make_bytecode([ AP2Action(100, AP2Action.STOP), JumpAction(101, 103), AP2Action(102, AP2Action.PLAY), AP2Action(103, AP2Action.STOP), ]) chunks_by_id, offset_map = self.__call_graph(bytecode) self.assertEqual(offset_map, {100: 0, 103: 1, 104: 2}) self.assertItemsEqual(chunks_by_id.keys(), {0, 1, 2}) 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, [2]) self.assertItemsEqual(chunks_by_id[2].previous_chunks, [1]) self.assertItemsEqual(chunks_by_id[2].next_chunks, []) # Also verify the code self.assertEqual(self.__equiv(chunks_by_id[0]), ["100: STOP", "101: JUMP, Offset To Jump To: 103"]) self.assertEqual(self.__equiv(chunks_by_id[1]), ["103: STOP"]) self.assertEqual(self.__equiv(chunks_by_id[2]), []) def test_dead_code_elimination_return(self) -> None: # Return case bytecode = self.__make_bytecode([ AP2Action(100, AP2Action.STOP), AP2Action(101, AP2Action.RETURN), AP2Action(102, AP2Action.STOP), ]) 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: 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([ PushAction(100, ["exception"]), AP2Action(101, AP2Action.THROW), AP2Action(102, AP2Action.STOP), ]) 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: PUSH\n 'exception'\nEND_PUSH", "101: THROW"]) self.assertEqual(self.__equiv(chunks_by_id[1]), []) def test_if_handling_basic(self) -> None: # If by itself case. bytecode = self.__make_bytecode([ # Beginning of the if statement. PushAction(100, [True]), IfAction(101, IfAction.IS_FALSE, 103), # False case (fall through from if). AP2Action(102, AP2Action.PLAY), # Line after the if statement. AP2Action(103, AP2Action.END), ]) chunks_by_id, offset_map = self.__call_graph(bytecode) self.assertEqual(offset_map, {100: 0, 102: 1, 103: 2, 104: 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, [2]) self.assertItemsEqual(chunks_by_id[2].previous_chunks, [0, 1]) self.assertItemsEqual(chunks_by_id[2].next_chunks, [3]) self.assertItemsEqual(chunks_by_id[3].previous_chunks, [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 FALSE, Offset To Jump To If True: 103"]) self.assertEqual(self.__equiv(chunks_by_id[1]), ["102: PLAY"]) self.assertEqual(self.__equiv(chunks_by_id[2]), ["103: END"]) self.assertEqual(self.__equiv(chunks_by_id[3]), []) def test_if_handling_basic_jump_to_end(self) -> None: # If by itself case. bytecode = self.__make_bytecode([ # Beginning of the if statement. PushAction(100, [True]), IfAction(101, IfAction.IS_FALSE, 103), # False case (fall through from if). AP2Action(102, AP2Action.PLAY), # Some code will jump to the end offset as a way of # "returning" early from a function. ]) chunks_by_id, offset_map = self.__call_graph(bytecode) self.assertEqual(offset_map, {100: 0, 102: 1, 103: 2}) self.assertItemsEqual(chunks_by_id.keys(), {0, 1, 2}) 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, [2]) self.assertItemsEqual(chunks_by_id[2].previous_chunks, [0, 1]) self.assertItemsEqual(chunks_by_id[2].next_chunks, []) # Also verify the code self.assertEqual(self.__equiv(chunks_by_id[0]), ["100: PUSH\n True\nEND_PUSH", "101: IF, Comparison: IS FALSE, Offset To Jump To If True: 103"]) self.assertEqual(self.__equiv(chunks_by_id[1]), ["102: PLAY"]) self.assertEqual(self.__equiv(chunks_by_id[2]), []) def test_if_handling_diamond(self) -> None: # If true-false diamond case. bytecode = self.__make_bytecode([ # Beginning of the if statement. PushAction(100, [True]), IfAction(101, IfAction.IS_TRUE, 104), # False case (fall through from if). AP2Action(102, AP2Action.STOP), JumpAction(103, 105), # True case. AP2Action(104, AP2Action.PLAY), # Line after the if statement. AP2Action(105, AP2Action.END), ]) chunks_by_id, offset_map = self.__call_graph(bytecode) self.assertEqual(offset_map, {100: 0, 102: 1, 104: 2, 105: 3, 106: 4}) self.assertItemsEqual(chunks_by_id.keys(), {0, 1, 2, 3, 4}) 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, [4]) self.assertItemsEqual(chunks_by_id[4].previous_chunks, [3]) self.assertItemsEqual(chunks_by_id[4].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: STOP", "103: JUMP, Offset To Jump To: 105"]) self.assertEqual(self.__equiv(chunks_by_id[2]), ["104: PLAY"]) self.assertEqual(self.__equiv(chunks_by_id[3]), ["105: END"]) self.assertEqual(self.__equiv(chunks_by_id[4]), []) def test_if_handling_diamond_jump_to_end(self) -> None: # If true-false diamond case. bytecode = self.__make_bytecode([ # Beginning of the if statement. PushAction(100, [True]), IfAction(101, IfAction.IS_TRUE, 104), # False case (fall through from if). AP2Action(102, AP2Action.STOP), JumpAction(103, 105), # True case. AP2Action(104, AP2Action.PLAY), ]) chunks_by_id, offset_map = self.__call_graph(bytecode) self.assertEqual(offset_map, {100: 0, 102: 1, 104: 2, 105: 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: STOP", "103: JUMP, Offset To Jump To: 105"]) self.assertEqual(self.__equiv(chunks_by_id[2]), ["104: PLAY"]) self.assertEqual(self.__equiv(chunks_by_id[3]), []) def test_if_handling_diamond_return_to_end(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.RETURN), # True case. PushAction(104, ['a']), AP2Action(105, AP2Action.RETURN), ]) 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: RETURN"]) self.assertEqual(self.__equiv(chunks_by_id[2]), ["104: PUSH\n 'a'\nEND_PUSH", "105: RETURN"]) self.assertEqual(self.__equiv(chunks_by_id[3]), []) def test_if_handling_switch(self) -> None: # Series of ifs (basically a switch statement). bytecode = self.__make_bytecode([ # Beginning of the first if statement. PushAction(100, [Register(0), 1]), IfAction(101, IfAction.NOT_EQUALS, 104), # False case (fall through from if). PushAction(102, ['a']), JumpAction(103, 113), # Beginning of the second if statement. PushAction(104, [Register(0), 2]), IfAction(105, IfAction.NOT_EQUALS, 108), # False case (fall through from if). PushAction(106, ['b']), JumpAction(107, 113), # Beginning of the third if statement. PushAction(108, [Register(0), 3]), IfAction(109, IfAction.NOT_EQUALS, 112), # False case (fall through from if). PushAction(110, ['c']), JumpAction(111, 113), # Beginning of default case. PushAction(112, ['d']), # Line after the switch statement. AP2Action(113, AP2Action.END), ]) chunks_by_id, offset_map = self.__call_graph(bytecode) self.assertEqual(offset_map, {100: 0, 102: 1, 104: 2, 106: 3, 108: 4, 110: 5, 112: 6, 113: 7, 114: 8}) self.assertItemsEqual(chunks_by_id.keys(), {0, 1, 2, 3, 4, 5, 6, 7, 8}) 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, [7]) self.assertItemsEqual(chunks_by_id[2].previous_chunks, [0]) self.assertItemsEqual(chunks_by_id[2].next_chunks, [3, 4]) self.assertItemsEqual(chunks_by_id[3].previous_chunks, [2]) self.assertItemsEqual(chunks_by_id[3].next_chunks, [7]) self.assertItemsEqual(chunks_by_id[4].previous_chunks, [2]) self.assertItemsEqual(chunks_by_id[4].next_chunks, [5, 6]) self.assertItemsEqual(chunks_by_id[5].previous_chunks, [4]) self.assertItemsEqual(chunks_by_id[5].next_chunks, [7]) self.assertItemsEqual(chunks_by_id[6].previous_chunks, [4]) self.assertItemsEqual(chunks_by_id[6].next_chunks, [7]) self.assertItemsEqual(chunks_by_id[7].previous_chunks, [1, 3, 5, 6]) self.assertItemsEqual(chunks_by_id[7].next_chunks, [8]) self.assertItemsEqual(chunks_by_id[8].previous_chunks, [7]) self.assertItemsEqual(chunks_by_id[8].next_chunks, []) # Also verify the code self.assertEqual(self.__equiv(chunks_by_id[0]), ["100: PUSH\n Register(0)\n 1\nEND_PUSH", "101: IF, Comparison: !=, Offset To Jump To If True: 104"]) self.assertEqual(self.__equiv(chunks_by_id[1]), ["102: PUSH\n 'a'\nEND_PUSH", "103: JUMP, Offset To Jump To: 113"]) self.assertEqual(self.__equiv(chunks_by_id[2]), ["104: PUSH\n Register(0)\n 2\nEND_PUSH", "105: IF, Comparison: !=, Offset To Jump To If True: 108"]) self.assertEqual(self.__equiv(chunks_by_id[3]), ["106: PUSH\n 'b'\nEND_PUSH", "107: JUMP, Offset To Jump To: 113"]) self.assertEqual(self.__equiv(chunks_by_id[4]), ["108: PUSH\n Register(0)\n 3\nEND_PUSH", "109: IF, Comparison: !=, Offset To Jump To If True: 112"]) self.assertEqual(self.__equiv(chunks_by_id[5]), ["110: PUSH\n 'c'\nEND_PUSH", "111: JUMP, Offset To Jump To: 113"]) 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]), []) class TestAFPDecompile(ExtendedTestCase): # Note that the offsets made up in these test functions are not realistic. Jump/If instructions # take up more than one opcode, and the end offset might be more than one byte past the last # action if that action takes up more than one byte. However, from the perspective of the # decompiler, it doesn't care about accurate sizes, only that the offsets are correct. def __make_bytecode(self, actions: Sequence[AP2Action]) -> ByteCode: return ByteCode( actions, actions[-1].offset + 1, ) def __call_decompile(self, bytecode: ByteCode) -> List[Statement]: # Just create a dummy compiler so we can access the internal method for testing. bcd = ByteCodeDecompiler(bytecode) bcd.decompile(verbose=self.verbose) return bcd.statements def __equiv(self, statements: List[Statement]) -> List[str]: return [str(x) for x in statements] def test_simple_bytecode(self) -> None: bytecode = self.__make_bytecode([ AP2Action(100, AP2Action.STOP), ]) statements = self.__call_decompile(bytecode) self.assertEqual(self.__equiv(statements), ['builtin_StopPlaying()']) def test_jump_handling(self) -> None: bytecode = self.__make_bytecode([ JumpAction(100, 102), JumpAction(101, 104), JumpAction(102, 101), JumpAction(103, 106), JumpAction(104, 103), JumpAction(105, 107), JumpAction(106, 105), AP2Action(107, AP2Action.STOP), ]) statements = self.__call_decompile(bytecode) self.assertEqual(self.__equiv(statements), ['builtin_StopPlaying()']) def test_dead_code_elimination_jump(self) -> None: # Jump case bytecode = self.__make_bytecode([ AP2Action(100, AP2Action.STOP), JumpAction(101, 103), AP2Action(102, AP2Action.PLAY), AP2Action(103, AP2Action.STOP), ]) statements = self.__call_decompile(bytecode) self.assertEqual(self.__equiv(statements), ['builtin_StopPlaying()', 'builtin_StopPlaying()']) def test_dead_code_elimination_return(self) -> None: # Return case bytecode = self.__make_bytecode([ PushAction(100, ["strval"]), AP2Action(101, AP2Action.RETURN), AP2Action(102, AP2Action.STOP), ]) statements = self.__call_decompile(bytecode) self.assertEqual(self.__equiv(statements), ["return 'strval'"]) 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), ]) statements = self.__call_decompile(bytecode) self.assertEqual(self.__equiv(statements), ['builtin_StopPlaying()']) def test_dead_code_elimination_throw(self) -> None: # Throw case bytecode = self.__make_bytecode([ PushAction(100, ["exception"]), AP2Action(101, AP2Action.THROW), AP2Action(102, AP2Action.STOP), ]) statements = self.__call_decompile(bytecode) self.assertEqual(self.__equiv(statements), ["throw 'exception'"]) def test_if_handling_basic_flow_to_end(self) -> None: # If by itself case. bytecode = self.__make_bytecode([ # Beginning of the if statement. PushAction(100, [True]), IfAction(101, IfAction.IS_FALSE, 103), # False case (fall through from if). AP2Action(102, AP2Action.PLAY), # Line after the if statement. AP2Action(103, AP2Action.END), ]) statements = self.__call_decompile(bytecode) self.assertEqual(self.__equiv(statements), ["if (True) {\n builtin_StartPlaying()\n}"]) def test_if_handling_basic_jump_to_end(self) -> None: # If by itself case. bytecode = self.__make_bytecode([ # Beginning of the if statement. PushAction(100, [True]), IfAction(101, IfAction.IS_FALSE, 103), # False case (fall through from if). AP2Action(102, AP2Action.PLAY), # Some code will jump to the end offset as a way of # "returning" early from a function. ]) statements = self.__call_decompile(bytecode) # TODO: The output should be optimized to remove the early return and move the # start playing section inside the if. self.assertEqual(self.__equiv(statements), ["if (not True) {\n return\n}", "builtin_StartPlaying()"]) def test_if_handling_diamond(self) -> None: # If true-false diamond case. bytecode = self.__make_bytecode([ # Beginning of the if statement. PushAction(100, [True]), IfAction(101, IfAction.IS_TRUE, 104), # False case (fall through from if). AP2Action(102, AP2Action.STOP), JumpAction(103, 105), # True case. AP2Action(104, AP2Action.PLAY), # Line after the if statement. AP2Action(105, AP2Action.END), ]) statements = self.__call_decompile(bytecode) self.assertEqual(self.__equiv(statements), ["if (True) {\n builtin_StartPlaying()\n} else {\n builtin_StopPlaying()\n}"]) def test_if_handling_diamond_jump_to_end(self) -> None: # If true-false diamond case. bytecode = self.__make_bytecode([ # Beginning of the if statement. PushAction(100, [True]), IfAction(101, IfAction.IS_TRUE, 104), # False case (fall through from if). AP2Action(102, AP2Action.STOP), JumpAction(103, 105), # True case. AP2Action(104, AP2Action.PLAY), ]) statements = self.__call_decompile(bytecode) # TODO: The output should be optimized to remove redundant return statements. self.assertEqual(self.__equiv(statements), ["if (True) {\n builtin_StartPlaying()\n return\n} else {\n builtin_StopPlaying()\n return\n}"]) def test_if_handling_diamond_return_to_end(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.RETURN), # True case. PushAction(104, ['a']), AP2Action(105, AP2Action.RETURN), ]) statements = self.__call_decompile(bytecode) self.assertEqual(self.__equiv(statements), ["if (True) {\n return 'a'\n} else {\n return 'b'\n}"]) def test_if_handling_switch(self) -> None: # Series of ifs (basically a switch statement). bytecode = self.__make_bytecode([ # Beginning of the first if statement. PushAction(100, [Register(0), 1]), IfAction(101, IfAction.NOT_EQUALS, 104), # False case (fall through from if). PushAction(102, ['a']), JumpAction(103, 113), # Beginning of the second if statement. PushAction(104, [Register(0), 2]), IfAction(105, IfAction.NOT_EQUALS, 108), # False case (fall through from if). PushAction(106, ['b']), JumpAction(107, 113), # Beginning of the third if statement. PushAction(108, [Register(0), 3]), IfAction(109, IfAction.NOT_EQUALS, 112), # False case (fall through from if). PushAction(110, ['c']), JumpAction(111, 113), # Beginning of default case. PushAction(112, ['d']), # Line after the switch statement. AP2Action(113, AP2Action.RETURN), ]) statements = self.__call_decompile(bytecode) # TODO: This should be optimized as an if/elseif/else chunk without so much indentation. self.assertEqual(self.__equiv(statements), [ "if (registers[0] != 1) {\n" " if (registers[0] != 2) {\n" " if (registers[0] != 3) {\n" " tempvar_0 = 'd'\n" " } else {\n" " tempvar_0 = 'c'\n" " }\n" " } else {\n" " tempvar_0 = 'b'\n" " }\n" "} else {\n" " tempvar_0 = 'a'\n" "}", "return tempvar_0" ]) 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). AP2Action(102, AP2Action.STOP), AP2Action(103, AP2Action.END), # True case. AP2Action(104, AP2Action.PLAY), AP2Action(105, AP2Action.END), ]) statements = self.__call_decompile(bytecode) # TODO: The output should be optimized to remove redundant return statements. self.assertEqual(self.__equiv(statements), ["if (True) {\n builtin_StartPlaying()\n return\n} else {\n builtin_StopPlaying()\n return\n}"]) def test_if_handling_or(self) -> None: # Two ifs that together make an or (if register == 1 or register == 3) bytecode = self.__make_bytecode([ # Beginning of the first if statement. PushAction(100, [Register(0), 1]), IfAction(101, IfAction.EQUALS, 104), # False case (circuit not broken, register is not equal to 1) PushAction(102, [Register(0), 2]), IfAction(103, IfAction.NOT_EQUALS, 106), # This is the true case AP2Action(104, AP2Action.PLAY), JumpAction(105, 107), # This is the false case AP2Action(106, AP2Action.STOP), # This is the fall-through after the if. PushAction(107, ['strval']), AP2Action(108, AP2Action.RETURN), ]) statements = self.__call_decompile(bytecode) # TODO: This should be optimized as a compound if statement. self.assertEqual(self.__equiv(statements), [ "if (registers[0] != 1) {\n" " if (registers[0] != 2) {\n" " builtin_StopPlaying()\n" " label_4:\n" " return 'strval'\n" " }\n" "}", "builtin_StartPlaying()", "goto label_4", ])