[lld] [lld-macho] Implement ObjC category merging (-objc_category_merging) (PR #85727)

via llvm-commits llvm-commits at lists.llvm.org
Mon Mar 18 19:21:08 PDT 2024


llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-lld

Author: None (alx32)

<details>
<summary>Changes</summary>

This change adds a flag to lld to enable category merging for MachoO + ObjC.
It adds the '-objc_category_merging' flag for enabling this option and uses the existing '-no_objc_category_merging' flag for disabling it.
In ld64, this optimization is enabled by default, but in lld, for now, we require explicitly passing the '-objc_category_merging' flag in order to enable it.

Behavior: if in the same link unit, multiple categories are extending the same class, then they get merged into a single category.
Ex: Cat1(method1+method2,protocol1) + Cat2(method3+method4,protocol2, property1) = Cat1_2(method1+method2+method3+method4, protocol1+protocol2, property1)

Notes on implementation decisions made in this diff:

There is a possibility to further 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 improvement may be done as a follow-up. This improved functionality is already present in ld64.
We do the merging on the raw inputSections - after dead-stripping (categories can't be dead stripped anyway).
The changes are mostly self-contained to ObjC.cpp, except for adding 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 cause the old symbols to show up in the linker map as being dead-stripped - even if dead-stripping is disabled. This flag allows us to match the ld64 behavior.

Note: This is a re-land of https://github.com/llvm/llvm-project/pull/82928 after fixing using already freed memory in `generatedSectionData`.  This issue was detected by ASAN build.

---

Patch is 81.80 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/85727.diff


7 Files Affected:

- (modified) lld/MachO/Driver.cpp (+9) 
- (modified) lld/MachO/InputSection.h (+1-1) 
- (modified) lld/MachO/ObjC.cpp (+920-3) 
- (modified) lld/MachO/ObjC.h (+12) 
- (modified) lld/MachO/Options.td (+6-4) 
- (added) lld/test/MachO/objc-category-merging-complete-test.s (+762) 
- (added) lld/test/MachO/objc-category-merging-extern-class-minimal.s (+155) 


``````````diff
diff --git a/lld/MachO/Driver.cpp b/lld/MachO/Driver.cpp
index 9edb6b9c60a1fe..36248925d65ad2 100644
--- a/lld/MachO/Driver.cpp
+++ b/lld/MachO/Driver.cpp
@@ -1437,6 +1437,8 @@ bool link(ArrayRef<const char *> argsArr, llvm::raw_ostream &stdoutOS,
     resetOutputSegments();
     resetWriter();
     InputFile::resetIdCount();
+
+    objc::doCleanup();
   };
 
   ctx->e.logName = args::getFilenameWithoutExe(argsArr[0]);
@@ -1979,9 +1981,16 @@ bool link(ArrayRef<const char *> argsArr, llvm::raw_ostream &stdoutOS,
     if (config->deadStrip)
       markLive();
 
+    // Categories are not subject to dead-strip. The __objc_catlist section is
+    // marked as NO_DEAD_STRIP and that propagates into all category data.
     if (args.hasArg(OPT_check_category_conflicts))
       objc::checkCategories();
 
+    // Category merging uses "->live = false" to erase old category data, so
+    // it has to run after dead-stripping (markLive).
+    if (args.hasArg(OPT_objc_category_merging, OPT_no_objc_category_merging))
+      objc::mergeCategories();
+
     // ICF assumes that all literals have been folded already, so we must run
     // foldIdenticalLiterals before foldIdenticalSections.
     foldIdenticalLiterals();
diff --git a/lld/MachO/InputSection.h b/lld/MachO/InputSection.h
index becb01017d633a..b25f0638f4c6cb 100644
--- a/lld/MachO/InputSection.h
+++ b/lld/MachO/InputSection.h
@@ -93,9 +93,9 @@ class InputSection {
   // .subsections_via_symbols, there is typically only one element here.
   llvm::TinyPtrVector<Defined *> symbols;
 
-protected:
   const Section §ion;
 
+protected:
   const Defined *getContainingSymbol(uint64_t off) const;
 };
 
diff --git a/lld/MachO/ObjC.cpp b/lld/MachO/ObjC.cpp
index 67254ec53a2145..eef40590936ed3 100644
--- a/lld/MachO/ObjC.cpp
+++ b/lld/MachO/ObjC.cpp
@@ -7,16 +7,19 @@
 //===----------------------------------------------------------------------===//
 
 #include "ObjC.h"
+#include "ConcatOutputSection.h"
 #include "InputFiles.h"
 #include "InputSection.h"
 #include "Layout.h"
 #include "OutputSegment.h"
+#include "SyntheticSections.h"
 #include "Target.h"
 
 #include "lld/Common/ErrorHandler.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/BinaryFormat/MachO.h"
 #include "llvm/Bitcode/BitcodeReader.h"
+#include "llvm/Support/TimeProfiler.h"
 
 using namespace llvm;
 using namespace llvm::MachO;
@@ -78,7 +81,8 @@ namespace {
   DO(Ptr, classMethods)                                                        \
   DO(Ptr, protocols)                                                           \
   DO(Ptr, instanceProps)                                                       \
-  DO(Ptr, classProps)
+  DO(Ptr, classProps)                                                          \
+  DO(uint32_t, size)
 
 CREATE_LAYOUT_CLASS(Category, FOR_EACH_CATEGORY_FIELD);
 
@@ -112,13 +116,19 @@ CREATE_LAYOUT_CLASS(ROClass, FOR_EACH_RO_CLASS_FIELD);
 #undef FOR_EACH_RO_CLASS_FIELD
 
 #define FOR_EACH_LIST_HEADER(DO)                                               \
-  DO(uint32_t, size)                                                           \
-  DO(uint32_t, count)
+  DO(uint32_t, structSize)                                                     \
+  DO(uint32_t, structCount)
 
 CREATE_LAYOUT_CLASS(ListHeader, FOR_EACH_LIST_HEADER);
 
 #undef FOR_EACH_LIST_HEADER
 
+#define FOR_EACH_PROTOCOL_LIST_HEADER(DO) DO(Ptr, protocolCount)
+
+CREATE_LAYOUT_CLASS(ProtocolListHeader, FOR_EACH_PROTOCOL_LIST_HEADER);
+
+#undef FOR_EACH_PROTOCOL_LIST_HEADER
+
 #define FOR_EACH_METHOD(DO)                                                    \
   DO(Ptr, name)                                                                \
   DO(Ptr, type)                                                                \
@@ -311,6 +321,8 @@ void ObjcCategoryChecker::parseClass(const Defined *classSym) {
 }
 
 void objc::checkCategories() {
+  TimeTraceScope timeScope("ObjcCategoryChecker");
+
   ObjcCategoryChecker checker;
   for (const InputSection *isec : inputSections) {
     if (isec->getName() == section_names::objcCatList)
@@ -320,3 +332,908 @@ void objc::checkCategories() {
       }
   }
 }
+
+namespace {
+
+class ObjcCategoryMerger {
+  // Information about an input category
+  struct InfoInputCategory {
+    ConcatInputSection *catListIsec;
+    ConcatInputSection *catBodyIsec;
+    uint32_t offCatListIsec = 0;
+
+    bool wasMerged = false;
+  };
+
+  // To write new (merged) categories or classes, we will try make limited
+  // assumptions about the alignment and the sections the various class/category
+  // info are stored in and . So we'll just reuse the same sections and
+  // alignment as already used in existing (input) categories. To do this we
+  // have InfoCategoryWriter which contains the various sections that the
+  // generated categories will be written to.
+  template <typename T> struct InfoWriteSection {
+    bool valid = false; // Data has been successfully collected from input
+    uint32_t align = 0;
+    Section *inputSection;
+    Reloc relocTemplate;
+    T *outputSection;
+  };
+
+  struct InfoCategoryWriter {
+    InfoWriteSection<ConcatOutputSection> catListInfo;
+    InfoWriteSection<ConcatOutputSection> catBodyInfo;
+    InfoWriteSection<CStringSection> catNameInfo;
+    InfoWriteSection<ConcatOutputSection> catPtrListInfo;
+  };
+
+  // Information about a pointer list in the original categories (method lists,
+  // protocol lists, etc)
+  struct PointerListInfo {
+    PointerListInfo(const char *_categoryPrefix, uint32_t _categoryOffset,
+                    uint32_t _pointersPerStruct)
+        : categoryPrefix(_categoryPrefix), categoryOffset(_categoryOffset),
+          pointersPerStruct(_pointersPerStruct) {}
+    const char *categoryPrefix;
+    uint32_t categoryOffset = 0;
+
+    uint32_t pointersPerStruct = 0;
+
+    uint32_t structSize = 0;
+    uint32_t structCount = 0;
+
+    std::vector<Symbol *> allPtrs;
+  };
+
+  // Full information about all the categories that extend a class. This will
+  // include all the additional methods, protocols, and properties that are
+  // contained in all the categories that extend a particular class.
+  struct ClassExtensionInfo {
+    ClassExtensionInfo(CategoryLayout &_catLayout) : catLayout(_catLayout){};
+
+    // Merged names of containers. Ex: base|firstCategory|secondCategory|...
+    std::string mergedContainerName;
+    std::string baseClassName;
+    Symbol *baseClass = nullptr;
+    CategoryLayout &catLayout;
+
+    // In case we generate new data, mark the new data as belonging to this file
+    ObjFile *objFileForMergeData = nullptr;
+
+    PointerListInfo instanceMethods = {
+        objc::symbol_names::categoryInstanceMethods,
+        /*_categoryOffset=*/catLayout.instanceMethodsOffset,
+        /*pointersPerStruct=*/3};
+    PointerListInfo classMethods = {
+        objc::symbol_names::categoryClassMethods,
+        /*_categoryOffset=*/catLayout.classMethodsOffset,
+        /*pointersPerStruct=*/3};
+    PointerListInfo protocols = {objc::symbol_names::categoryProtocols,
+                                 /*_categoryOffset=*/catLayout.protocolsOffset,
+                                 /*pointersPerStruct=*/0};
+    PointerListInfo instanceProps = {
+        objc::symbol_names::listProprieties,
+        /*_categoryOffset=*/catLayout.instancePropsOffset,
+        /*pointersPerStruct=*/2};
+    PointerListInfo classProps = {
+        objc::symbol_names::klassPropList,
+        /*_categoryOffset=*/catLayout.classPropsOffset,
+        /*pointersPerStruct=*/2};
+  };
+
+public:
+  ObjcCategoryMerger(std::vector<ConcatInputSection *> &_allInputSections);
+  void doMerge();
+  static void doCleanup();
+
+private:
+  void collectAndValidateCategoriesData();
+  void
+  mergeCategoriesIntoSingleCategory(std::vector<InfoInputCategory> &categories);
+
+  void eraseISec(ConcatInputSection *isec);
+  void eraseMergedCategories();
+
+  void generateCatListForNonErasedCategories(
+      std::map<ConcatInputSection *, std::set<uint64_t>>
+          catListToErasedOffsets);
+  template <typename T>
+  void collectSectionWriteInfoFromIsec(const InputSection *isec,
+                                       InfoWriteSection<T> &catWriteInfo);
+  void collectCategoryWriterInfoFromCategory(const InfoInputCategory &catInfo);
+  void parseCatInfoToExtInfo(const InfoInputCategory &catInfo,
+                             ClassExtensionInfo &extInfo);
+
+  void parseProtocolListInfo(const ConcatInputSection *isec, uint32_t secOffset,
+                             PointerListInfo &ptrList);
+
+  void parsePointerListInfo(const ConcatInputSection *isec, uint32_t secOffset,
+                            PointerListInfo &ptrList);
+
+  void emitAndLinkPointerList(Defined *parentSym, uint32_t linkAtOffset,
+                              const ClassExtensionInfo &extInfo,
+                              const PointerListInfo &ptrList);
+
+  void emitAndLinkProtocolList(Defined *parentSym, uint32_t linkAtOffset,
+                               const ClassExtensionInfo &extInfo,
+                               const PointerListInfo &ptrList);
+
+  Defined *emitCategory(const ClassExtensionInfo &extInfo);
+  Defined *emitCatListEntrySec(const std::string &forCateogryName,
+                               const std::string &forBaseClassName,
+                               ObjFile *objFile);
+  Defined *emitCategoryBody(const std::string &name, const Defined *nameSym,
+                            const Symbol *baseClassSym,
+                            const std::string &baseClassName, ObjFile *objFile);
+  Defined *emitCategoryName(const std::string &name, ObjFile *objFile);
+  void createSymbolReference(Defined *refFrom, const Symbol *refTo,
+                             uint32_t offset, const Reloc &relocTemplate);
+  Symbol *tryGetSymbolAtIsecOffset(const ConcatInputSection *isec,
+                                   uint32_t offset);
+  Defined *tryGetDefinedAtIsecOffset(const ConcatInputSection *isec,
+                                     uint32_t offset);
+  void tryEraseDefinedAtIsecOffset(const ConcatInputSection *isec,
+                                   uint32_t offset);
+
+  // Allocate a null-terminated StringRef backed by generatedSectionData
+  StringRef newStringData(const char *str);
+  // Allocate section data, backed by generatedSectionData
+  SmallVector<uint8_t> &newSectionData(uint32_t size);
+
+  CategoryLayout catLayout;
+  ClassLayout classLayout;
+  ROClassLayout roClassLayout;
+  ListHeaderLayout listHeaderLayout;
+  MethodLayout methodLayout;
+  ProtocolListHeaderLayout protocolListHeaderLayout;
+
+  InfoCategoryWriter infoCategoryWriter;
+  std::vector<ConcatInputSection *> &allInputSections;
+  // Map of base class Symbol to list of InfoInputCategory's for it
+  DenseMap<const Symbol *, std::vector<InfoInputCategory>> categoryMap;
+
+  // 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.
+  // Need this to be 'static' so the data survives past the ObjcCategoryMerger
+  // object, as the data will be read by the Writer when the final binary is
+  // generated.
+  static SmallVector<std::unique_ptr<SmallVector<uint8_t>>>
+      generatedSectionData;
+};
+
+SmallVector<std::unique_ptr<SmallVector<uint8_t>>>
+    ObjcCategoryMerger::generatedSectionData;
+
+ObjcCategoryMerger::ObjcCategoryMerger(
+    std::vector<ConcatInputSection *> &_allInputSections)
+    : catLayout(target->wordSize), classLayout(target->wordSize),
+      roClassLayout(target->wordSize), listHeaderLayout(target->wordSize),
+      methodLayout(target->wordSize),
+      protocolListHeaderLayout(target->wordSize),
+      allInputSections(_allInputSections) {}
+
+// This is a template so that it can be used both for CStringSection and
+// ConcatOutputSection
+template <typename T>
+void ObjcCategoryMerger::collectSectionWriteInfoFromIsec(
+    const InputSection *isec, InfoWriteSection<T> &catWriteInfo) {
+
+  catWriteInfo.inputSection = const_cast<Section *>(&isec->section);
+  catWriteInfo.align = isec->align;
+  catWriteInfo.outputSection = dyn_cast_or_null<T>(isec->parent);
+
+  assert(catWriteInfo.outputSection &&
+         "outputSection may not be null in collectSectionWriteInfoFromIsec.");
+
+  if (isec->relocs.size())
+    catWriteInfo.relocTemplate = isec->relocs[0];
+
+  catWriteInfo.valid = true;
+}
+
+Symbol *
+ObjcCategoryMerger::tryGetSymbolAtIsecOffset(const ConcatInputSection *isec,
+                                             uint32_t offset) {
+  const Reloc *reloc = isec->getRelocAt(offset);
+
+  if (!reloc)
+    return nullptr;
+
+  return reloc->referent.get<Symbol *>();
+}
+
+Defined *
+ObjcCategoryMerger::tryGetDefinedAtIsecOffset(const ConcatInputSection *isec,
+                                              uint32_t offset) {
+  Symbol *sym = tryGetSymbolAtIsecOffset(isec, offset);
+  return dyn_cast_or_null<Defined>(sym);
+}
+
+// Given an ConcatInputSection or CStringInputSection and an offset, if there is
+// a symbol(Defined) at that offset, then erase the symbol (mark it not live)
+void ObjcCategoryMerger::tryEraseDefinedAtIsecOffset(
+    const ConcatInputSection *isec, uint32_t offset) {
+  const Reloc *reloc = isec->getRelocAt(offset);
+
+  if (!reloc)
+    return;
+
+  Defined *sym = dyn_cast_or_null<Defined>(reloc->referent.get<Symbol *>());
+  if (!sym)
+    return;
+
+  if (auto *cisec = dyn_cast_or_null<ConcatInputSection>(sym->isec))
+    eraseISec(cisec);
+  else if (auto *csisec = dyn_cast_or_null<CStringInputSection>(sym->isec)) {
+    uint32_t totalOffset = sym->value + reloc->addend;
+    StringPiece &piece = csisec->getStringPiece(totalOffset);
+    piece.live = false;
+  } else {
+    llvm_unreachable("erased symbol has to be Defined or CStringInputSection");
+  }
+}
+
+void ObjcCategoryMerger::collectCategoryWriterInfoFromCategory(
+    const InfoInputCategory &catInfo) {
+
+  if (!infoCategoryWriter.catListInfo.valid)
+    collectSectionWriteInfoFromIsec<ConcatOutputSection>(
+        catInfo.catListIsec, infoCategoryWriter.catListInfo);
+  if (!infoCategoryWriter.catBodyInfo.valid)
+    collectSectionWriteInfoFromIsec<ConcatOutputSection>(
+        catInfo.catBodyIsec, infoCategoryWriter.catBodyInfo);
+
+  if (!infoCategoryWriter.catNameInfo.valid) {
+    lld::macho::Defined *catNameSym =
+        tryGetDefinedAtIsecOffset(catInfo.catBodyIsec, catLayout.nameOffset);
+    assert(catNameSym && "Category does not have a valid name Symbol");
+
+    collectSectionWriteInfoFromIsec<CStringSection>(
+        catNameSym->isec, infoCategoryWriter.catNameInfo);
+  }
+
+  // Collect writer info from all the category lists (we're assuming they all
+  // would provide the same info)
+  if (!infoCategoryWriter.catPtrListInfo.valid) {
+    for (uint32_t off = catLayout.instanceMethodsOffset;
+         off <= catLayout.classPropsOffset; off += target->wordSize) {
+      if (Defined *ptrList =
+              tryGetDefinedAtIsecOffset(catInfo.catBodyIsec, off)) {
+        collectSectionWriteInfoFromIsec<ConcatOutputSection>(
+            ptrList->isec, infoCategoryWriter.catPtrListInfo);
+        // we've successfully collected data, so we can break
+        break;
+      }
+    }
+  }
+}
+
+// Parse a protocol list that might be linked to ConcatInputSection at a given
+// offset. The format of the protocol list is different than other lists (prop
+// lists, method lists) so we need to parse it differently
+void ObjcCategoryMerger::parseProtocolListInfo(const ConcatInputSection *isec,
+                                               uint32_t secOffset,
+                                               PointerListInfo &ptrList) {
+  if (!isec || (secOffset + target->wordSize > isec->data.size()))
+    assert("Tried to read pointer list beyond protocol section end");
+
+  const Reloc *reloc = isec->getRelocAt(secOffset);
+  if (!reloc)
+    return;
+
+  auto *ptrListSym = dyn_cast_or_null<Defined>(reloc->referent.get<Symbol *>());
+  assert(ptrListSym && "Protocol list reloc does not have a valid Defined");
+
+  // Theoretically protocol count can be either 32b or 64b, depending on
+  // platform pointer size, but to simplify implementation we always just read
+  // the lower 32b which should be good enough.
+  uint32_t protocolCount = *reinterpret_cast<const uint32_t *>(
+      ptrListSym->isec->data.data() + listHeaderLayout.structSizeOffset);
+
+  ptrList.structCount += protocolCount;
+  ptrList.structSize = target->wordSize;
+
+  uint32_t expectedListSize =
+      (protocolCount * target->wordSize) +
+      /*header(count)*/ protocolListHeaderLayout.totalSize +
+      /*extra null value*/ target->wordSize;
+  assert(expectedListSize == ptrListSym->isec->data.size() &&
+         "Protocol list does not match expected size");
+
+  uint32_t off = protocolListHeaderLayout.totalSize;
+  for (uint32_t inx = 0; inx < protocolCount; ++inx) {
+    const Reloc *reloc = ptrListSym->isec->getRelocAt(off);
+    assert(reloc && "No reloc found at protocol list offset");
+
+    auto *listSym = dyn_cast_or_null<Defined>(reloc->referent.get<Symbol *>());
+    assert(listSym && "Protocol list reloc does not have a valid Defined");
+
+    ptrList.allPtrs.push_back(listSym);
+    off += target->wordSize;
+  }
+  assert((ptrListSym->isec->getRelocAt(off) == nullptr) &&
+         "expected null terminating protocol");
+  assert(off + /*extra null value*/ target->wordSize == expectedListSize &&
+         "Protocol list end offset does not match expected size");
+}
+
+// Parse a pointer list that might be linked to ConcatInputSection at a given
+// offset. This can be used for instance methods, class methods, instance props
+// and class props since they have the same format.
+void ObjcCategoryMerger::parsePointerListInfo(const ConcatInputSection *isec,
+                                              uint32_t secOffset,
+                                              PointerListInfo &ptrList) {
+  assert(ptrList.pointersPerStruct == 2 || ptrList.pointersPerStruct == 3);
+  assert(isec && "Trying to parse pointer list from null isec");
+  assert(secOffset + target->wordSize <= isec->data.size() &&
+         "Trying to read pointer list beyond section end");
+
+  const Reloc *reloc = isec->getRelocAt(secOffset);
+  if (!reloc)
+    return;
+
+  auto *ptrListSym = dyn_cast_or_null<Defined>(reloc->referent.get<Symbol *>());
+  assert(ptrListSym && "Reloc does not have a valid Defined");
+
+  uint32_t thisStructSize = *reinterpret_cast<const uint32_t *>(
+      ptrListSym->isec->data.data() + listHeaderLayout.structSizeOffset);
+  uint32_t thisStructCount = *reinterpret_cast<const uint32_t *>(
+      ptrListSym->isec->data.data() + listHeaderLayout.structCountOffset);
+  assert(thisStructSize == ptrList.pointersPerStruct * target->wordSize);
+
+  assert(!ptrList.structSize || (thisStructSize == ptrList.structSize));
+
+  ptrList.structCount += thisStructCount;
+  ptrList.structSize = thisStructSize;
+
+  uint32_t expectedListSize =
+      listHeaderLayout.totalSize + (thisStructSize * thisStructCount);
+  assert(expectedListSize == ptrListSym->isec->data.size() &&
+         "Pointer list does not match expected size");
+
+  for (uint32_t off = listHeaderLayout.totalSize; off < expectedListSize;
+       off += target->wordSize) {
+    const Reloc *reloc = ptrListSym->isec->getRelocAt(off);
+    assert(reloc && "No reloc found at pointer list offset");
+
+    auto *listSym = dyn_cast_or_null<Defined>(reloc->referent.get<Symbol *>());
+    assert(listSym && "Reloc does not have a valid Defined");
+
+    ptrList.allPtrs.push_back(listSym);
+  }
+}
+
+// Here we parse all the information of an input category (catInfo) and
+// append the parsed info into the structure which will contain all the
+// information about how a class is extended (extInfo)
+void ObjcCategoryMerger::parseCatInfoToExtInfo(const InfoInputCategory &catInfo,
+                                               ClassExtensionInfo &extInfo) {
+  const Reloc *catNameReloc =
+      catInfo.catBodyIsec->getRelocAt(catLayout.nameOffset);
+
+  // Parse name
+  assert(catNameReloc && "Category does not have a reloc at 'nameOffset'");
+
+  // is this the first category we are parsing?
+  if (extInfo.mergedContainerName.empty())
+    extInfo.objFileForMergeData =
+        dyn_cast_or_null<ObjFile>(catInfo.catBod...
[truncated]

``````````

</details>


https://github.com/llvm/llvm-project/pull/85727


More information about the llvm-commits mailing list