[lld-macho] icf objc stubs

Author: kyulee-com

lld/MachO/SyntheticSections.cpp

@@ -825,10 +825,43 @@ StringRef ObjCStubsSection::getMethName(Symbol *sym) {
   return methname;
 }
 
+void ObjCStubsSection::initialize() {
+  // Do not fold selrefs without ICF.
+  if (config->icfLevel == ICFLevel::none)
+    return;
+
+  // Search methnames already referenced in __objc_selrefs
+  // Map the name to the corresponding selref entry
+  // which we will reuse when creating objc stubs.
+  for (ConcatInputSection *isec : inputSections) {
+    if (isec->shouldOmitFromOutput())
+      continue;
+    if (isec->getName() != "__objc_selrefs")
+      continue;
+    // We expect a single relocation per selref entry to __objc_methname that
+    // might be aggregated.
+    assert(isec->relocs.size() == 1);
+    auto Reloc = isec->relocs[0];
+    if (const auto *sym = Reloc.referent.dyn_cast<Symbol *>()) {
+      if (const auto *d = dyn_cast<Defined>(sym)) {
+        auto *cisec = cast<CStringInputSection>(d->isec);
+        auto methname = cisec->getStringRefAtOffset(d->value);
+        methnameToselref[methname] = isec;
+      }
+    }
+  }
+}
+
 void ObjCStubsSection::addEntry(Symbol *sym) {
   StringRef methname = getMethName(sym);
-  offsets.push_back(
-      in.objcMethnameSection->getStringOffset(methname).outSecOff);
+  // We will create a selref entry for each unique methname.
+  if (!methnameToselref.count(methname) &&
+      !methnameToidxOffsetPair.count(methname)) {
+    auto methnameOffset =
+        in.objcMethnameSection->getStringOffset(methname).outSecOff;
+    auto selIndex = methnameToidxOffsetPair.size();
+    methnameToidxOffsetPair[methname] = {selIndex, methnameOffset};
+  }
 
   auto stubSize = config->objcStubsMode == ObjCStubsMode::fast
                       ? target->objcStubsFastSize
@@ -863,10 +896,12 @@ void ObjCStubsSection::setUp() {
       in.stubs->addEntry(objcMsgSend);
   }
 
-  size_t size = offsets.size() * target->wordSize;
+  size_t size = methnameToidxOffsetPair.size() * target->wordSize;
   uint8_t *selrefsData = bAlloc().Allocate<uint8_t>(size);
-  for (size_t i = 0, n = offsets.size(); i < n; ++i)
-    write64le(&selrefsData[i * target->wordSize], offsets[i]);
+  for (const auto &[methname, idxOffsetPair] : methnameToidxOffsetPair) {
+    const auto &[idx, offset] = idxOffsetPair;
+    write64le(&selrefsData[idx * target->wordSize], offset);
+  }
 
   in.objcSelrefs =
       makeSyntheticInputSection(segment_names::data, section_names::objcSelrefs,
@@ -875,12 +910,13 @@ void ObjCStubsSection::setUp() {
                                 /*align=*/target->wordSize);
   in.objcSelrefs->live = true;
 
-  for (size_t i = 0, n = offsets.size(); i < n; ++i) {
+  for (const auto &[methname, idxOffsetPair] : methnameToidxOffsetPair) {
+    const auto &[idx, offset] = idxOffsetPair;
     in.objcSelrefs->relocs.push_back(
         {/*type=*/target->unsignedRelocType,
          /*pcrel=*/false, /*length=*/3,
-         /*offset=*/static_cast<uint32_t>(i * target->wordSize),
-         /*addend=*/offsets[i] * in.objcMethnameSection->align,
+         /*offset=*/static_cast<uint32_t>(idx * target->wordSize),
+         /*addend=*/offset * in.objcMethnameSection->align,
          /*referent=*/in.objcMethnameSection->isec});
   }
 
@@ -904,8 +940,22 @@ void ObjCStubsSection::writeTo(uint8_t *buf) const {
   uint64_t stubOffset = 0;
   for (size_t i = 0, n = symbols.size(); i < n; ++i) {
     Defined *sym = symbols[i];
+    uint64_t selBaseAddr;
+    uint64_t selIndex;
+
+    auto methname = getMethName(sym);
+    auto j = methnameToselref.find(methname);
+    if (j != methnameToselref.end()) {
+      selBaseAddr = j->second->getVA(0);
+      selIndex = 0;
+    } else {
+      auto k = methnameToidxOffsetPair.find(methname);
+      assert(k != methnameToidxOffsetPair.end());
+      selBaseAddr = in.objcSelrefs->getVA();
+      selIndex = k->second.first;
+    }
     target->writeObjCMsgSendStub(buf + stubOffset, sym, in.objcStubs->addr,
-                                 stubOffset, in.objcSelrefs->getVA(), i,
+                                 stubOffset, selBaseAddr, selIndex,
                                  objcMsgSend);
   }
 }

lld/MachO/SyntheticSections.h

@@ -324,6 +324,7 @@ class StubHelperSection final : public SyntheticSection {
 class ObjCStubsSection final : public SyntheticSection {
 public:
   ObjCStubsSection();
+  void initialize();
   void addEntry(Symbol *sym);
   uint64_t getSize() const override;
   bool isNeeded() const override { return !symbols.empty(); }
@@ -337,7 +338,12 @@ class ObjCStubsSection final : public SyntheticSection {
 
 private:
   std::vector<Defined *> symbols;
-  std::vector<uint32_t> offsets;
+  // Existing mapping from methname to selref (0 index is assumed).
+  llvm::StringMap<InputSection *> methnameToselref;
+  // Newly created mapping from methname to the pair of index (selref) and
+  // offset (methname).
+  llvm::MapVector<StringRef, std::pair<uint32_t, uint32_t>>
+      methnameToidxOffsetPair;
   Symbol *objcMsgSend = nullptr;
 };
 

lld/MachO/Writer.cpp

@@ -720,6 +720,7 @@ static void addNonWeakDefinition(const Defined *defined) {
 
 void Writer::scanSymbols() {
   TimeTraceScope timeScope("Scan symbols");
+  in.objcStubs->initialize();
   for (Symbol *sym : symtab->getSymbols()) {
     if (auto *defined = dyn_cast<Defined>(sym)) {
       if (!defined->isLive())

lld/test/MachO/objc-selrefs.s

@@ -24,7 +24,6 @@
 # SELREFS-NEXT: __TEXT:__objc_methname:length
 # SELREFS-EMPTY:
 
-## We don't yet support dedup'ing implicitly-defined selrefs.
 # RUN: %lld -dylib -arch arm64 -lSystem --icf=all -o %t/explicit-and-implicit \
 # RUN:   %t/explicit-selrefs-1.o %t/explicit-selrefs-2.o %t/implicit-selrefs.o
 # RUN: llvm-otool -vs __DATA __objc_selrefs %t/explicit-and-implicit \
@@ -43,8 +42,6 @@
 # EXPLICIT-AND-IMPLICIT:      Contents of (__DATA,__objc_selrefs) section
 # EXPLICIT-AND-IMPLICIT-NEXT: __TEXT:__objc_methname:foo
 # EXPLICIT-AND-IMPLICIT-NEXT: __TEXT:__objc_methname:bar
-# NOTE: Ideally this wouldn't exist, but while it does it needs to point to the deduplicated string
-# EXPLICIT-AND-IMPLICIT-NEXT: __TEXT:__objc_methname:foo
 # EXPLICIT-AND-IMPLICIT-NEXT: __TEXT:__objc_methname:length
 
 #--- explicit-selrefs-1.s

lld/test/MachO/x86-64-objc-stubs.s

@@ -10,6 +10,9 @@
 
 # WARNING: warning: -objc_stubs_small is not yet implemented, defaulting to -objc_stubs_fast
 
+# RUN: %lld -arch x86_64 -lSystem -o %t-icf.out --icf=all %t.o
+# RUN: llvm-otool -vs __DATA __objc_selrefs %t-icf.out | FileCheck %s --check-prefix=ICF
+
 # CHECK: Sections:
 # CHECK: __got            {{[0-9a-f]*}} [[#%x, GOTSTART:]] DATA
 # CHECK: __objc_selrefs   {{[0-9a-f]*}} [[#%x, SELSTART:]] DATA
@@ -21,6 +24,13 @@
 # CHECK-NEXT: [[#%x, FOOSELREF:]]  __TEXT:__objc_methname:foo
 # CHECK-NEXT: [[#%x, LENGTHSELREF:]]  __TEXT:__objc_methname:length
 
+# ICF: Contents of (__DATA,__objc_selrefs) section
+
+# ICF-NEXT: {{[0-9a-f]*}}  __TEXT:__objc_methname:foo
+# ICF-NEXT: {{[0-9a-f]*}}  __TEXT:__objc_methname:bar
+# ICF-NEXT: {{[0-9a-f]*}}  __TEXT:__objc_methname:length
+# ICF-EMPTY:
+
 # CHECK: Contents of (__TEXT,__objc_stubs) section
 
 # CHECK-NEXT: _objc_msgSend$foo: