From 064feab51a4338b2a7bb76dca3d15b2027e36f78 Mon Sep 17 00:00:00 2001
From: MS <disinvite@users.noreply.github.com>
Date: Tue, 26 Mar 2024 21:01:37 -0400
Subject: [PATCH] Detect when we exceed original function size (#736)

---
 tools/isledecomp/isledecomp/compare/core.py | 11 ++++++++++-
 tools/isledecomp/isledecomp/compare/db.py   | 16 ++++++++++++++++
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/tools/isledecomp/isledecomp/compare/core.py b/tools/isledecomp/isledecomp/compare/core.py
index ce992f41..06887e4d 100644
--- a/tools/isledecomp/isledecomp/compare/core.py
+++ b/tools/isledecomp/isledecomp/compare/core.py
@@ -384,7 +384,16 @@ def _find_vtordisp(self):
                     self._db.set_function_pair(orig_addr, recomp_addr)
 
     def _compare_function(self, match: MatchInfo) -> DiffReport:
-        orig_raw = self.orig_bin.read(match.orig_addr, match.size)
+        # Detect when the recomp function size would cause us to read
+        # enough bytes from the original function that we cross into
+        # the next annotated function.
+        next_orig = self._db.get_next_orig_addr(match.orig_addr)
+        if next_orig is not None:
+            orig_size = min(next_orig - match.orig_addr, match.size)
+        else:
+            orig_size = match.size
+
+        orig_raw = self.orig_bin.read(match.orig_addr, orig_size)
         recomp_raw = self.recomp_bin.read(match.recomp_addr, match.size)
 
         # It's unlikely that a function other than an adjuster thunk would
diff --git a/tools/isledecomp/isledecomp/compare/db.py b/tools/isledecomp/isledecomp/compare/db.py
index 50cac87f..71f10373 100644
--- a/tools/isledecomp/isledecomp/compare/db.py
+++ b/tools/isledecomp/isledecomp/compare/db.py
@@ -73,6 +73,7 @@ def matchinfo_factory(_, row):
 
 
 class CompareDb:
+    # pylint: disable=too-many-public-methods
     def __init__(self):
         self._db = sqlite3.connect(":memory:")
         self._db.executescript(_SETUP_SQL)
@@ -348,6 +349,21 @@ def _match_on(self, compare_type: SymbolType, addr: int, name: str) -> bool:
 
         return self.set_pair(addr, recomp_addr, compare_type)
 
+    def get_next_orig_addr(self, addr: int) -> Optional[int]:
+        """Return the original address (matched or not) that follows
+        the one given. If our recomp function size would cause us to read
+        too many bytes for the original function, we can adjust it."""
+        result = self._db.execute(
+            """SELECT orig_addr
+            FROM `symbols`
+            WHERE orig_addr > ?
+            ORDER BY orig_addr
+            LIMIT 1""",
+            (addr,),
+        ).fetchone()
+
+        return result[0] if result is not None else None
+
     def match_function(self, addr: int, name: str) -> bool:
         did_match = self._match_on(SymbolType.FUNCTION, addr, name)
         if not did_match: