Refactor returned data structure for extensibility

tools/ghidra_scripts/import_functions_and_types_from_pdb.py

@@ -125,16 +125,15 @@ def add_python_path(path: str):
 # We need to quote the types here because they might not exist when running without Ghidra
 def import_function_into_ghidra(
     api: "FlatProgramAPI",
-    match_info: "MatchInfo",
-    signature: "FunctionSignature",
+    pdb_function: "PdbFunction",
     type_importer: "PdbTypeImporter",
 ):
-    hex_original_address = f"{match_info.orig_addr:x}"
+    hex_original_address = f"{pdb_function.match_info.orig_addr:x}"
 
     # Find the Ghidra function at that address
     ghidra_address = getAddressFactory().getAddress(hex_original_address)
     # pylint: disable=possibly-used-before-assignment
-    function_importer = PdbFunctionImporter(api, match_info, signature, type_importer)
+    function_importer = PdbFunctionImporter(api, pdb_function, type_importer)
 
     ghidra_function = getFunctionAt(ghidra_address)
     if ghidra_function is None:
@@ -165,7 +164,7 @@ def import_function_into_ghidra(
 
 
 def process_functions(extraction: "PdbFunctionExtractor"):
-    func_signatures = extraction.get_function_list()
+    pdb_functions = extraction.get_function_list()
 
     if not GLOBALS.running_from_ghidra:
         logger.info("Completed the dry run outside Ghidra.")
@@ -175,12 +174,13 @@ def process_functions(extraction: "PdbFunctionExtractor"):
     # pylint: disable=possibly-used-before-assignment
     type_importer = PdbTypeImporter(api, extraction)
 
-    for match_info, signature in func_signatures:
+    for pdb_func in pdb_functions:
+        func_name = pdb_func.match_info.name
         try:
-            import_function_into_ghidra(api, match_info, signature, type_importer)
+            import_function_into_ghidra(api, pdb_func, type_importer)
             GLOBALS.statistics.successes += 1
         except Lego1Exception as e:
-            log_and_track_failure(match_info.name, e)
+            log_and_track_failure(func_name, e)
         except RuntimeError as e:
             cause = e.args[0]
             if CancelledException is not None and isinstance(cause, CancelledException):
@@ -188,10 +188,10 @@ def process_functions(extraction: "PdbFunctionExtractor"):
                 logging.critical("Import aborted by the user.")
                 return
 
-            log_and_track_failure(match_info.name, cause, unexpected=True)
+            log_and_track_failure(func_name, cause, unexpected=True)
             logger.error(traceback.format_exc())
         except Exception as e:  # pylint: disable=broad-exception-caught
-            log_and_track_failure(match_info.name, e, unexpected=True)
+            log_and_track_failure(func_name, e, unexpected=True)
             logger.error(traceback.format_exc())
 
 
@@ -257,7 +257,6 @@ def main():
     from isledecomp.compare import Compare as IsleCompare
 
     reload_module("isledecomp.compare.db")
-    from isledecomp.compare.db import MatchInfo
 
     reload_module("lego_util.exceptions")
     from lego_util.exceptions import Lego1Exception
@@ -265,7 +264,7 @@ def main():
     reload_module("lego_util.pdb_extraction")
     from lego_util.pdb_extraction import (
         PdbFunctionExtractor,
-        FunctionSignature,
+        PdbFunction,
     )
 
     if GLOBALS.running_from_ghidra:

tools/ghidra_scripts/lego_util/function_importer.py

@@ -11,10 +11,8 @@
 from ghidra.program.model.listing import ParameterImpl
 from ghidra.program.model.symbol import SourceType
 
-from isledecomp.compare.db import MatchInfo
-
 from lego_util.pdb_extraction import (
-    FunctionSignature,
+    PdbFunction,
     CppRegisterSymbol,
     CppStackSymbol,
 )
@@ -37,28 +35,28 @@ class PdbFunctionImporter:
     def __init__(
         self,
         api: FlatProgramAPI,
-        match_info: MatchInfo,
-        signature: FunctionSignature,
+        func: PdbFunction,
         type_importer: "PdbTypeImporter",
     ):
         self.api = api
-        self.match_info = match_info
-        self.signature = signature
+        self.match_info = func.match_info
+        self.signature = func.signature
+        self.is_stub = func.is_stub
         self.type_importer = type_importer
 
-        if signature.class_type is not None:
+        if self.signature.class_type is not None:
             # Import the base class so the namespace exists
-            self.type_importer.import_pdb_type_into_ghidra(signature.class_type)
+            self.type_importer.import_pdb_type_into_ghidra(self.signature.class_type)
 
-        assert match_info.name is not None
+        assert self.match_info.name is not None
 
-        colon_split = sanitize_name(match_info.name).split("::")
+        colon_split = sanitize_name(self.match_info.name).split("::")
         self.name = colon_split.pop()
         namespace_hierachy = colon_split
         self.namespace = get_ghidra_namespace(api, namespace_hierachy)
 
         self.return_type = type_importer.import_pdb_type_into_ghidra(
-            signature.return_type
+            self.signature.return_type
         )
         self.arguments = [
             ParameterImpl(
@@ -66,7 +64,7 @@ def __init__(
                 type_importer.import_pdb_type_into_ghidra(type_name),
                 api.getCurrentProgram(),
             )
-            for (index, type_name) in enumerate(signature.arglist)
+            for (index, type_name) in enumerate(self.signature.arglist)
         ]
 
     @property

tools/ghidra_scripts/lego_util/pdb_extraction.py

@@ -38,6 +38,13 @@ class FunctionSignature:
     stack_symbols: list[CppStackOrRegisterSymbol]
 
 
+@dataclass
+class PdbFunction:
+    match_info: MatchInfo
+    signature: FunctionSignature
+    is_stub: bool
+
+
 class PdbFunctionExtractor:
     """
     Extracts all information on a given function from the parsed PDB
@@ -121,7 +128,7 @@ def get_func_signature(self, fn: SymbolsEntry) -> Optional[FunctionSignature]:
             stack_symbols=stack_symbols,
         )
 
-    def get_function_list(self) -> list[tuple[MatchInfo, FunctionSignature]]:
+    def get_function_list(self) -> list[PdbFunction]:
         handled = (
             self.handle_matched_function(match)
             for match in self.compare.get_functions()
@@ -130,11 +137,11 @@ def get_function_list(self) -> list[tuple[MatchInfo, FunctionSignature]]:
 
     def handle_matched_function(
         self, match_info: MatchInfo
-    ) -> Optional[tuple[MatchInfo, FunctionSignature]]:
+    ) -> Optional[PdbFunction]:
         assert match_info.orig_addr is not None
         match_options = self.compare.get_match_options(match_info.orig_addr)
         assert match_options is not None
-        if match_options.get("skip", False) or match_options.get("stub", False):
+        if match_options.get("skip", False):
             return None
 
         function_data = next(
@@ -163,4 +170,10 @@ def handle_matched_function(
         if function_signature is None:
             return None
 
-        return match_info, function_signature
+        is_stub = match_options.get("stub", False)
+
+        # TODO: Remove when implementing stubs
+        if is_stub:
+            return None
+
+        return PdbFunction(match_info, function_signature, is_stub)