[lld] 5fa24ac - [lld-macho] Category Merger: add support for addrsig references (#90903)

[via llvm-commits]

Author: alx32
Date: 2024-05-06T09:45:32-07:00
New Revision: 5fa24ac277172045a81603bb56e66b5f4a27f99e

URL: https://github.com/llvm/llvm-project/commit/5fa24ac277172045a81603bb56e66b5f4a27f99e
DIFF: https://github.com/llvm/llvm-project/commit/5fa24ac277172045a81603bb56e66b5f4a27f99e.diff

LOG: [lld-macho] Category Merger: add support for addrsig references (#90903)

When generating categories, clang sometimes will generate references in
the `.addrsig` section to the various category data items. Since we may
erase such items after merging them, we also need to remove them from
the `.addrsig` section - otherwise this will cause runtime asserts with
the `.addrsig` section trying to access invalid data.

Implementation wise, we use a hashset to keep track of all erased
`InputSection`'s and then go through all `.addrsig` sections and remove
references to any erased `InputSection`.

Added: 
    

Modified: 
    lld/MachO/ObjC.cpp
    lld/test/MachO/objc-category-merging-extern-class-minimal.s

Removed: 
    


################################################################################
diff  --git a/lld/MachO/ObjC.cpp b/lld/MachO/ObjC.cpp
index 4760fffebe3b30..c60016f98ff492 100644
--- a/lld/MachO/ObjC.cpp
+++ b/lld/MachO/ObjC.cpp
@@ -420,6 +420,7 @@ class ObjcCategoryMerger {
   mergeCategoriesIntoSingleCategory(std::vector<InfoInputCategory> &categories);
 
   void eraseISec(ConcatInputSection *isec);
+  void removeRefsToErasedIsecs();
   void eraseMergedCategories();
 
   void generateCatListForNonErasedCategories(
@@ -478,6 +479,8 @@ class ObjcCategoryMerger {
   std::vector<ConcatInputSection *> &allInputSections;
   // Map of base class Symbol to list of InfoInputCategory's for it
   DenseMap<const Symbol *, std::vector<InfoInputCategory>> categoryMap;
+  // Set for tracking InputSection erased via eraseISec
+  DenseSet<InputSection *> erasedIsecs;
 
   // Normally, the binary data comes from the input files, but since we're
   // generating binary data ourselves, we use the below array to store it in.
@@ -518,6 +521,8 @@ void ObjcCategoryMerger::collectSectionWriteInfoFromIsec(
 Symbol *
 ObjcCategoryMerger::tryGetSymbolAtIsecOffset(const ConcatInputSection *isec,
                                              uint32_t offset) {
+  if (!isec)
+    return nullptr;
   const Reloc *reloc = isec->getRelocAt(offset);
 
   if (!reloc)
@@ -1141,6 +1146,8 @@ void ObjcCategoryMerger::generateCatListForNonErasedCategories(
 }
 
 void ObjcCategoryMerger::eraseISec(ConcatInputSection *isec) {
+  erasedIsecs.insert(isec);
+
   isec->live = false;
   for (auto &sym : isec->symbols)
     sym->used = false;
@@ -1175,6 +1182,7 @@ void ObjcCategoryMerger::eraseMergedCategories() {
         continue;
 
       eraseISec(catInfo.catBodyIsec);
+
       tryEraseDefinedAtIsecOffset(catInfo.catBodyIsec, catLayout.nameOffset);
       tryEraseDefinedAtIsecOffset(catInfo.catBodyIsec,
                                   catLayout.instanceMethodsOffset);
@@ -1188,6 +1196,33 @@ void ObjcCategoryMerger::eraseMergedCategories() {
                                   catLayout.instancePropsOffset);
     }
   }
+
+  removeRefsToErasedIsecs();
+}
+
+// The compiler may generate references to categories inside the addrsig
+// section. This function will erase these references.
+void ObjcCategoryMerger::removeRefsToErasedIsecs() {
+  for (InputSection *isec : inputSections) {
+    if (isec->getName() != section_names::addrSig)
+      continue;
+
+    auto removeRelocs = [this](Reloc &r) {
+      auto *isec = dyn_cast_or_null<ConcatInputSection>(
+          r.referent.dyn_cast<InputSection *>());
+      if (!isec) {
+        Defined *sym =
+            dyn_cast_or_null<Defined>(r.referent.dyn_cast<Symbol *>());
+        if (sym)
+          isec = dyn_cast<ConcatInputSection>(sym->isec());
+      }
+      if (!isec)
+        return false;
+      return erasedIsecs.count(isec) > 0;
+    };
+
+    llvm::erase_if(isec->relocs, removeRelocs);
+  }
 }
 
 void ObjcCategoryMerger::doMerge() {

diff  --git a/lld/test/MachO/objc-category-merging-extern-class-minimal.s b/lld/test/MachO/objc-category-merging-extern-class-minimal.s
index ea79f29a421c5c..796993799f2db0 100644
--- a/lld/test/MachO/objc-category-merging-extern-class-minimal.s
+++ b/lld/test/MachO/objc-category-merging-extern-class-minimal.s
@@ -153,3 +153,6 @@ L_OBJC_IMAGE_INFO:
 	.long	0
 	.long	96
 .subsections_via_symbols
+
+.addrsig
+.addrsig_sym __OBJC_$_CATEGORY_MyBaseClass_$_Category01