From 7ca9e6920fb6a8e32004773fc35bfde04b09223a Mon Sep 17 00:00:00 2001 From: Jennifer Taylor Date: Thu, 15 Jul 2021 00:14:29 +0000 Subject: [PATCH] Enable many more types of switch statements to be detected. --- bemani/format/afp/decompile.py | 114 +++++++++++++++++++++++++++-- bemani/tests/test_afp_decompile.py | 2 +- 2 files changed, 109 insertions(+), 7 deletions(-) diff --git a/bemani/format/afp/decompile.py b/bemani/format/afp/decompile.py index 9bc973c..cec7a77 100644 --- a/bemani/format/afp/decompile.py +++ b/bemani/format/afp/decompile.py @@ -3111,11 +3111,35 @@ class ByteCodeDecompiler(VerboseOutput): statement.false_statements and isinstance( statement.true_statements[-1], - (BreakStatement, ContinueStatement, ReturnStatement, NullReturnStatement, ThrowStatement, GotoStatement), + (ReturnStatement, NullReturnStatement, ThrowStatement, GotoStatement), ) and not isinstance( statement.false_statements[-1], - (BreakStatement, ContinueStatement, ReturnStatement, NullReturnStatement, ThrowStatement, GotoStatement), + (ReturnStatement, NullReturnStatement, ThrowStatement, GotoStatement), + ) + ): + # We need to walk both halves still, but once we're done, swap the true/false + # locations of the statements, so that the false value can be eliminated at + # some later optimization stage. + changed = True + false_statements, _ = update_ifs(statement.true_statements) + true_statements, _ = update_ifs(statement.false_statements) + + # Now, append the if statement, and follow up with the body. + statement.true_statements = true_statements + statement.false_statements = false_statements + statement.cond = statement.cond.invert() + new_statements.append(statement) + elif ( + statement.true_statements and + statement.false_statements and + isinstance( + statement.true_statements[-1], + (BreakStatement, ContinueStatement), + ) and + not isinstance( + statement.false_statements[-1], + (BreakStatement, ContinueStatement), ) ): # We need to walk both halves still, but once we're done, hoist the false @@ -3283,8 +3307,6 @@ class ByteCodeDecompiler(VerboseOutput): # First, is the current combination a valid combined or statement? candidate_true_expr = get_compound_if(candidate_statements) candidate_false_expr = candidate_true_expr.invert().simplify() - - true_cond = AndIf(conditional, candidate_true_expr).simplify() false_cond = AndIf(conditional, candidate_false_expr).simplify() if false_cond in paths: @@ -3482,7 +3504,7 @@ class ByteCodeDecompiler(VerboseOutput): def get_lhs(statement: IfStatement) -> Optional[Expression]: if not isinstance(statement.cond, TwoParameterIf): return None - if statement.cond.comp not in {TwoParameterIf.EQUALS, TwoParameterIf.NOT_EQUALS}: + if statement.cond.comp not in {TwoParameterIf.EQUALS, TwoParameterIf.NOT_EQUALS, TwoParameterIf.STRICT_EQUALS, TwoParameterIf.STRICT_NOT_EQUALS}: return None if not isinstance(statement.cond.conditional1, (Variable, Register, Member)): return None @@ -3573,7 +3595,7 @@ class ByteCodeDecompiler(VerboseOutput): cond = cast(TwoParameterIf, statement.cond) # If its already switched, leave it alone. - if cond.comp == TwoParameterIf.EQUALS: + if cond.comp in {TwoParameterIf.EQUALS, TwoParameterIf.STRICT_EQUALS}: new_batch.append(statement) return statement @@ -3681,6 +3703,85 @@ class ByteCodeDecompiler(VerboseOutput): statements = self.__walk(statements, replace_if_with_switch) return statements, changed + def __convert_if_gotos(self, statements: Sequence[Statement]) -> Tuple[List[Statement], bool]: + # Find if statements in the middle of a chunk of code whose last true statement is + # a return/goto/throw. Take all the following statements and put them in the if + # statement's false clause. We do this because it allows us to recognize other + # optimizations we could do on the code, including switch folding and redundant + # statement eliminiation. It also often helps get rid of gotos that were only used + # to skip the false path of an if statement. + new_statements: List[Statement] = [] + changed: bool = False + + i = 0 + while i < len(statements): + cur_statement = statements[i] + + if isinstance(cur_statement, IfStatement): + if ( + cur_statement.true_statements and + not cur_statement.false_statements and + isinstance(cur_statement.true_statements[-1], (NullReturnStatement, ReturnStatement, ThrowStatement, GotoStatement)) + ): + # This is a candidate! Take all following statements up until the possible label we jump to and + # make them into the false chunk. + last_true_statement = cur_statement.true_statements[-1] + stop_at_label: Optional[int] = None + if isinstance(last_true_statement, GotoStatement): + stop_at_label = last_true_statement.location + + # We skip past this statement because its the one we're updating. + i += 1 + + false_statements: List[Statement] = [] + while i < len(statements): + false_statement = statements[i] + if stop_at_label is not None and isinstance(false_statement, DefineLabelStatement): + if stop_at_label == false_statement.location: + # Exit early, the rest of the code including this + # label is not part of the else case. + break + + false_statements.append(false_statement) + i += 1 + + # Add all of these to the false case. + if false_statements: + cur_statement.false_statements = false_statements + changed = True + else: + # Skip past this statement, we have nothing to do aside from walk its children. + i += 1 + + # Regardless of whether we modified the if statement, recurse down its true and false path. + cur_statement.true_statements, new_changed = self.__convert_if_gotos(cur_statement.true_statements) + changed = changed or new_changed + + cur_statement.false_statements, new_changed = self.__convert_if_gotos(cur_statement.false_statements) + changed = changed or new_changed + + new_statements.append(cur_statement) + + elif isinstance(cur_statement, SwitchStatement): + for case in cur_statement.cases: + case.statements, new_changed = self.__convert_if_gotos(case.statements) + changed = changed or new_changed + new_statements.append(cur_statement) + i += 1 + + elif isinstance(cur_statement, DoWhileStatement): + cur_statement.body, new_changed = self.__convert_if_gotos(cur_statement.body) + changed = changed or new_changed + + new_statements.append(cur_statement) + i += 1 + + else: + new_statements.append(cur_statement) + i += 1 + + return new_statements, changed + def _optimize_code(self, statements: Sequence[Statement]) -> List[Statement]: statements = list(statements) @@ -3693,6 +3794,7 @@ class ByteCodeDecompiler(VerboseOutput): self.__remove_goto_return, self.__eliminate_useless_returns, self.__convert_loops, + self.__convert_if_gotos, self.__swap_empty_ifs, self.__drop_unneeded_else, self.__swap_ugly_ifexprs, diff --git a/bemani/tests/test_afp_decompile.py b/bemani/tests/test_afp_decompile.py index 6ab8c69..0fbd47b 100644 --- a/bemani/tests/test_afp_decompile.py +++ b/bemani/tests/test_afp_decompile.py @@ -660,7 +660,7 @@ class TestAFPDecompile(ExtendedTestCase): # "returning" early from a function. ]) statements = self.__call_decompile(bytecode) - self.assertEqual(self.__equiv(statements), [f"if (not True) {OPEN_BRACKET}{os.linesep} return{os.linesep}{CLOSE_BRACKET}", "builtin_StartPlaying()"]) + 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.