A bunch of shitty hacks to get slightly more accurate alignments

vtable_alignment.py

@@ -1,5 +1,6 @@
 import click
 import json
+import numpy as np
 
 from vtable_diff import extract_opcode_data
 from utils import eprint
@@ -64,6 +65,12 @@ def needleman_wunsch(old_seq, new_seq, similarity, gap_penalty):
     return alignment, mat[n][m]
 
 
+class Placeholder:
+    def __init__(self, old, new):
+        self.old = old
+        self.new = new
+
+
 class Similarity:
     def __init__(self, similarity_json_file):
         with open(similarity_json_file) as f:
@@ -75,10 +82,13 @@ class Similarity:
         self.new_opcodes = {
             opcode: idx for (idx, opcode) in enumerate(data["new_opcodes"])
         }
-        self.matrix = data["matrix"]
+        self.matrix = np.array(data["matrix"])
         self.warnings = set()
 
     def lookup(self, old_opcode, new_opcode):
+        if isinstance(old_opcode, Placeholder) or isinstance(new_opcode, Placeholder):
+            return -9999
+
         if old_opcode not in self.old_opcodes:
             self.warnings.add(
                 f"WARNING: Could not find old opcode {hex(old_opcode)} in similarity matrix"
@@ -93,6 +103,47 @@ class Similarity:
         j = self.new_opcodes[new_opcode]
         return self.matrix[i][j]
 
+    def get_confident_matches(self, threshold=0.1):
+        """
+        Returns matches in the form of [(old,new), ...] that are confidently
+        above the score threshold and where all pairs in the matching prefer
+        each other over any other opcodes.
+        """
+        scores = {}
+        new_opcodes = np.array([op for op in self.new_opcodes])
+        for old_op, i in self.old_opcodes.items():
+            top_n_idxs = np.argpartition(-self.matrix[i], 2)[:2]
+            scores[old_op] = [(new_opcodes[j], self.matrix[i][j]) for j in top_n_idxs]
+            scores[old_op].sort(key=lambda x: x[1], reverse=True)
+
+        transposed_matrix = self.matrix.transpose()
+        rev_scores = {}
+        old_opcodes = np.array([op for op in self.old_opcodes])
+        for new_op, j in self.new_opcodes.items():
+            top_n_idxs = np.argpartition(-transposed_matrix[j], 2)[:2]
+            rev_scores[new_op] = [
+                (old_opcodes[i], transposed_matrix[j][i]) for i in top_n_idxs
+            ]
+            rev_scores[new_op].sort(key=lambda x: x[1], reverse=True)
+
+        # Accept matches if they pass the threshold in the old => new direction,
+        # and if the new => old direction is a best match
+        matches = []
+        for old_op, top_matches in scores.items():
+            new_op = top_matches[0][0]
+            if top_matches[0][1] - top_matches[1][1] >= threshold:
+                rev_top_matches = rev_scores[new_op]
+                if (
+                    rev_top_matches[0][0] == old_op
+                    and rev_top_matches[0][1] - rev_top_matches[1][1] >= threshold
+                ):
+                    matches.append((old_op, new_op))
+
+        return matches
+
+    def clear_warnings(self):
+        self.warnings = set()
+
     def print_warnings(self):
         for warning in self.warnings:
             eprint(warning)
@@ -142,11 +193,61 @@ def vtable_alignment(old_exe, new_exe, similarity_json_file):
     new_seq = [opcode for opcode in new_opcodes_db.values()]
 
     similarity = Similarity(similarity_json_file)
+
+    matches = similarity.get_confident_matches()
+
+    # Ensure matches at least exist somewhere in the seq
+    old_seq_set = set(old_seq)
+    new_seq_set = set(new_seq)
+    matches = list(
+        filter(lambda x: x[0] in old_seq_set and x[1] in new_seq_set, matches)
+    )
+
+    eprint(f"Found {len(matches)} confident matches")
+
+    eprint("Running initial alignment...")
     alignment, score = needleman_wunsch(old_seq, new_seq, similarity, -1)
 
     similarity.print_warnings()
     eprint(f"Alignment score: {score}")
 
+    # Check for mismatches
+    matched_old = {match[0]: match[1] for match in matches}
+
+    mismatched_old = set()
+    mismatched_new = dict()
+    for old, new in alignment:
+        if old in matched_old and matched_old[old] != new:
+            truth = matched_old[old]
+            eprint(f"Mismatch detected! {hex(old)} => {hex(truth)}, got {hex(new)}")
+            mismatched_old.add(old)
+            mismatched_new[truth] = Placeholder(old, truth)
+
+    eprint("Attempting to fix mismatches")
+    old_seq = list(filter(lambda x: x not in mismatched_old, old_seq))
+    new_seq = list(
+        map(lambda x: mismatched_new[x] if x in mismatched_new else x, new_seq)
+    )
+
+    alignment, score = needleman_wunsch(old_seq, new_seq, similarity, -1)
+
+    old_seq = []
+    new_seq = []
+    for old, target in alignment:
+        if isinstance(target, Placeholder):
+            old_seq.append(target.old)
+            new_seq.append(target.new)
+        else:
+            if old is not None:
+                old_seq.append(old)
+            if target is not None:
+                new_seq.append(target)
+
+    similarity.clear_warnings()
+    alignment, score = needleman_wunsch(old_seq, new_seq, similarity, -1)
+    similarity.print_warnings()
+    eprint(f"Alignment score: {score}")
+
     diff = []
     for old, new in alignment:
         if old is None: