Fix Ghidra comment read fidelity

declib/decompilers/ghidra/interface.py

@@ -843,25 +843,46 @@ def _get_comment(self, addr) -> Optional[Comment]:
         return comments.get(addr, None)
 
     def _comments(self) -> Dict[int, Comment]:
+        from .compat.imports import CodeUnit
+
+        # Ghidra stores multiple comment slots per code unit, while declib has one
+        # portable Comment per address. Preserve all populated text and label the
+        # slots when more than one has to collapse into the same Comment.
+        ordered_types = (
+            (CodeUnit.PLATE_COMMENT, "PLATE", False),
+            (CodeUnit.PRE_COMMENT, "PRE", True),
+            (CodeUnit.EOL_COMMENT, "EOL", False),
+            (CodeUnit.POST_COMMENT, "POST", False),
+            (CodeUnit.REPEATABLE_COMMENT, "REPEATABLE", False),
+        )
+
         comments = {}
-        funcs_code_units = self.__function_code_units()
-        for code_units in funcs_code_units:
-            for code_unit in code_units:
-                # TODO: this could be bad if we have multiple comments at the same address (pre and eol)
-                # eol comment
-                eol_cmt = code_unit.getComment(0)
-                if eol_cmt:
-                    addr = int(code_unit.getAddress().getOffset())
-                    comments[addr] = Comment(
-                        addr=addr, comment=str(eol_cmt)
-                    )
-                # pre comment
-                pre_cmt = code_unit.getComment(1)
-                if pre_cmt:
-                    addr = int(code_unit.getAddress().getOffset())
-                    comments[addr] = Comment(
-                        addr=addr, comment=str(pre_cmt), decompiled=True
-                    )
+        listing = self.currentProgram.getListing()
+        for func in self.currentProgram.getFunctionManager().getFunctions(True):
+            for code_unit in listing.getCodeUnits(func.getBody(), True):
+                comment_entries = []
+                for cmt_type, label, is_decompiled_type in ordered_types:
+                    text = code_unit.getComment(cmt_type)
+                    if not text:
+                        continue
+                    comment_entries.append((label, str(text), is_decompiled_type))
+
+                if not comment_entries:
+                    continue
+
+                should_prefix = len(comment_entries) > 1
+                parts = [
+                    f"[{label}] {text}" if should_prefix else text
+                    for label, text, _ in comment_entries
+                ]
+                has_decompiled = any(is_decompiled for _, _, is_decompiled in comment_entries)
+                has_disassembly = any(not is_decompiled for _, _, is_decompiled in comment_entries)
+                addr = int(code_unit.getAddress().getOffset())
+                comments[addr] = Comment(
+                    addr=addr,
+                    comment="\n".join(parts),
+                    decompiled=has_decompiled and not has_disassembly,
+                )
 
         return comments
 
@@ -1425,13 +1446,3 @@ def __gtypedefs(self):
             for typedef in self.currentProgram.getDataTypeManager().getAllDataTypes()
             if isinstance(typedef, TypedefDB)
         ]
-
-
-    def __function_code_units(self):
-        """
-        Returns a list of code units for each function in the program.
-        """
-        return [
-            [code_unit for code_unit in self.currentProgram.getListing().getCodeUnits(func.getBody(), True)]
-            for func in self.currentProgram.getFunctionManager().getFunctions(True)
-        ]

tests/test_decompilers.py

@@ -104,6 +104,114 @@ def test_ghidra_server_dispatch_requirement(self):
                             self.assertIs(deci.requires_main_thread_dispatch, expected)
                             self.assertIs(server.requires_main_thread, expected)
 
+    def test_ghidra_comment_slots(self):
+        import declib.decompilers.ghidra as ghidra_package
+
+        class CodeUnit:
+            EOL_COMMENT = 0
+            PRE_COMMENT = 1
+            POST_COMMENT = 2
+            PLATE_COMMENT = 3
+            REPEATABLE_COMMENT = 4
+
+        class JavaComment:
+            def __str__(self):
+                return "post"
+
+        pyghidra = types.ModuleType("pyghidra")
+        pyghidra_core = types.ModuleType("pyghidra.core")
+        pyghidra_core._analyze_program = MagicMock()
+        pyghidra_core._get_language = MagicMock()
+        pyghidra_core._get_compiler_spec = MagicMock()
+        jpype = types.ModuleType("jpype")
+        jpype.JClass = type
+        compat_imports = types.ModuleType("declib.decompilers.ghidra.compat.imports")
+        compat_imports.CodeUnit = CodeUnit
+
+        cases = (
+            ("empty", {}, None, None),
+            ("plate", {CodeUnit.PLATE_COMMENT: "plate"}, "plate", False),
+            ("pre", {CodeUnit.PRE_COMMENT: "pre"}, "pre", True),
+            ("eol", {CodeUnit.EOL_COMMENT: "eol"}, "eol", False),
+            ("post", {CodeUnit.POST_COMMENT: JavaComment()}, "post", False),
+            ("repeatable", {CodeUnit.REPEATABLE_COMMENT: "repeatable"}, "repeatable", False),
+            (
+                "plate-and-pre",
+                {CodeUnit.PLATE_COMMENT: "plate", CodeUnit.PRE_COMMENT: "pre"},
+                "[PLATE] plate\n[PRE] pre",
+                False,
+            ),
+            (
+                "all-slots",
+                {
+                    CodeUnit.PLATE_COMMENT: "plate",
+                    CodeUnit.PRE_COMMENT: "pre",
+                    CodeUnit.EOL_COMMENT: "eol",
+                    CodeUnit.POST_COMMENT: "post",
+                    CodeUnit.REPEATABLE_COMMENT: "repeatable",
+                },
+                "[PLATE] plate\n[PRE] pre\n[EOL] eol\n[POST] post\n[REPEATABLE] repeatable",
+                False,
+            ),
+        )
+
+        modules = {
+            "pyghidra": pyghidra,
+            "pyghidra.core": pyghidra_core,
+            "jpype": jpype,
+            "declib.decompilers.ghidra.compat.imports": compat_imports,
+        }
+        with patch.object(ghidra_package, "interface", create=True):
+            with patch.dict(sys.modules, modules):
+                sys.modules.pop("declib.decompilers.ghidra.interface", None)
+                from declib.decompilers.ghidra.interface import GhidraDecompilerInterface
+
+                class TestableGhidraDecompilerInterface(GhidraDecompilerInterface):
+                    @property
+                    def currentProgram(self):
+                        return self._program_for_test
+
+                code_units = []
+                for index, (_, slot_comments, _, _) in enumerate(cases):
+                    code_unit = MagicMock()
+                    code_unit.getComment.side_effect = slot_comments.get
+                    code_unit.getAddress.return_value.getOffset.return_value = 0x401000 + index * 0x10
+                    code_units.append(code_unit)
+
+                body = object()
+                func = MagicMock()
+                func.getBody.return_value = body
+                program = MagicMock()
+                program.getFunctionManager.return_value.getFunctions.return_value = [func]
+                program.getListing.return_value.getCodeUnits.return_value = code_units
+
+                deci = object.__new__(TestableGhidraDecompilerInterface)
+                deci._program_for_test = program
+                comments = deci._comments()
+
+                expected_addresses = {
+                    0x401000 + index * 0x10
+                    for index, (_, _, expected_text, _) in enumerate(cases)
+                    if expected_text is not None
+                }
+                self.assertEqual(set(comments), expected_addresses)
+                for index, (name, _, expected_text, expected_decompiled) in enumerate(cases):
+                    with self.subTest(case=name):
+                        addr = 0x401000 + index * 0x10
+                        if expected_text is None:
+                            self.assertNotIn(addr, comments)
+                            continue
+
+                        comment = comments[addr]
+                        self.assertEqual(comment.addr, addr)
+                        self.assertIsNone(comment.func_addr)
+                        self.assertEqual(comment.comment, expected_text)
+                        self.assertEqual(comment.decompiled, expected_decompiled)
+
+                program.getFunctionManager.return_value.getFunctions.assert_called_once_with(True)
+                func.getBody.assert_called_once_with()
+                program.getListing.return_value.getCodeUnits.assert_called_once_with(body, True)
+
     def test_readme_example(self):
         # TODO: add angr
         for dec_name in [IDA_DECOMPILER, GHIDRA_DECOMPILER, BINJA_DECOMPILER]: