Simplify logic for related vtable searching

src/smgdc/angr/vtable_disamb.py

@@ -33,8 +33,12 @@ def read_cstring(
 
 class VtableDisambiguator(angr.Analysis):
     syms_by_addr: dict[int, set[Symbol]]
+
+    # given a vtable symbol, returns childrens' vtable symbols in an undefined order
     subclass_map: dict[Symbol, set[Symbol]]
-    superclass_map: dict[Symbol, set[Symbol]]
+
+    # given a vtable symbol, returns the parents' vtable symbols in ascending order
+    superclass_map: dict[Symbol, list[Symbol]]
 
     # we use this to test if a function belongs to a class with a vtable
     # this way we can prune static functions from the candidate pool
@@ -59,7 +63,7 @@ class VtableDisambiguator(angr.Analysis):
 
         # we associate a type with their subclasses so we can check them to disambiguate duplicates on the base
         self.subclass_map = collections.defaultdict(set)
-        self.superclass_map = collections.defaultdict(set)
+        self.superclass_map = collections.defaultdict(list)
 
         for vtsym in vtable_syms:
             if vtsym.is_import:
@@ -71,7 +75,7 @@ class VtableDisambiguator(angr.Analysis):
                 if not svtsym:
                     continue
                 self.subclass_map[svtsym].add(vtsym)
-                self.superclass_map[vtsym].add(svtsym)
+                self.superclass_map[vtsym].append(svtsym)
 
     def dump_class_parents(self, typeinfo_ptr: int) -> Iterable[tuple[int, str]]:
         """
@@ -164,21 +168,19 @@ class VtableDisambiguator(angr.Analysis):
     def get_possible_vtable_set_candidates(
         self, vtsym: Symbol, vtidx: int
     ) -> Iterable[set[Symbol]]:
-        # Yield the ordered list of subclasses for the given class.
-        yield self.subclass_map[vtsym]
-
         # It's possible that this class inherited a base version of a method that is specialized
-        # in a different part of the class hierarchy, so yield lists of subclasses of parents
-        # that include the vtable index too.
-        #
-        # FIXME: This can iterate over unrelated classes; properly filter this down to the
-        # subclasses of the base that implemented this.
-        max_vtsize = 4 * (vtidx + 2)
-        for parent_vtsym in sorted(
-            filter(lambda vt: vt.size > max_vtsize, self.superclass_map[vtsym]),
-            key=lambda vt: vt.size,
-        ):
-            yield self.subclass_map[parent_vtsym]
+        # in a different part of the class hierarchy, so find the base implementor then return
+        # it along with its descendants.
+        psym = self.get_parent_vt_implementing(vtsym, vtidx)
+        yield {psym} | self.subclass_map[psym]
+
+    def get_parent_vt_implementing(self, vtsym: Symbol, vtidx: int) -> Symbol:
+        # walk up parents until we reach the first one to implement this method
+        for super_vtsym, gp_vtsym in itertools.pairwise(self.superclass_map[vtsym]):
+            vt_first, *_ = self.get_vfptrs_from_table(gp_vtsym)
+            if len(vt_first) < vtidx:
+                return super_vtsym
+        return self.superclass_map[vtsym][-1]
 
     def get_vfptrs_from_table(self, vtsym: Symbol) -> list[list[int]]:
         # returns a list of addresses for each vtable present on the class