# vim: set fileencoding=utf-8 import os import unittest from typing import Dict, List, Sequence, Tuple, Union from bemani.tests.helpers import ExtendedTestCase from bemani.format.afp.decompile import ( ByteCodeDecompiler, ByteCode, BitVector, ByteCodeChunk, ControlFlow, ) from bemani.format.afp.types import ( AP2Action, IfAction, JumpAction, PushAction, AddNumVariableAction, Register, Variable, ArithmeticExpression, UNDEFINED, Statement, DefineLabelStatement, GotoStatement, ReturnStatement, PlayMovieStatement, StopMovieStatement, NextFrameStatement, PreviousFrameStatement, IfStatement, ForStatement, IsUndefinedIf, IsBooleanIf, TwoParameterIf, AndIf, OrIf, ) OPEN_BRACKET = "{" CLOSE_BRACKET = "}" 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( None, 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, optimize=True) # Call it, return the data in an easier to test fashion. chunks, offset_map = bcd._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]), [f"100: PUSH{os.linesep} 'exception'{os.linesep}END_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.COMP_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]), [ f"100: PUSH{os.linesep} True{os.linesep}END_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.COMP_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]), [ f"100: PUSH{os.linesep} True{os.linesep}END_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.COMP_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]), [ f"100: PUSH{os.linesep} True{os.linesep}END_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.COMP_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]), [ f"100: PUSH{os.linesep} True{os.linesep}END_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.COMP_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]), [ f"100: PUSH{os.linesep} True{os.linesep}END_PUSH", "101: IF, Comparison: IS TRUE, Offset To Jump To If True: 104", ], ) self.assertEqual( self.__equiv(chunks_by_id[1]), [f"102: PUSH{os.linesep} 'b'{os.linesep}END_PUSH", "103: RETURN"], ) self.assertEqual( self.__equiv(chunks_by_id[2]), [f"104: PUSH{os.linesep} 'a'{os.linesep}END_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.COMP_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.COMP_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.COMP_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]), [ f"100: PUSH{os.linesep} Register(0){os.linesep} 1{os.linesep}END_PUSH", "101: IF, Comparison: !=, Offset To Jump To If True: 104", ], ) self.assertEqual( self.__equiv(chunks_by_id[1]), [ f"102: PUSH{os.linesep} 'a'{os.linesep}END_PUSH", "103: JUMP, Offset To Jump To: 113", ], ) self.assertEqual( self.__equiv(chunks_by_id[2]), [ f"104: PUSH{os.linesep} Register(0){os.linesep} 2{os.linesep}END_PUSH", "105: IF, Comparison: !=, Offset To Jump To If True: 108", ], ) self.assertEqual( self.__equiv(chunks_by_id[3]), [ f"106: PUSH{os.linesep} 'b'{os.linesep}END_PUSH", "107: JUMP, Offset To Jump To: 113", ], ) self.assertEqual( self.__equiv(chunks_by_id[4]), [ f"108: PUSH{os.linesep} Register(0){os.linesep} 3{os.linesep}END_PUSH", "109: IF, Comparison: !=, Offset To Jump To If True: 112", ], ) self.assertEqual( self.__equiv(chunks_by_id[5]), [ f"110: PUSH{os.linesep} 'c'{os.linesep}END_PUSH", "111: JUMP, Offset To Jump To: 113", ], ) self.assertEqual( self.__equiv(chunks_by_id[6]), [f"112: PUSH{os.linesep} 'd'{os.linesep}END_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.COMP_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]), [ f"100: PUSH{os.linesep} True{os.linesep}END_PUSH", "101: IF, Comparison: IS TRUE, Offset To Jump To If True: 104", ], ) self.assertEqual( self.__equiv(chunks_by_id[1]), [f"102: PUSH{os.linesep} 'b'{os.linesep}END_PUSH", "103: END"], ) self.assertEqual( self.__equiv(chunks_by_id[2]), [f"104: PUSH{os.linesep} 'a'{os.linesep}END_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( None, 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, optimize=True) 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.COMP_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), [f"if (True) {OPEN_BRACKET}{os.linesep} builtin_StartPlaying(){os.linesep}{CLOSE_BRACKET}"], ) 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.COMP_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) self.assertEqual( self.__equiv(statements), [f"if (True) {OPEN_BRACKET}{os.linesep} builtin_StartPlaying(){os.linesep}{CLOSE_BRACKET}"], ) 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.COMP_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), [ f"if (True) {OPEN_BRACKET}{os.linesep} builtin_StartPlaying(){os.linesep}{CLOSE_BRACKET} else {OPEN_BRACKET}{os.linesep} builtin_StopPlaying(){os.linesep}{CLOSE_BRACKET}" ], ) 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.COMP_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) self.assertEqual( self.__equiv(statements), [ f"if (True) {OPEN_BRACKET}{os.linesep} builtin_StartPlaying(){os.linesep}{CLOSE_BRACKET} else {OPEN_BRACKET}{os.linesep} builtin_StopPlaying(){os.linesep}{CLOSE_BRACKET}" ], ) 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.COMP_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), [ f"if (True) {OPEN_BRACKET}{os.linesep} return 'a'{os.linesep}{CLOSE_BRACKET} else {OPEN_BRACKET}{os.linesep} return 'b'{os.linesep}{CLOSE_BRACKET}" ], ) 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.COMP_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.COMP_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.COMP_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) self.assertEqual( self.__equiv(statements), [ f"switch (registers[0]) {OPEN_BRACKET}{os.linesep}" f" case 1:{os.linesep}" f" tempvar_0 = 'a'{os.linesep}" f" break{os.linesep}" f" case 2:{os.linesep}" f" tempvar_0 = 'b'{os.linesep}" f" break{os.linesep}" f" case 3:{os.linesep}" f" tempvar_0 = 'c'{os.linesep}" f" break{os.linesep}" f" default:{os.linesep}" f" tempvar_0 = 'd'{os.linesep}" f" break{os.linesep}" "}", "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.COMP_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) self.assertEqual( self.__equiv(statements), [ f"if (True) {OPEN_BRACKET}{os.linesep} builtin_StartPlaying(){os.linesep}{CLOSE_BRACKET} else {OPEN_BRACKET}{os.linesep} builtin_StopPlaying(){os.linesep}{CLOSE_BRACKET}" ], ) 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.COMP_EQUALS, 104), # False case (circuit not broken, register is not equal to 1) PushAction(102, [Register(0), 2]), IfAction(103, IfAction.COMP_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) self.assertEqual( self.__equiv(statements), [ f"if (registers[0] == 1 || registers[0] == 2) {OPEN_BRACKET}{os.linesep}" f" builtin_StartPlaying(){os.linesep}" f"{CLOSE_BRACKET} else {OPEN_BRACKET}{os.linesep}" f" builtin_StopPlaying(){os.linesep}" f"{CLOSE_BRACKET}", "return 'strval'", ], ) def test_basic_while(self) -> None: # A basic while statement. bytecode = self.__make_bytecode( [ # Define exit condition variable. PushAction(100, ["finished", False]), AP2Action(101, AP2Action.DEFINE_LOCAL), # Check exit condition. PushAction(102, ["finished"]), AP2Action(103, AP2Action.GET_VARIABLE), IfAction(104, IfAction.COMP_IS_TRUE, 107), # Loop code. AP2Action(105, AP2Action.NEXT_FRAME), # Loop finished jump back to beginning. JumpAction(106, 102), # End of loop. AP2Action(107, AP2Action.END), ] ) statements = self.__call_decompile(bytecode) self.assertEqual( self.__equiv(statements), [ "local finished = False", f"while (not finished) {OPEN_BRACKET}{os.linesep}" f" builtin_GotoNextFrame(){os.linesep}" "}", ], ) def test_advanced_while(self) -> None: # A basic while statement. bytecode = self.__make_bytecode( [ # Define exit condition variable. PushAction(100, ["finished", False]), AP2Action(101, AP2Action.DEFINE_LOCAL), # Check exit condition. PushAction(102, ["finished"]), AP2Action(103, AP2Action.GET_VARIABLE), IfAction(104, IfAction.COMP_IS_TRUE, 112), # Loop code with a continue statement. PushAction(105, ["some_condition"]), AP2Action(106, AP2Action.GET_VARIABLE), IfAction(107, IfAction.COMP_IS_FALSE, 110), AP2Action(108, AP2Action.NEXT_FRAME), # Continue statement. JumpAction(109, 102), # Exit early. AP2Action(110, AP2Action.STOP), # Break statement. JumpAction(111, 112), # End of loop. AP2Action(112, AP2Action.END), ] ) statements = self.__call_decompile(bytecode) self.assertEqual( self.__equiv(statements), [ "local finished = False", f"while (not finished) {OPEN_BRACKET}{os.linesep}" f" if (not some_condition) {OPEN_BRACKET}{os.linesep}" f" builtin_StopPlaying(){os.linesep}" f" break{os.linesep}" f" {CLOSE_BRACKET}{os.linesep}" f" builtin_GotoNextFrame(){os.linesep}" "}", ], ) def test_basic_for(self) -> None: # A basic for statement. bytecode = self.__make_bytecode( [ # Define exit condition variable. PushAction(100, ["i", 0]), AP2Action(101, AP2Action.DEFINE_LOCAL), # Check exit condition. PushAction(102, [10, "i"]), AP2Action(103, AP2Action.GET_VARIABLE), IfAction(104, IfAction.COMP_LT_EQUALS, 109), # Loop code. AP2Action(105, AP2Action.NEXT_FRAME), # Increment, also the continue point. PushAction(106, ["i"]), AddNumVariableAction(107, 1), # Loop finished jump back to beginning. JumpAction(108, 102), # End of loop. AP2Action(109, AP2Action.END), ] ) statements = self.__call_decompile(bytecode) self.assertEqual( self.__equiv(statements), [ f"for (local i = 0; i < 10; i = i + 1) {OPEN_BRACKET}{os.linesep}" f" builtin_GotoNextFrame(){os.linesep}" "}" ], ) def test_advanced_for(self) -> None: # A basic for statement. bytecode = self.__make_bytecode( [ # Define exit condition variable. PushAction(100, ["i", 0]), AP2Action(101, AP2Action.DEFINE_LOCAL), # Check exit condition. PushAction(102, [10, "i"]), AP2Action(103, AP2Action.GET_VARIABLE), IfAction(104, IfAction.COMP_LT_EQUALS, 115), # Loop code with a continue statement. PushAction(105, ["some_condition"]), AP2Action(106, AP2Action.GET_VARIABLE), IfAction(107, IfAction.COMP_IS_FALSE, 110), AP2Action(108, AP2Action.NEXT_FRAME), # Continue statement. JumpAction(109, 112), # Exit early. AP2Action(110, AP2Action.STOP), # Break statement. JumpAction(111, 115), # Increment, also the continue point. PushAction(112, ["i"]), AddNumVariableAction(113, 1), # Loop finished jump back to beginning. JumpAction(114, 102), # End of loop. AP2Action(115, AP2Action.END), ] ) statements = self.__call_decompile(bytecode) self.assertEqual( self.__equiv(statements), [ f"for (local i = 0; i < 10; i = i + 1) {OPEN_BRACKET}{os.linesep}" f" if (not some_condition) {OPEN_BRACKET}{os.linesep}" f" builtin_StopPlaying(){os.linesep}" f" break{os.linesep}" f" {CLOSE_BRACKET}{os.linesep}" f" builtin_GotoNextFrame(){os.linesep}" "}" ], ) class TestIfExprs(ExtendedTestCase): def test_simple(self) -> None: self.assertEqual(str(IsUndefinedIf(Variable("a"))), "a is UNDEFINED") self.assertEqual(str(IsBooleanIf(Variable("a"))), "a") self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.EQUALS, Variable("b"))), "a == b", ) self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.NOT_EQUALS, Variable("b"))), "a != b", ) self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.LT, Variable("b"))), "a < b", ) self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.GT, Variable("b"))), "a > b", ) self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.LT_EQUALS, Variable("b"))), "a <= b", ) self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.GT_EQUALS, Variable("b"))), "a >= b", ) self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.STRICT_EQUALS, Variable("b"))), "a === b", ) self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.STRICT_NOT_EQUALS, Variable("b"))), "a !== b", ) self.assertEqual( str( AndIf( TwoParameterIf(Variable("a"), TwoParameterIf.LT, Variable("b")), TwoParameterIf(Variable("c"), TwoParameterIf.GT, Variable("d")), ) ), "a < b && c > d", ) self.assertEqual( str( OrIf( TwoParameterIf(Variable("a"), TwoParameterIf.LT, Variable("b")), TwoParameterIf(Variable("c"), TwoParameterIf.GT, Variable("d")), ) ), "a < b || c > d", ) def test_invert_simple(self) -> None: self.assertEqual(str(IsUndefinedIf(Variable("a")).invert()), "a is not UNDEFINED") self.assertEqual(str(IsBooleanIf(Variable("a")).invert()), "not a") self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.EQUALS, Variable("b")).invert()), "a != b", ) self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.NOT_EQUALS, Variable("b")).invert()), "a == b", ) self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.LT, Variable("b")).invert()), "a >= b", ) self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.GT, Variable("b")).invert()), "a <= b", ) self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.LT_EQUALS, Variable("b")).invert()), "a > b", ) self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.GT_EQUALS, Variable("b")).invert()), "a < b", ) self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.STRICT_EQUALS, Variable("b")).invert()), "a !== b", ) self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.STRICT_NOT_EQUALS, Variable("b")).invert()), "a === b", ) self.assertEqual( str( AndIf( TwoParameterIf(Variable("a"), TwoParameterIf.LT, Variable("b")), TwoParameterIf(Variable("c"), TwoParameterIf.GT, Variable("d")), ).invert() ), "a >= b || c <= d", ) self.assertEqual( str( OrIf( TwoParameterIf(Variable("a"), TwoParameterIf.LT, Variable("b")), TwoParameterIf(Variable("c"), TwoParameterIf.GT, Variable("d")), ).invert(), ), "a >= b && c <= d", ) def test_invert_double(self) -> None: self.assertEqual(str(IsUndefinedIf(Variable("a")).invert().invert()), "a is UNDEFINED") self.assertEqual(str(IsBooleanIf(Variable("a")).invert().invert()), "a") self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.EQUALS, Variable("b")).invert().invert()), "a == b", ) self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.NOT_EQUALS, Variable("b")).invert().invert()), "a != b", ) self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.LT, Variable("b")).invert().invert()), "a < b", ) self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.GT, Variable("b")).invert().invert()), "a > b", ) self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.LT_EQUALS, Variable("b")).invert().invert()), "a <= b", ) self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.GT_EQUALS, Variable("b")).invert().invert()), "a >= b", ) self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.STRICT_EQUALS, Variable("b")).invert().invert()), "a === b", ) self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.STRICT_NOT_EQUALS, Variable("b")).invert().invert()), "a !== b", ) self.assertEqual( str( AndIf( TwoParameterIf(Variable("a"), TwoParameterIf.LT, Variable("b")), TwoParameterIf(Variable("c"), TwoParameterIf.GT, Variable("d")), ) .invert() .invert() ), "a < b && c > d", ) self.assertEqual( str( OrIf( TwoParameterIf(Variable("a"), TwoParameterIf.LT, Variable("b")), TwoParameterIf(Variable("c"), TwoParameterIf.GT, Variable("d")), ) .invert() .invert() ), "a < b || c > d", ) def test_swap_simple(self) -> None: self.assertEqual(str(IsUndefinedIf(Variable("a")).swap()), "a is UNDEFINED") self.assertEqual(str(IsBooleanIf(Variable("a")).swap()), "a") self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.EQUALS, Variable("b")).swap()), "b == a", ) self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.NOT_EQUALS, Variable("b")).swap()), "b != a", ) self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.LT, Variable("b")).swap()), "b > a", ) self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.GT, Variable("b")).swap()), "b < a", ) self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.LT_EQUALS, Variable("b")).swap()), "b >= a", ) self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.GT_EQUALS, Variable("b")).swap()), "b <= a", ) self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.STRICT_EQUALS, Variable("b")).swap()), "b === a", ) self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.STRICT_NOT_EQUALS, Variable("b")).swap()), "b !== a", ) self.assertEqual( str( AndIf( TwoParameterIf(Variable("a"), TwoParameterIf.LT, Variable("b")), TwoParameterIf(Variable("c"), TwoParameterIf.GT, Variable("d")), ).swap() ), "c > d && a < b", ) self.assertEqual( str( OrIf( TwoParameterIf(Variable("a"), TwoParameterIf.LT, Variable("b")), TwoParameterIf(Variable("c"), TwoParameterIf.GT, Variable("d")), ).swap(), ), "c > d || a < b", ) def test_swap_double(self) -> None: self.assertEqual(str(IsUndefinedIf(Variable("a")).swap().swap()), "a is UNDEFINED") self.assertEqual(str(IsBooleanIf(Variable("a")).swap().swap()), "a") self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.EQUALS, Variable("b")).swap().swap()), "a == b", ) self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.NOT_EQUALS, Variable("b")).swap().swap()), "a != b", ) self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.LT, Variable("b")).swap().swap()), "a < b", ) self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.GT, Variable("b")).swap().swap()), "a > b", ) self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.LT_EQUALS, Variable("b")).swap().swap()), "a <= b", ) self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.GT_EQUALS, Variable("b")).swap().swap()), "a >= b", ) self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.STRICT_EQUALS, Variable("b")).swap().swap()), "a === b", ) self.assertEqual( str(TwoParameterIf(Variable("a"), TwoParameterIf.STRICT_NOT_EQUALS, Variable("b")).swap().swap()), "a !== b", ) self.assertEqual( str( AndIf( TwoParameterIf(Variable("a"), TwoParameterIf.LT, Variable("b")), TwoParameterIf(Variable("c"), TwoParameterIf.GT, Variable("d")), ) .swap() .swap() ), "a < b && c > d", ) self.assertEqual( str( OrIf( TwoParameterIf(Variable("a"), TwoParameterIf.LT, Variable("b")), TwoParameterIf(Variable("c"), TwoParameterIf.GT, Variable("d")), ) .swap() .swap() ), "a < b || c > d", ) def test_simplify_noop(self) -> None: self.assertEqual(str(IsUndefinedIf(Variable("a")).simplify()), "a is UNDEFINED") self.assertEqual(str(IsUndefinedIf(Variable("a")).invert().simplify()), "a is not UNDEFINED") self.assertEqual(str(IsBooleanIf(Variable("a")).simplify()), "a") self.assertEqual(str(IsBooleanIf(Variable("a")).invert().simplify()), "not a") self.assertEqual( str( AndIf( TwoParameterIf(Variable("a"), TwoParameterIf.LT, Variable("b")), TwoParameterIf(Variable("c"), TwoParameterIf.GT, Variable("d")), ).simplify() ), "a < b && c > d", ) self.assertEqual( str( OrIf( TwoParameterIf(Variable("a"), TwoParameterIf.LT, Variable("b")), TwoParameterIf(Variable("c"), TwoParameterIf.GT, Variable("d")), ).simplify() ), "a < b || c > d", ) self.assertEqual( str( AndIf( TwoParameterIf(Variable("a"), TwoParameterIf.LT, Variable("b")), TwoParameterIf(Variable("a"), TwoParameterIf.GT_EQUALS, Variable("c")), ).simplify() ), "a < b && a >= c", ) self.assertEqual( str( OrIf( TwoParameterIf(Variable("a"), TwoParameterIf.LT, Variable("b")), TwoParameterIf(Variable("a"), TwoParameterIf.GT, Variable("b")), ).simplify() ), "a < b || a > b", ) def test_simplify_basic(self) -> None: self.assertEqual(str(IsUndefinedIf(UNDEFINED).simplify()), "True") self.assertEqual(str(IsUndefinedIf(UNDEFINED).invert().simplify()), "False") self.assertEqual(str(IsBooleanIf(True).simplify()), "True") self.assertEqual(str(IsBooleanIf(True).invert().simplify()), "False") self.assertEqual(str(IsBooleanIf(False).simplify()), "False") self.assertEqual(str(IsBooleanIf(False).invert().simplify()), "True") def test_simplify_compound(self) -> None: self.assertEqual( str( AndIf( TwoParameterIf(Variable("a"), TwoParameterIf.LT, Variable("b")), IsBooleanIf(True), ).simplify() ), "a < b", ) self.assertEqual( str( AndIf( TwoParameterIf(Variable("a"), TwoParameterIf.LT, Variable("b")), IsBooleanIf(False), ).simplify() ), "False", ) self.assertEqual( str( AndIf( IsBooleanIf(True), TwoParameterIf(Variable("a"), TwoParameterIf.LT, Variable("b")), ).simplify() ), "a < b", ) self.assertEqual( str( AndIf( IsBooleanIf(False), TwoParameterIf(Variable("a"), TwoParameterIf.LT, Variable("b")), ).simplify() ), "False", ) self.assertEqual( str( OrIf( TwoParameterIf(Variable("a"), TwoParameterIf.LT, Variable("b")), IsBooleanIf(True), ).simplify() ), "True", ) self.assertEqual( str( OrIf( TwoParameterIf(Variable("a"), TwoParameterIf.LT, Variable("b")), IsBooleanIf(False), ).simplify() ), "a < b", ) self.assertEqual( str( OrIf( IsBooleanIf(True), TwoParameterIf(Variable("a"), TwoParameterIf.LT, Variable("b")), ).simplify() ), "True", ) self.assertEqual( str( OrIf( IsBooleanIf(False), TwoParameterIf(Variable("a"), TwoParameterIf.LT, Variable("b")), ).simplify() ), "a < b", ) def test_simplify_equivalent(self) -> None: self.assertEqual( str( AndIf( TwoParameterIf(Variable("a"), TwoParameterIf.LT, Variable("b")), TwoParameterIf(Variable("a"), TwoParameterIf.GT_EQUALS, Variable("b")), ).simplify() ), "False", ) self.assertEqual( str( OrIf( TwoParameterIf(Variable("a"), TwoParameterIf.LT, Variable("b")), TwoParameterIf(Variable("a"), TwoParameterIf.GT_EQUALS, Variable("b")), ).simplify() ), "True", ) def test_equals_associativity(self) -> None: self.assertEqual( AndIf( IsBooleanIf(Variable("a")), AndIf( IsBooleanIf(Variable("b")), IsBooleanIf(Variable("c")), ), ), AndIf( AndIf( IsBooleanIf(Variable("a")), IsBooleanIf(Variable("b")), ), IsBooleanIf(Variable("c")), ), ) self.assertEqual( OrIf( IsBooleanIf(Variable("a")), OrIf( IsBooleanIf(Variable("b")), IsBooleanIf(Variable("c")), ), ), OrIf( OrIf( IsBooleanIf(Variable("a")), IsBooleanIf(Variable("b")), ), IsBooleanIf(Variable("c")), ), ) def test_equals_commutativity(self) -> None: self.assertEqual( AndIf( IsBooleanIf(Variable("a")), IsBooleanIf(Variable("b")), ), AndIf( IsBooleanIf(Variable("b")), IsBooleanIf(Variable("a")), ), ) self.assertEqual( OrIf( IsBooleanIf(Variable("a")), IsBooleanIf(Variable("b")), ), OrIf( IsBooleanIf(Variable("b")), IsBooleanIf(Variable("a")), ), ) def test_simplify_identity(self) -> None: self.assertEqual( str( AndIf( IsBooleanIf(Variable("a")), IsBooleanIf(True), ).simplify(), ), "a", ) self.assertEqual( str( OrIf( IsBooleanIf(Variable("a")), IsBooleanIf(False), ).simplify(), ), "a", ) def test_simplify_annihilation(self) -> None: self.assertEqual( str( AndIf( IsBooleanIf(Variable("a")), IsBooleanIf(False), ).simplify(), ), "False", ) self.assertEqual( str( OrIf( IsBooleanIf(Variable("a")), IsBooleanIf(True), ).simplify(), ), "True", ) def test_simplify_idempotence(self) -> None: self.assertEqual( str( AndIf( IsBooleanIf(Variable("a")), IsBooleanIf(Variable("a")), ).simplify(), ), "a", ) self.assertEqual( str( OrIf( IsBooleanIf(Variable("a")), IsBooleanIf(Variable("a")), ).simplify(), ), "a", ) def test_simplify_complementation(self) -> None: self.assertEqual( str( AndIf( IsBooleanIf(Variable("a")), IsBooleanIf(Variable("a")).invert(), ).simplify(), ), "False", ) self.assertEqual( str( OrIf( IsBooleanIf(Variable("a")), IsBooleanIf(Variable("a")).invert(), ).simplify(), ), "True", ) def test_simplify_elimination(self) -> None: self.assertEqual( str( OrIf( AndIf( IsBooleanIf(Variable("x")), IsBooleanIf(Variable("y")), ), AndIf( IsBooleanIf(Variable("x")), IsBooleanIf(Variable("y")).invert(), ), ).simplify(), ), "x", ) self.assertEqual( str( AndIf( OrIf( IsBooleanIf(Variable("x")), IsBooleanIf(Variable("y")), ), OrIf( IsBooleanIf(Variable("x")), IsBooleanIf(Variable("y")).invert(), ), ).simplify(), ), "x", ) def test_simplify_absorption(self) -> None: self.assertEqual( str( AndIf( IsBooleanIf(Variable("x")), OrIf( IsBooleanIf(Variable("x")), IsBooleanIf(Variable("y")), ), ).simplify(), ), "x", ) self.assertEqual( str( OrIf( IsBooleanIf(Variable("x")), AndIf( IsBooleanIf(Variable("x")), IsBooleanIf(Variable("y")), ), ).simplify(), ), "x", ) def test_simplify_negative_absorption(self) -> None: self.assertEqual( str( AndIf( IsBooleanIf(Variable("x")), OrIf( IsBooleanIf(Variable("x")).invert(), IsBooleanIf(Variable("y")), ), ).simplify(), ), "x && y", ) self.assertEqual( str( OrIf( IsBooleanIf(Variable("x")), AndIf( IsBooleanIf(Variable("x")).invert(), IsBooleanIf(Variable("y")), ), ).simplify(), ), "x || y", ) class TestAFPOptimize(ExtendedTestCase): def __optimize_code(self, statements: Sequence[Statement]) -> List[str]: bcd = ByteCodeDecompiler( ByteCode( None, [], -1, ), optimize=True, ) with bcd.debugging(verbose=self.verbose): return bcd._pretty_print(bcd._optimize_code(statements), prefix="").split(os.linesep) def test_no_flow(self) -> None: statements: List[Statement] = [ PlayMovieStatement(), StopMovieStatement(), ] self.assertEqual( self.__optimize_code(statements), [ "builtin_StartPlaying();", "builtin_StopPlaying();", ], ) def test_basic_flow(self) -> None: statements: List[Statement] = [ PlayMovieStatement(), IfStatement( IsBooleanIf( Variable("a"), ), [ NextFrameStatement(), ], [ PreviousFrameStatement(), ], ), StopMovieStatement(), ] self.assertEqual( self.__optimize_code(statements), [ "builtin_StartPlaying();", "if (a)", "{", " builtin_GotoNextFrame();", "}", "else", "{", " builtin_GotoPreviousFrame();", "}", "builtin_StopPlaying();", ], ) def test_compound_or_basic(self) -> None: statements: List[Statement] = [ IfStatement( TwoParameterIf( Variable("a"), TwoParameterIf.NOT_EQUALS, 1, ), [ IfStatement( TwoParameterIf( Variable("a"), TwoParameterIf.NOT_EQUALS, 2, ), [ StopMovieStatement(), DefineLabelStatement(1000), ReturnStatement("strval"), ], [], ), ], [], ), PlayMovieStatement(), GotoStatement(1000), ] self.assertEqual( self.__optimize_code(statements), [ "if (a == 1 || a == 2)", "{", " builtin_StartPlaying();", "}", "else", "{", " builtin_StopPlaying();", "}", "return 'strval';", ], ) def test_compound_or_alternate(self) -> None: statements: List[Statement] = [ IfStatement( TwoParameterIf( Variable("a"), TwoParameterIf.NOT_EQUALS, 1, ), [ IfStatement( TwoParameterIf( Variable("a"), TwoParameterIf.NOT_EQUALS, 2, ), [ StopMovieStatement(), GotoStatement(1000), ], [], ), ], [], ), PlayMovieStatement(), DefineLabelStatement(1000), ReturnStatement("strval"), ] self.assertEqual( self.__optimize_code(statements), [ "if (a == 1 || a == 2)", "{", " builtin_StartPlaying();", "}", "else", "{", " builtin_StopPlaying();", "}", "return 'strval';", ], ) def test_compound_or_inside_if(self) -> None: statements: List[Statement] = [ IfStatement( IsBooleanIf( Variable("debug"), ).invert(), [ IfStatement( TwoParameterIf( Variable("a"), TwoParameterIf.NOT_EQUALS, 1, ), [ IfStatement( TwoParameterIf( Variable("a"), TwoParameterIf.NOT_EQUALS, 2, ), [ StopMovieStatement(), DefineLabelStatement(1000), ReturnStatement("strval"), ], [], ), ], [], ), PlayMovieStatement(), GotoStatement(1000), ], [], ), ] self.assertEqual( self.__optimize_code(statements), [ "if (not debug)", "{", " if (a == 1 || a == 2)", " {", " builtin_StartPlaying();", " }", " else", " {", " builtin_StopPlaying();", " }", " return 'strval';", "}", ], ) def test_compound_or_inside_while(self) -> None: statements: List[Statement] = [ ForStatement( "x", 0, TwoParameterIf( Variable("x"), TwoParameterIf.LT, 10, ), ArithmeticExpression( Variable("x"), "+", 1, ), [ IfStatement( TwoParameterIf( Variable("a"), TwoParameterIf.NOT_EQUALS, 1, ), [ IfStatement( TwoParameterIf( Variable("a"), TwoParameterIf.NOT_EQUALS, 2, ), [ StopMovieStatement(), DefineLabelStatement(1000), ReturnStatement("strval"), ], [], ), ], [], ), PlayMovieStatement(), GotoStatement(1000), ], local=True, ), ] self.assertEqual( self.__optimize_code(statements), [ "for (local x = 0; x < 10; x = x + 1)", "{", " if (a == 1 || a == 2)", " {", " builtin_StartPlaying();", " }", " else", " {", " builtin_StopPlaying();", " }", " return 'strval';", "}", ], ) def test_compound_or_with_inner_if(self) -> None: statements: List[Statement] = [ IfStatement( TwoParameterIf( Variable("a"), TwoParameterIf.NOT_EQUALS, 1, ), [ IfStatement( TwoParameterIf( Variable("a"), TwoParameterIf.NOT_EQUALS, 2, ), [ IfStatement( TwoParameterIf( Variable("x"), TwoParameterIf.EQUALS, 5, ), [ StopMovieStatement(), ], [], ), DefineLabelStatement(1000), ReturnStatement("strval"), ], [], ), ], [], ), IfStatement( TwoParameterIf( Variable("x"), TwoParameterIf.EQUALS, 10, ), [ PlayMovieStatement(), ], [], ), GotoStatement(1000), ] self.assertEqual( self.__optimize_code(statements), [ "if (a == 1 || a == 2)", "{", " if (x == 10)", " {", " builtin_StartPlaying();", " }", "}", "else", "{", " if (x == 5)", " {", " builtin_StopPlaying();", " }", "}", "return 'strval';", ], ) def test_compound_or_with_inner_compound_or(self) -> None: statements: List[Statement] = [ IfStatement( TwoParameterIf( Variable("a"), TwoParameterIf.NOT_EQUALS, 1, ), [ IfStatement( TwoParameterIf( Variable("a"), TwoParameterIf.NOT_EQUALS, 2, ), [ IfStatement( TwoParameterIf( Variable("x"), TwoParameterIf.NOT_EQUALS, 10, ), [ IfStatement( TwoParameterIf( Variable("x"), TwoParameterIf.NOT_EQUALS, 20, ), [ StopMovieStatement(), GotoStatement(1000), ], [], ), ], [], ), PlayMovieStatement(), GotoStatement(1000), ], [], ), ], [], ), NextFrameStatement(), PreviousFrameStatement(), DefineLabelStatement(1000), ReturnStatement("strval"), ] self.assertEqual( self.__optimize_code(statements), [ "if (a == 1 || a == 2)", "{", " builtin_GotoNextFrame();", " builtin_GotoPreviousFrame();", "}", "else", "{", " if (x == 10 || x == 20)", " {", " builtin_StartPlaying();", " }", " else", " {", " builtin_StopPlaying();", " }", "}", "return 'strval';", ], ) def test_compound_or_triple(self) -> None: statements: List[Statement] = [ IfStatement( TwoParameterIf( Variable("a"), TwoParameterIf.NOT_EQUALS, 1, ), [ IfStatement( TwoParameterIf( Variable("a"), TwoParameterIf.NOT_EQUALS, 2, ), [ IfStatement( TwoParameterIf( Variable("a"), TwoParameterIf.NOT_EQUALS, 3, ), [ StopMovieStatement(), DefineLabelStatement(1000), ReturnStatement("strval"), ], [], ), ], [], ), ], [], ), PlayMovieStatement(), GotoStatement(1000), ] self.assertEqual( self.__optimize_code(statements), [ "if (a == 1 || a == 2 || a == 3)", "{", " builtin_StartPlaying();", "}", "else", "{", " builtin_StopPlaying();", "}", "return 'strval';", ], ) def test_compound_or_quad(self) -> None: statements: List[Statement] = [ IfStatement( TwoParameterIf( Variable("a"), TwoParameterIf.NOT_EQUALS, 1, ), [ IfStatement( TwoParameterIf( Variable("a"), TwoParameterIf.NOT_EQUALS, 2, ), [ IfStatement( TwoParameterIf( Variable("a"), TwoParameterIf.NOT_EQUALS, 3, ), [ IfStatement( TwoParameterIf( Variable("a"), TwoParameterIf.NOT_EQUALS, 4, ), [ StopMovieStatement(), DefineLabelStatement(1000), ReturnStatement("strval"), ], [], ), ], [], ), ], [], ), ], [], ), PlayMovieStatement(), GotoStatement(1000), ] self.assertEqual( self.__optimize_code(statements), [ "if (a == 1 || a == 2 || a == 3 || a == 4)", "{", " builtin_StartPlaying();", "}", "else", "{", " builtin_StopPlaying();", "}", "return 'strval';", ], ) def test_compound_or_quint(self) -> None: # Okay, at this point I believe that the algorithm works... statements: List[Statement] = [ IfStatement( TwoParameterIf( Variable("a"), TwoParameterIf.NOT_EQUALS, 1, ), [ IfStatement( TwoParameterIf( Variable("a"), TwoParameterIf.NOT_EQUALS, 2, ), [ IfStatement( TwoParameterIf( Variable("a"), TwoParameterIf.NOT_EQUALS, 3, ), [ IfStatement( TwoParameterIf( Variable("a"), TwoParameterIf.NOT_EQUALS, 4, ), [ IfStatement( TwoParameterIf( Variable("a"), TwoParameterIf.NOT_EQUALS, 5, ), [ StopMovieStatement(), DefineLabelStatement(1000), ReturnStatement("strval"), ], [], ), ], [], ), ], [], ), ], [], ), ], [], ), PlayMovieStatement(), GotoStatement(1000), ] self.assertEqual( self.__optimize_code(statements), [ "if (a == 1 || a == 2 || a == 3 || a == 4 || a == 5)", "{", " builtin_StartPlaying();", "}", "else", "{", " builtin_StopPlaying();", "}", "return 'strval';", ], )