# 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';",
            ]
        )
