diff --git a/tools/count_misordered_functions.py b/tools/count_misordered_functions.py new file mode 100755 index 00000000..8416169b --- /dev/null +++ b/tools/count_misordered_functions.py @@ -0,0 +1,280 @@ +#!/usr/bin/env python3 + +import argparse +import bisect +import re +import subprocess +import sys +from collections import defaultdict +from dataclasses import dataclass, field +from pathlib import Path +from typing import Callable, NamedTuple, TypeVar + +from colorama import Fore, Style + +from common.util import config, tools, utils + +ORIGINAL_BASE = 0x7100000000 + +T = TypeVar("T") + + +class Symbol(NamedTuple): + path: Path + address: int + is_weak: bool + + +class Function(NamedTuple): + orig_addr: int + decomp_addr: int + name: str + obj: str + + +@dataclass +class MatchedFunctions: + functions: list[Function] = field(default_factory=list) + not_in_binary: int = 0 + shared_inline: int = 0 + ambiguous: int = 0 + + +@dataclass +class Report: + misordered_objects: list[str] + stray_functions: list[str] + misordered_functions: list[str] + + SECTIONS = { + "misordered objects": "misordered_objects", + "functions outside their object's main block in the original": "stray_functions", + "misordered functions within their object": "misordered_functions", + } + SUMMARY_NAMES = ("objects", "stray functions", "functions in objects") + + def sections(self) -> dict[str, list[str]]: + return {title: getattr(self, attribute) for title, attribute in self.SECTIONS.items()} + + def save(self, path: Path) -> None: + path.write_text("".join(f"# {title}\n" + "".join(f"{line}\n" for line in lines) + for title, lines in self.sections().items())) + + @classmethod + def load(cls, path: Path) -> "Report": + sections: dict[str, list[str]] = {} + lines: list[str] = [] + for line in path.read_text().splitlines(): + if line.startswith("# "): + lines = sections.setdefault(line[2:], []) + else: + lines.append(line) + return cls(**{attribute: sections.get(title, []) for title, attribute in cls.SECTIONS.items()}) + + +def get_function_symbols(paths: list[Path]) -> dict[str, list[Symbol]]: + objdump = tools.find_tool("llvm-objdump") + symbols: dict[str, list[Symbol]] = defaultdict(list) + for i in range(0, len(paths), 500): + output = subprocess.run([objdump, "--syms", *paths[i:i + 500]], check=True, capture_output=True, + text=True).stdout + path = paths[i] + for line in output.splitlines(): + if line.endswith("file format elf64-littleaarch64"): + path = Path(line.partition(":\t")[0]) + continue + # address, 7 flag columns (binding, weak, ..., type), sectionsize name + if len(line) < 26 or line[23] != "F" or line[25:].startswith("*UND*"): + continue + name = line.split()[-1] + symbols[name].append(Symbol(path, int(line[:16], 16), line[18] == "w")) + return symbols + + +def get_link_objects(build_dir: Path) -> list[Path]: + link_step = f"build {config.get_build_target()}:" + for line in (build_dir / "build.ninja").read_text().splitlines(): + if line.startswith(link_step): + inputs = line.split()[3:] + return [build_dir / path for path in inputs[:inputs.index("||")] if path.endswith(".obj")] + utils.fail(f"link step not found in {build_dir / 'build.ninja'}") + + +def get_object_name(path: Path, build_dir: Path) -> str: + return re.sub(r"CMakeFiles/[^/]+\.dir/", "", str(path.relative_to(build_dir))).removesuffix(".obj") + + +def match_functions(decomp_elf: Path, build_dir: Path) -> MatchedFunctions: + binary_symbols = get_function_symbols([decomp_elf]) + object_symbols = get_function_symbols(get_link_objects(build_dir)) + + matched = MatchedFunctions() + for info in utils.get_functions(): + if not info.decomp_name: + continue + in_binary = binary_symbols.get(info.decomp_name, []) + definitions = object_symbols.get(info.decomp_name, []) + if not in_binary: + matched.not_in_binary += 1 + elif len(in_binary) == 1 and len(definitions) == 1: + obj = get_object_name(definitions[0].path, build_dir) + matched.functions.append(Function(info.addr, in_binary[0].address, info.decomp_name, obj)) + elif all(definition.is_weak for definition in definitions): + matched.shared_inline += 1 + else: + matched.ambiguous += 1 + + matched.functions.sort() + return matched + + +def split_into_runs(functions: list[Function], + is_separator: Callable[[Function], bool]) -> dict[str, list[list[Function]]]: + runs: dict[str, list[list[Function]]] = defaultdict(list) + separators_seen = 0 + separators_seen_after_object: dict[str, int] = {} + for function in functions: + if separators_seen_after_object.get(function.obj) != separators_seen: + runs[function.obj].append([]) + runs[function.obj][-1].append(function) + if is_separator(function): + separators_seen += 1 + separators_seen_after_object[function.obj] = separators_seen + return runs + + +def split_into_blocks(functions: list[Function]) -> dict[str, list[list[Function]]]: + # A stray function must not split the object it lands in, so only functions in their + # object's longest run separate the other objects' functions. + runs = split_into_runs(functions, lambda function: True) + main_run_functions = {function for object_runs in runs.values() for function in max(object_runs, key=len)} + return split_into_runs(functions, lambda function: function in main_run_functions) + + +def longest_increasing_subsequence(values: list[int]) -> set[int]: + tail_values: list[int] = [] + tail_indices: list[int] = [] + predecessors: list[int] = [] + for i, value in enumerate(values): + position = bisect.bisect_right(tail_values, value) + predecessors.append(tail_indices[position - 1] if position else -1) + if position == len(tail_values): + tail_values.append(value) + tail_indices.append(i) + else: + tail_values[position] = value + tail_indices[position] = i + + indices: set[int] = set() + i = tail_indices[-1] if tail_indices else -1 + while i != -1: + indices.add(i) + i = predecessors[i] + return indices + + +def out_of_order(items: list[T], key: Callable[[T], int]) -> list[T]: + in_order = longest_increasing_subsequence([key(item) for item in items]) + return [item for i, item in enumerate(items) if i not in in_order] + + +def describe(function: Function) -> str: + return f"0x{function.orig_addr + ORIGINAL_BASE:016x} {function.obj} {function.name}" + + +def colorize(line: str) -> str: + address, obj, *name = line.split(" ", 2) + text = f"{address} {Fore.CYAN}{obj}{Fore.RESET}" + if name: + text += f" {utils.format_symbol_name(name[0])}" + return text + + +def print_comparison(baseline: Report, report: Report) -> bool: + deltas: list[int] = [] + for (title, old_lines), new_lines in zip(baseline.sections().items(), report.sections().values()): + old_set, new_set = set(old_lines), set(new_lines) + fixed = [line for line in old_lines if line not in new_set] + added = [line for line in new_lines if line not in old_set] + delta = len(new_lines) - len(old_lines) + deltas.append(delta) + color = Fore.RED if delta > 0 else Fore.GREEN if delta < 0 else "" + print(f"{Style.BRIGHT}{title}: {len(old_lines)} -> {len(new_lines)} " + f"{color}({delta:+d}){Style.RESET_ALL}") + for line in fixed: + print(f"{Fore.GREEN}-{Fore.RESET} {colorize(line)}") + for line in added: + print(f"{Fore.RED}+{Fore.RESET} {colorize(line)}") + print() + + changes = ", ".join(f"{name} {delta:+d}" for name, delta in zip(Report.SUMMARY_NAMES, deltas)) + print(f"order change: {changes}") + return any(delta > 0 for delta in deltas) + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Measure how closely object and function order in the decomp binary follows the original executable") + parser.add_argument("-v", "--verbose", action="store_true", help="list the misplaced objects and functions") + parser.add_argument("--save", type=Path, metavar="FILE", help="save the misplaced objects and functions to FILE") + parser.add_argument("--baseline", type=Path, metavar="FILE", + help="compare against a file written by --save and fail if any misplacement count grew") + args = parser.parse_args() + + decomp_elf = config.get_decomp_elf().resolve() + matched = match_functions(decomp_elf, decomp_elf.parent) + + main_runs: dict[str, list[Function]] = {} + split_objects: list[str] = [] + stray_functions: list[Function] = [] + for obj, runs in split_into_blocks(matched.functions).items(): + main_run = max(runs, key=len) + main_runs[obj] = main_run + if len(runs) > 1: + split_objects.append(obj) + stray_functions += [function for run in runs if run is not main_run for function in run] + stray_functions.sort() + + objects_in_original_order = sorted(main_runs, key=lambda obj: main_runs[obj][0].orig_addr) + misordered_objects = out_of_order(objects_in_original_order, + lambda obj: min(function.decomp_addr for function in main_runs[obj])) + + misordered_functions: list[Function] = [] + for obj in objects_in_original_order: + misordered_functions += out_of_order(main_runs[obj], lambda function: function.decomp_addr) + + report = Report( + misordered_objects=[f"0x{main_runs[obj][0].orig_addr + ORIGINAL_BASE:016x} {obj}" for obj in misordered_objects], + stray_functions=[describe(function) for function in stray_functions], + misordered_functions=[describe(function) for function in misordered_functions], + ) + + if args.save: + report.save(args.save) + + if args.verbose: + for title, lines in report.sections().items(): + print(f"{Style.BRIGHT}{title}{Style.RESET_ALL}") + for line in lines: + print(colorize(line)) + print() + + print(f"{len(matched.functions):>7d} compared functions") + print(f"{Style.DIM}{matched.not_in_binary:>7d} skipped, not in the decomp binary") + print(f"{matched.shared_inline:>7d} skipped, inline in several objects") + print(f"{matched.ambiguous:>7d} skipped, ambiguous name{Style.RESET_ALL}") + print(f"{len(main_runs):>7d} objects") + print(f"{len(misordered_objects):>7d} {Fore.RED}misordered objects{Fore.RESET}") + print(f"{len(split_objects):>7d} {Fore.YELLOW}objects split in the original{Fore.RESET} " + f"({len(stray_functions)} stray functions)") + print(f"{len(misordered_functions):>7d} {Fore.YELLOW}misordered functions within objects{Fore.RESET}") + + if args.baseline: + print() + if print_comparison(Report.load(args.baseline), report): + sys.stdout.flush() + utils.fail("this change makes the object or function order worse") + + +if __name__ == "__main__": + main()