diff --git a/bemani/format/afp/decompile.py b/bemani/format/afp/decompile.py index 5c19e09..fb7c19b 100644 --- a/bemani/format/afp/decompile.py +++ b/bemani/format/afp/decompile.py @@ -1679,28 +1679,28 @@ class ByteCodeDecompiler(VerboseOutput): stack: List[Any] = [] def make_if_expr(action: IfAction, negate: bool) -> IfExpr: - if action.comparison in ["IS UNDEFINED", "IS NOT UNDEFINED"]: + if action.comparison in [IfAction.IS_UNDEFINED, IfAction.IS_NOT_UNDEFINED]: conditional = stack.pop() - return IsUndefinedIf(conditional, negate=negate != (action.comparison == "IS UNDEFINED")) - if action.comparison in ["IS TRUE", "IS FALSE"]: + return IsUndefinedIf(conditional, negate=negate != (action.comparison == IfAction.IS_UNDEFINED)) + if action.comparison in [IfAction.IS_TRUE, IfAction.IS_FALSE]: conditional = stack.pop() - return IsBooleanIf(conditional, negate=negate != (action.comparison == "IS FALSE")) - if action.comparison in ["==", "!="]: + return IsBooleanIf(conditional, negate=negate != (action.comparison == IfAction.IS_FALSE)) + if action.comparison in [IfAction.EQUALS, IfAction.NOT_EQUALS]: conditional2 = stack.pop() conditional1 = stack.pop() - return IsEqualIf(conditional1, conditional2, negate=negate != (action.comparison == "!=")) - if action.comparison in ["STRICT ==", "STRICT !="]: + return IsEqualIf(conditional1, conditional2, negate=negate != (action.comparison == IfAction.NOT_EQUALS)) + if action.comparison in [IfAction.STRICT_EQUALS, IfAction.STRICT_NOT_EQUALS]: conditional2 = stack.pop() conditional1 = stack.pop() - return IsStrictEqualIf(conditional1, conditional2, negate=negate != (action.comparison == "STRICT !=")) - if action.comparison in ["<", ">"]: + return IsStrictEqualIf(conditional1, conditional2, negate=negate != (action.comparison == IfAction.STRICT_NOT_EQUALS)) + if action.comparison in [IfAction.LT, IfAction.GT]: conditional2 = stack.pop() conditional1 = stack.pop() - return MagnitudeIf(conditional1, conditional2, negate=negate != (action.comparison == "<")) - if action.comparison in ["<=", ">="]: + return MagnitudeIf(conditional1, conditional2, negate=negate != (action.comparison == IfAction.LT)) + if action.comparison in [IfAction.LT_EQUALS, IfAction.GT_EQUALS]: conditional2 = stack.pop() conditional1 = stack.pop() - return MagnitudeEqualIf(conditional1, conditional2, negate=negate != (action.comparison == "<=")) + return MagnitudeEqualIf(conditional1, conditional2, negate=negate != (action.comparison == IfAction.LT_EQUALS)) raise Exception(f"TODO: {action}") diff --git a/bemani/format/afp/swf.py b/bemani/format/afp/swf.py index c77a204..1bddabf 100644 --- a/bemani/format/afp/swf.py +++ b/bemani/format/afp/swf.py @@ -768,30 +768,14 @@ class SWF(TrackedCoverage, VerboseOutput): offset_ptr += 3 self.vprint(f"{prefix} {lineno}: Offset If True: {jump_if_true_offset}") - actions.append(IfAction(lineno, "IS TRUE", jump_if_true_offset)) + actions.append(IfAction(lineno, IfAction.IS_TRUE, jump_if_true_offset)) elif opcode == AP2Action.IF2: if2_type, jump_if_true_offset = struct.unpack(">Bh", datachunk[(offset_ptr + 1):(offset_ptr + 4)]) jump_if_true_offset += (lineno + 4) offset_ptr += 4 - if2_typestr = { - 0: "==", - 1: "!=", - 2: "<", - 3: ">", - 4: "<=", - 5: ">=", - 6: "IS FALSE", - 7: "BITAND", - 8: "BITNOTAND", - 9: "STRICT ==", - 10: "STRICT !=", - 11: "IS UNDEFINED", - 12: "IS NOT UNDEFINED", - }[if2_type] - - self.vprint(f"{prefix} {lineno}: {action_name} {if2_typestr}, Offset If True: {jump_if_true_offset}") - actions.append(IfAction(lineno, if2_typestr, jump_if_true_offset)) + self.vprint(f"{prefix} {lineno}: {action_name} {IfAction.comparison_to_str(if2_type)}, Offset If True: {jump_if_true_offset}") + actions.append(IfAction(lineno, if2_type, jump_if_true_offset)) elif opcode == AP2Action.JUMP: jump_offset = struct.unpack(">h", datachunk[(offset_ptr + 1):(offset_ptr + 3)])[0] jump_offset += (lineno + 3) diff --git a/bemani/format/afp/types/ap2.py b/bemani/format/afp/types/ap2.py index 888b491..f8d4109 100644 --- a/bemani/format/afp/types/ap2.py +++ b/bemani/format/afp/types/ap2.py @@ -2674,19 +2674,54 @@ class StoreRegisterAction(AP2Action): class IfAction(AP2Action): - def __init__(self, offset: int, comparison: str, jump_if_true_offset: int) -> None: + EQUALS = 0 + NOT_EQUALS = 1 + LT = 2 + GT = 3 + LT_EQUALS = 4 + GT_EQUALS = 5 + IS_FALSE = 6 + BITAND = 7 + NOT_BITAND = 8 + STRICT_EQUALS = 9 + STRICT_NOT_EQUALS = 10 + IS_UNDEFINED = 11 + IS_NOT_UNDEFINED = 12 + IS_TRUE = 1000 + + def __init__(self, offset: int, comparison: int, jump_if_true_offset: int) -> None: super().__init__(offset, AP2Action.IF) self.comparison = comparison self.jump_if_true_offset = jump_if_true_offset + @classmethod + def comparison_to_str(cls, comparison: int) -> str: + return { + cls.EQUALS: "==", + cls.NOT_EQUALS: "!=", + cls.LT: "<", + cls.GT: ">", + cls.LT_EQUALS: "<=", + cls.GT_EQUALS: ">=", + cls.IS_FALSE: "IS FALSE", + cls.BITAND: "BITAND", + cls.NOT_BITAND: "BITNOTAND", + cls.STRICT_EQUALS: "STRICT ==", + cls.STRICT_NOT_EQUALS: "STRICT !=", + cls.IS_UNDEFINED: "IS UNDEFINED", + cls.IS_NOT_UNDEFINED: "IS NOT UNDEFINED", + cls.IS_TRUE: "IS TRUE", + }[comparison] + def as_dict(self, *args: Any, **kwargs: Any) -> Dict[str, Any]: return { **super().as_dict(*args, **kwargs), + 'comparison': IfAction.comparison_to_str(self.comparison), 'jump_if_true_offset': self.jump_if_true_offset, } def __repr__(self) -> str: - return f"{self.offset}: {AP2Action.action_to_name(self.opcode)}, Comparison: {self.comparison}, Offset To Jump To If True: {self.jump_if_true_offset}" + return f"{self.offset}: {AP2Action.action_to_name(self.opcode)}, Comparison: {IfAction.comparison_to_str(self.comparison)}, Offset To Jump To If True: {self.jump_if_true_offset}" class JumpAction(AP2Action):