From d034f02df18642b0870c3df1d9292a99962b4b1b Mon Sep 17 00:00:00 2001
From: disinvite <disinvite@users.noreply.github.com>
Date: Fri, 19 Jan 2024 23:44:49 -0500
Subject: [PATCH] vtable enforce

---
 .github/workflows/build.yml |  6 +++
 tools/vtable/vtable.py      | 99 +++++++++++++++++++++++++++++++++++++
 2 files changed, 105 insertions(+)
 create mode 100644 tools/vtable/vtable.py

diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index 419033cc..51b0759f 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -145,6 +145,12 @@ jobs:
       run: |
         python3 tools/verexp/verexp.py legobin/LEGO1.DLL build/LEGO1.DLL
 
+    - name: Check Vtables
+      shell: bash
+      run: |
+        python3 tools/vtable/vtable.py legobin/ISLE.EXE build/ISLE.EXE build/ISLE.PDB .
+        python3 tools/vtable/vtable.py legobin/LEGO1.DLL build/LEGO1.DLL build/LEGO1.PDB .
+
     - name: Upload Artifact
       uses: actions/upload-artifact@master
       with:
diff --git a/tools/vtable/vtable.py b/tools/vtable/vtable.py
new file mode 100644
index 00000000..bdee080b
--- /dev/null
+++ b/tools/vtable/vtable.py
@@ -0,0 +1,99 @@
+import os
+import argparse
+import logging
+from typing import List
+import colorama
+from isledecomp.bin import Bin as IsleBin
+from isledecomp.compare import Compare as IsleCompare
+from isledecomp.utils import print_diff
+
+# Ignore all compare-db messages.
+logging.getLogger("isledecomp.compare").addHandler(logging.NullHandler())
+
+colorama.init()
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Comparing vtables.")
+    parser.add_argument(
+        "original", metavar="original-binary", help="The original binary"
+    )
+    parser.add_argument(
+        "recompiled", metavar="recompiled-binary", help="The recompiled binary"
+    )
+    parser.add_argument(
+        "pdb", metavar="recompiled-pdb", help="The PDB of the recompiled binary"
+    )
+    parser.add_argument(
+        "decomp_dir", metavar="decomp-dir", help="The decompiled source tree"
+    )
+    parser.add_argument(
+        "--verbose", "-v", action="store_true", help="Show more detailed information"
+    )
+    parser.add_argument(
+        "--no-color", "-n", action="store_true", help="Do not color the output"
+    )
+
+    (args, _) = parser.parse_known_args()
+
+    if not os.path.isfile(args.original):
+        parser.error(f"Original binary {args.original} does not exist")
+
+    if not os.path.isfile(args.recompiled):
+        parser.error(f"Recompiled binary {args.recompiled} does not exist")
+
+    if not os.path.isfile(args.pdb):
+        parser.error(f"Symbols PDB {args.pdb} does not exist")
+
+    if not os.path.isdir(args.decomp_dir):
+        parser.error(f"Source directory {args.decomp_dir} does not exist")
+
+    return args
+
+
+def show_vtable_diff(udiff: List[str], verbose: bool = False, plain: bool = False):
+    lines = [
+        line
+        for line in udiff
+        if verbose or line.startswith("+") or line.startswith("-")
+    ]
+    print_diff(lines, plain)
+
+
+def print_summary(vtable_count: int, problem_count: int):
+    if problem_count == 0:
+        print(f"Vtables found: {vtable_count}.\n100% match.")
+        return
+
+    print(f"Vtables found: {vtable_count}.\nVtables not matching: {problem_count}.")
+
+
+def main():
+    args = parse_args()
+    vtable_count = 0
+    problem_count = 0
+
+    with IsleBin(args.original) as orig_bin, IsleBin(args.recompiled) as recomp_bin:
+        engine = IsleCompare(orig_bin, recomp_bin, args.pdb, args.decomp_dir)
+
+        for tbl_match in engine.compare_vtables():
+            vtable_count += 1
+            if tbl_match.ratio < 1:
+                problem_count += 1
+
+                udiff = list(tbl_match.udiff)
+
+                print(
+                    tbl_match.name,
+                    f": orig 0x{tbl_match.orig_addr:x}, recomp 0x{tbl_match.recomp_addr:x}",
+                )
+                show_vtable_diff(udiff, args.verbose, args.no_color)
+                print()
+
+        print_summary(vtable_count, problem_count)
+
+    return 1 if problem_count > 0 else 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())