[lld-macho] Implement ObjC category merging (-merge-objc-categories)

This change adds a flag to lld to enable category merging for MachoO + ObjC. If in the same link unit, multiple categories are extending the same class, then they get merged into a single cateogry.

Notes on implemetation decisions made in this diff:
1. There is a possibility to even improve the current implementation by directly merging the category data into the base class (if the base class is present in the link unit) - this may be done as a follow-up.
2. We do the merging as early as possible, on the raw inputSections.
3. We add a new flag for ObjFile (isLinkerGenerated) and create such an ObjFile to which all new linker-generated date belongs.
4. We add a new flag (linkerOptimizeReason) to ConcatInputSection and StringPiece to mark that this data has been optimized away. Another way to do it would have been to just mark the pieces as not 'live' but this would require some work-arounds in the actual live symbol determination logic and would also cause symbols to incorrectly show up as 'dead-stripped' when that's not the cause that they are not present.

Author: Alex B

lld/MachO/Driver.cpp

@@ -1973,15 +1973,20 @@ bool link(ArrayRef<const char *> argsArr, llvm::raw_ostream &stdoutOS,
     }
 
     gatherInputSections();
+
+    // Run category checking & merging before anything else, it operates
+    // directly on inputSections.
+    if (args.hasArg(OPT_check_category_conflicts))
+      objc::checkCategories();
+
+    if (args.hasArg(OPT_merge_objc_categories))
+      objc::mergeCategories();
+
     if (config->callGraphProfileSort)
       priorityBuilder.extractCallGraphProfile();
 
     if (config->deadStrip)
       markLive();
-
-    if (args.hasArg(OPT_check_category_conflicts))
-      objc::checkCategories();
-
     // ICF assumes that all literals have been folded already, so we must run
     // foldIdenticalLiterals before foldIdenticalSections.
     foldIdenticalLiterals();

lld/MachO/InputFiles.cpp

@@ -965,21 +965,23 @@ void ObjFile::parseLinkerOptions(SmallVectorImpl<StringRef> &LCLinkerOptions) {
 SmallVector<StringRef> macho::unprocessedLCLinkerOptions;
 ObjFile::ObjFile(MemoryBufferRef mb, uint32_t modTime, StringRef archiveName,
                  bool lazy, bool forceHidden, bool compatArch,
-                 bool builtFromBitcode)
+                 bool builtFromBitcode, bool isLinkerGenerated)
     : InputFile(ObjKind, mb, lazy), modTime(modTime), forceHidden(forceHidden),
-      builtFromBitcode(builtFromBitcode) {
+      builtFromBitcode(builtFromBitcode), isLinkerGenerated(isLinkerGenerated) {
   this->archiveName = std::string(archiveName);
   this->compatArch = compatArch;
-  if (lazy) {
-    if (target->wordSize == 8)
-      parseLazy<LP64>();
-    else
-      parseLazy<ILP32>();
-  } else {
-    if (target->wordSize == 8)
-      parse<LP64>();
-    else
-      parse<ILP32>();
+  if (!isLinkerGenerated) {
+    if (lazy) {
+      if (target->wordSize == 8)
+        parseLazy<LP64>();
+      else
+        parseLazy<ILP32>();
+    } else {
+      if (target->wordSize == 8)
+        parse<LP64>();
+      else
+        parse<ILP32>();
+    }
   }
 }
 
@@ -1103,6 +1105,8 @@ void ObjFile::parseDebugInfo() {
 }
 
 ArrayRef<data_in_code_entry> ObjFile::getDataInCode() const {
+  if (!mb.getBufferSize())
+    return {};
   const auto *buf = reinterpret_cast<const uint8_t *>(mb.getBufferStart());
   const load_command *cmd = findCommand(buf, LC_DATA_IN_CODE);
   if (!cmd)
@@ -1113,6 +1117,8 @@ ArrayRef<data_in_code_entry> ObjFile::getDataInCode() const {
 }
 
 ArrayRef<uint8_t> ObjFile::getOptimizationHints() const {
+  if (!mb.getBufferSize())
+    return {};
   const auto *buf = reinterpret_cast<const uint8_t *>(mb.getBufferStart());
   if (auto *cmd =
           findCommand<linkedit_data_command>(buf, LC_LINKER_OPTIMIZATION_HINT))

lld/MachO/InputFiles.h

@@ -161,7 +161,7 @@ class ObjFile final : public InputFile {
 public:
   ObjFile(MemoryBufferRef mb, uint32_t modTime, StringRef archiveName,
           bool lazy = false, bool forceHidden = false, bool compatArch = true,
-          bool builtFromBitcode = false);
+          bool builtFromBitcode = false, bool isLinkerGenerated = false);
   ArrayRef<llvm::MachO::data_in_code_entry> getDataInCode() const;
   ArrayRef<uint8_t> getOptimizationHints() const;
   template <class LP> void parse();
@@ -181,6 +181,7 @@ class ObjFile final : public InputFile {
   const uint32_t modTime;
   bool forceHidden;
   bool builtFromBitcode;
+  bool isLinkerGenerated;
   std::vector<ConcatInputSection *> debugSections;
   std::vector<CallGraphEntry> callGraph;
   llvm::DenseMap<ConcatInputSection *, FDE> fdes;

lld/MachO/InputSection.h

@@ -24,6 +24,11 @@
 namespace lld {
 namespace macho {
 
+enum LinkerOptReason : uint8_t {
+  NotOptimized,
+  CategoryMerging,
+};
+
 class InputFile;
 class OutputSection;
 
@@ -60,6 +65,7 @@ class InputSection {
   // Whether the data at \p off in this InputSection is live.
   virtual bool isLive(uint64_t off) const = 0;
   virtual void markLive(uint64_t off) = 0;
+  virtual bool isLinkOptimizedAway() const { return false; }
   virtual InputSection *canonical() { return this; }
   virtual const InputSection *canonical() const { return this; }
 
@@ -93,9 +99,9 @@ class InputSection {
   // .subsections_via_symbols, there is typically only one element here.
   llvm::TinyPtrVector<Defined *> symbols;
 
-protected:
   const Section &section;
 
+protected:
   const Defined *getContainingSymbol(uint64_t off) const;
 };
 
@@ -114,7 +120,12 @@ class ConcatInputSection final : public InputSection {
   bool isLive(uint64_t off) const override { return live; }
   void markLive(uint64_t off) override { live = true; }
   bool isCoalescedWeak() const { return wasCoalesced && symbols.empty(); }
-  bool shouldOmitFromOutput() const { return !live || isCoalescedWeak(); }
+  bool isLinkOptimizedAway() const override {
+    return linkerOptimizeReason != LinkerOptReason::NotOptimized;
+  }
+  bool shouldOmitFromOutput() const {
+    return isLinkOptimizedAway() || !live || isCoalescedWeak();
+  }
   void writeTo(uint8_t *buf);
 
   void foldIdentical(ConcatInputSection *redundant);
@@ -141,6 +152,11 @@ class ConcatInputSection final : public InputSection {
   // first and not copied to the output.
   bool wasCoalesced = false;
   bool live = !config->deadStrip;
+  // Flag to specify if a linker optimzation flagged this section to be
+  // discarded. Need a separate flag from live as live specifically means
+  // 'dead-stripped' which is rellevant in contexts such as linker map
+  // generation
+  LinkerOptReason linkerOptimizeReason = LinkerOptReason::NotOptimized;
   bool hasCallSites = false;
   // This variable has two usages. Initially, it represents the input order.
   // After assignAddresses is called, it represents the offset from the
@@ -176,10 +192,18 @@ struct StringPiece {
   // Only set if deduplicating literals
   uint32_t hash : 31;
   // Offset from the start of the containing output section.
-  uint64_t outSecOff = 0;
+  uint64_t outSecOff : 56;
+  LinkerOptReason linkerOptimizeReason : 8;
+
+  bool shouldOmitFromOutput() const {
+    return !live || linkerOptimizeReason != LinkerOptReason::NotOptimized;
+  }
 
   StringPiece(uint64_t off, uint32_t hash)
-      : inSecOff(off), live(!config->deadStrip), hash(hash) {}
+      : inSecOff(off), live(!config->deadStrip), hash(hash) {
+    outSecOff = 0;
+    linkerOptimizeReason = LinkerOptReason::NotOptimized;
+  }
 };
 
 static_assert(sizeof(StringPiece) == 16, "StringPiece is too big!");

lld/MachO/MapFile.cpp

@@ -80,7 +80,7 @@ static MapInfo gatherMapInfo() {
           if (d->isec && d->getFile() == file &&
               !isa<CStringInputSection>(d->isec)) {
             isReferencedFile = true;
-            if (!d->isLive())
+            if (!d->isLive() && (!d->isec || !d->isec->isLinkOptimizedAway()))
               info.deadSymbols.push_back(d);
           }
       }
@@ -93,6 +93,8 @@ static MapInfo gatherMapInfo() {
           if (auto isec = dyn_cast<CStringInputSection>(subsec.isec)) {
             auto &liveCStrings = info.liveCStringsForSection[isec->parent];
             for (const auto &[i, piece] : llvm::enumerate(isec->pieces)) {
+              if (piece.linkerOptimizeReason != LinkerOptReason::NotOptimized)
+                continue;
               if (piece.live)
                 liveCStrings.push_back({isec->parent->addr + piece.outSecOff,
                                         {fileIndex, isec->getStringRef(i)}});
@@ -203,6 +205,8 @@ void macho::writeMapFile() {
     for (const OutputSection *osec : seg->getSections()) {
       if (auto *concatOsec = dyn_cast<ConcatOutputSection>(osec)) {
         for (const InputSection *isec : concatOsec->inputs) {
+          if (isec->isLinkOptimizedAway())
+            continue;
           for (Defined *sym : isec->symbols)
             if (!(isPrivateLabel(sym->getName()) && sym->size == 0))
               os << format("0x%08llX\t0x%08llX\t[%3u] %s\n", sym->getVA(),

lld/MachO/MarkLive.cpp

@@ -259,6 +259,8 @@ void markLive() {
           dyn_cast_or_null<DylibSymbol>(symtab->find("dyld_stub_binder")))
     marker->addSym(stubBinder);
   for (ConcatInputSection *isec : inputSections) {
+    if (isec->isLinkOptimizedAway())
+      continue;
     // Sections marked no_dead_strip
     if (isec->getFlags() & S_ATTR_NO_DEAD_STRIP) {
       marker->enqueue(isec, 0);