diff --git a/bemani/format/afp/decompile.py b/bemani/format/afp/decompile.py index 6d76a5a..97488f1 100644 --- a/bemani/format/afp/decompile.py +++ b/bemani/format/afp/decompile.py @@ -1519,17 +1519,28 @@ class ByteCodeDecompiler(VerboseOutput): [false_action] if false_action else [], ) + if last_action.opcode in [AP2Action.RETURN, AP2Action.THROW, AP2Action.END]: + if len(chunk.next_chunks) != 1: + raise Exception(f"Logic error, chunkd ID {chunk.id} returns, throws or end to multiple blocks!") + if chunk.next_chunks[0] != offset_map[self.bytecode.end_offset]: + raise Exception(f"Expected chunk ID {chunk.id} to jump to return block but jumped elsewhere!") + # We will convert this later. + self.vprint("Severing link to return address.") + chunk.next_chunks = [] + # At this point, all chunks in our list should point only to other chunks in our list. for chunk in loop.chunks: for n in chunk.next_chunks: if n not in internal_jump_points: - raise Exception(f"Found unconverted next chunk {n} in chunk ID {chunk.id}!") + raise Exception(f"Found unconverted next chunk {n} in chunk ID {chunk.id}, for loop ID {loop.id} with break point {break_point}!") if isinstance(chunk, ByteCodeChunk): last_action = chunk.actions[-1] if isinstance(last_action, AP2Action): if last_action.opcode == AP2Action.IF and len(chunk.next_chunks) != 2: raise Exception(f"Somehow messed up the next pointers on if statement in chunk ID {chunk.id}!") - if last_action.opcode in [AP2Action.JUMP, AP2Action.RETURN, AP2Action.THROW, AP2Action.END] and len(chunk.next_chunks) != 1: + if last_action.opcode == AP2Action.JUMP and len(chunk.next_chunks) != 1: + raise Exception(f"Somehow messed up the next pointers on control flow statement in chunk ID {chunk.id}!") + if last_action.opcode in [AP2Action.RETURN, AP2Action.THROW, AP2Action.END] and len(chunk.next_chunks) != 0: raise Exception(f"Somehow messed up the next pointers on control flow statement in chunk ID {chunk.id}!") else: if len(chunk.next_chunks) > 1: @@ -1579,6 +1590,27 @@ class ByteCodeDecompiler(VerboseOutput): self.vprint(f"Found loop with header {header} and blocks {', '.join(str(b) for b in blocks)}.") + # Now, make sure we scoop up any remaining if/else bodies not found in the backwards walk. + changed: bool = True + while changed: + changed = False + + for b in blocks: + # Explicitly exclude the header here, as it will only point at the break + # location which will usually pass the following dominator test. + if b == header: + continue + add_id: Optional[int] = None + for cid, doms in dominators.items(): + if dominators[b] == doms - {cid} and cid not in blocks and cid != header: + add_id = cid + break + if add_id is not None: + self.vprint(f"Chunk {cid} should be included in loop list!") + blocks.add(add_id) + changed = True + break + # We found a loop! if header in loops: raise Exception(f"Logic error, loop with header {header} was already found!") @@ -3100,6 +3132,33 @@ class ByteCodeDecompiler(VerboseOutput): return self.__walk(statements, swap_empty_ifs) + def __verify_balanced_labels(self, statements: Sequence[Statement]) -> None: + gotos: Set[int] = set() + labels: Set[int] = set() + + # Gather gotos and labels and make sure they're balanced. + def gather_gotos_and_labels(statement: Statement) -> Optional[Statement]: + nonlocal gotos + nonlocal labels + + if isinstance(statement, GotoStatement): + gotos.add(statement.location) + elif isinstance(statement, DefineLabelStatement): + labels.add(statement.location) + return statement + + self.__walk(statements, gather_gotos_and_labels) + + unmatched_gotos = gotos - labels + unmatched_labels = labels - gotos + + if unmatched_gotos: + formatted_labels = ", ".join(f"label_{x}" for x in unmatched_gotos) + raise Exception(f"Logic error, gotos found jumping to the following labels which don't exist: {formatted_labels}") + if unmatched_labels: + formatted_labels = ", ".join(f"label_{x}" for x in unmatched_labels) + raise Exception(f"Logic error, labels found with no gotos pointing at them: {formatted_labels}") + def __pretty_print(self, statements: Sequence[Statement], prefix: str = "") -> str: output: List[str] = [] @@ -3151,6 +3210,9 @@ class ByteCodeDecompiler(VerboseOutput): break statements = self.__swap_empty_ifs(statements) + # Let's sanity check the code for a few things that might trip us up. + self.__verify_balanced_labels(statements) + # Finally, let's save the code! self.__statements = statements