Clean up types on If statement, stop using strings to pass information.

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"))
+                return IsUndefinedIf(conditional, negate=negate != (action.comparison == IfAction.IS_UNDEFINED))
-            if action.comparison in ["IS TRUE", "IS FALSE"]:
+            if action.comparison in [IfAction.IS_TRUE, IfAction.IS_FALSE]:
                 conditional = stack.pop()
-                return IsBooleanIf(conditional, negate=negate != (action.comparison == "IS FALSE"))
+                return IsBooleanIf(conditional, negate=negate != (action.comparison == IfAction.IS_FALSE))
-            if action.comparison in ["==", "!="]:
+            if action.comparison in [IfAction.EQUALS, IfAction.NOT_EQUALS]:
                 conditional2 = stack.pop()
                 conditional1 = stack.pop()
-                return IsEqualIf(conditional1, conditional2, negate=negate != (action.comparison == "!="))
+                return IsEqualIf(conditional1, conditional2, negate=negate != (action.comparison == IfAction.NOT_EQUALS))
-            if action.comparison in ["STRICT ==", "STRICT !="]:
+            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 !="))
+                return IsStrictEqualIf(conditional1, conditional2, negate=negate != (action.comparison == IfAction.STRICT_NOT_EQUALS))
-            if action.comparison in ["<", ">"]:
+            if action.comparison in [IfAction.LT, IfAction.GT]:
                 conditional2 = stack.pop()
                 conditional1 = stack.pop()
-                return MagnitudeIf(conditional1, conditional2, negate=negate != (action.comparison == "<"))
+                return MagnitudeIf(conditional1, conditional2, negate=negate != (action.comparison == IfAction.LT))
-            if action.comparison in ["<=", ">="]:
+            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}")
 

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 = {
+                self.vprint(f"{prefix}      {lineno}: {action_name} {IfAction.comparison_to_str(if2_type)}, Offset If True: {jump_if_true_offset}")
-                    0: "==",
+                actions.append(IfAction(lineno, if2_type, jump_if_true_offset))
-                    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))
             elif opcode == AP2Action.JUMP:
                 jump_offset = struct.unpack(">h", datachunk[(offset_ptr + 1):(offset_ptr + 3)])[0]
                 jump_offset += (lineno + 3)

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):