Compare non-matching functions against expected output

This makes it possible to catch regressions for non-matching functions,
especially those that only have minor issues.

This also reclassifies some minor non-matchings as major non-matchings
whenever it's really not obvious to see that they are equivalent.

3 files changed, 101 insertions, 13 deletions

diff --git a/tools/check.py b/tools/check.py
index 312c97e2..cf432319 100755
--- a/tools/check.py
+++ b/tools/check.py
@@ -11,8 +11,9 @@ import utils
 config: Dict[str, Any] = {}
 diff_settings.apply(config, {})
 
-base_elf = ELFFile((Path(__file__).parent.parent / config["baseimg"]).open("rb"))
-my_elf = ELFFile((Path(__file__).parent.parent / config["myimg"]).open("rb"))
+root = Path(__file__).parent.parent
+base_elf = ELFFile((root / config["baseimg"]).open("rb"))
+my_elf = ELFFile((root / config["myimg"]).open("rb"))
 my_symtab = my_elf.get_section_by_name(".symtab")
 if not my_symtab:
     utils.fail(f'{config["myimg"]} has no symbol table')
@@ -46,15 +47,7 @@ def get_fn_from_my_elf(name: str, size: int) -> bytes:
     return my_elf.stream.read(size)
 
 
-def check_function(addr: int, size: int, name: str) -> bool:
-    try:
-        base_fn = get_fn_from_base_elf(addr, size)
-    except KeyError:
-        utils.print_error(f"couldn't find base function 0x{addr:016x} for {utils.format_symbol_name_for_msg(name)}")
-        return False
-
-    my_fn = get_fn_from_my_elf(name, size)
-
+def check_function_ex(addr: int, size: int, base_fn: bytes, my_fn: bytes) -> bool:
     md = cs.Cs(cs.CS_ARCH_ARM64, cs.CS_MODE_ARM)
     md.detail = True
     adrp_pair_registers: Set[int] = set()
@@ -123,8 +116,23 @@ def check_function(addr: int, size: int, name: str) -> bool:
     return True
 
 
+def check_function(addr: int, size: int, name: str, base_fn=None) -> bool:
+    if base_fn is None:
+        try:
+            base_fn = get_fn_from_base_elf(addr, size)
+        except KeyError:
+            utils.print_error(f"couldn't find base function 0x{addr:016x} for {utils.format_symbol_name_for_msg(name)}")
+            return False
+
+    my_fn = get_fn_from_my_elf(name, size)
+    return check_function_ex(addr, size, base_fn, my_fn)
+
+
 def main() -> None:
     failed = False
+
+    nonmatching_fns_with_dump = {p.stem: p.read_bytes() for p in (root / "expected").glob("*.bin")}
+
     for func in utils.get_functions():
         if not func.decomp_name:
             continue
@@ -145,6 +153,12 @@ def main() -> None:
                 utils.print_note(
                     f"function {utils.format_symbol_name_for_msg(func.decomp_name)} is marked as non-matching but matches")
 
+            fn_dump = nonmatching_fns_with_dump.get(func.decomp_name, None)
+            if fn_dump is not None and not check_function(func.addr, len(fn_dump), func.decomp_name, fn_dump):
+                utils.print_error(
+                    f"function {utils.format_symbol_name_for_msg(func.decomp_name)} does not match expected output")
+                failed = True
+
     if failed:
         sys.exit(1)
 
diff --git a/tools/dump_function.py b/tools/dump_function.py
new file mode 100755
index 00000000..c26901a1
--- /dev/null
+++ b/tools/dump_function.py
@@ -0,0 +1,62 @@
+#!/usr/bin/env python3
+
+import argparse
+from elftools.elf.elffile import ELFFile
+import diff_settings
+from pathlib import Path
+from typing import Any, Dict
+import utils
+
+config: Dict[str, Any] = {}
+diff_settings.apply(config, {})
+
+root = Path(__file__).parent.parent
+my_elf = ELFFile((root / config["myimg"]).open("rb"))
+my_symtab = my_elf.get_section_by_name(".symtab")
+if not my_symtab:
+    utils.fail(f'{config["myimg"]} has no symbol table')
+
+
+def get_file_offset(elf, addr: int) -> int:
+    for seg in elf.iter_segments():
+        if seg.header["p_type"] != "PT_LOAD":
+            continue
+        if seg["p_vaddr"] <= addr < seg["p_vaddr"] + seg["p_filesz"]:
+            return addr - seg["p_vaddr"] + seg["p_offset"]
+    assert False
+
+
+def get_symbol_file_offset_and_size(elf, table, name: str) -> (int, int):
+    syms = table.get_symbol_by_name(name)
+    if not syms or len(syms) != 1:
+        raise KeyError(name)
+    return get_file_offset(elf, syms[0]["st_value"]), syms[0]["st_size"]
+
+
+def get_fn_from_my_elf(name: str) -> bytes:
+    offset, size = get_symbol_file_offset_and_size(my_elf, my_symtab, name)
+    my_elf.stream.seek(offset)
+    return my_elf.stream.read(size)
+
+
+def dump_fn(name: str) -> None:
+    expected_dir = root / "expected"
+    try:
+        fn = get_fn_from_my_elf(name)
+        path = expected_dir / f"{name}.bin"
+        path.parent.mkdir(exist_ok=True)
+        path.write_bytes(fn)
+    except KeyError:
+        utils.fail("could not find function")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("function_name", help="Name of the function to dump")
+    args = parser.parse_args()
+
+    dump_fn(args.function_name)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tools/progress.py b/tools/progress.py
index f18819ca..976f94ab 100755
--- a/tools/progress.py
+++ b/tools/progress.py
@@ -3,6 +3,7 @@ import argparse
 from collections import defaultdict
 from colorama import Back, Fore, Style
 import enum
+from pathlib import Path
 import utils
 from utils import FunctionStatus
 import typing as tp
@@ -14,6 +15,8 @@ parser.add_argument("--print-eq", "-e", action="store_true",
                     help="Print non-matching functions with minor issues")
 parser.add_argument("--print-ok", "-m", action="store_true",
                     help="Print matching functions")
+parser.add_argument("--hide-nonmatchings-with-dumps", "-H", help="Hide non-matching functions that have expected "
+                                                                 "output dumps", action="store_true")
 args = parser.parse_args()
 
 
@@ -31,6 +34,15 @@ counts: tp.DefaultDict[FunctionStatus, int] = defaultdict(int)
 ai_counts: tp.DefaultDict[AIClassType, int] = defaultdict(int)
 ai_counts_done: tp.DefaultDict[AIClassType, int] = defaultdict(int)
 
+nonmatching_fns_with_dump = {p.stem for p in (Path(__file__).parent.parent / "expected").glob("*.bin")}
+
+
+def should_hide_nonmatching(name: str) -> bool:
+    if not args.hide_nonmatchings_with_dumps:
+        return False
+    return name in nonmatching_fns_with_dump
+
+
 for info in utils.get_functions():
     code_size_total += info.size
     num_total += 1
@@ -57,10 +69,10 @@ for info in utils.get_functions():
     code_size[info.status] += info.size
 
     if info.status == FunctionStatus.NonMatching:
-        if args.print_nm:
+        if args.print_nm and not should_hide_nonmatching(info.decomp_name):
             print(f"{Fore.RED}NM{Fore.RESET} {utils.format_symbol_name(info.decomp_name)}")
     elif info.status == FunctionStatus.Equivalent:
-        if args.print_eq:
+        if args.print_eq and not should_hide_nonmatching(info.decomp_name):
             print(f"{Fore.YELLOW}EQ{Fore.RESET} {utils.format_symbol_name(info.decomp_name)}")
     elif info.status == FunctionStatus.Matching:
         if args.print_ok: