[lld-macho][arm64] implement -objc_stubs_small (#78665)

kyulee-com · web-flow · commit 77e204c7b04f · 2024-01-23T07:31:34.000-08:00
This patch implements `-objc_stubs_small` targeting arm64, aiming to
align with ld64's behavior.
1. `-objc_stubs_fast`: As previously implemented, this always uses the
Global Offset Table (GOT) to invoke `objc_msgSend`. The alignment of the
objc stub is 32 bytes.
2. `-objc_stubs_small`: This behavior depends on whether `objc_msgSend`
is defined. If it is, it directly jumps to `objc_msgSend`. If not, it
creates another stub to indirectly jump to `objc_msgSend`, minimizing
the size. The alignment of the objc stub in this case is 4 bytes.
diff --git a/lld/MachO/Arch/ARM64.cpp b/lld/MachO/Arch/ARM64.cpp
@@ -37,9 +37,9 @@ struct ARM64 : ARM64Common {
                             uint64_t entryAddr) const override;
 
   void writeObjCMsgSendStub(uint8_t *buf, Symbol *sym, uint64_t stubsAddr,
-                            uint64_t stubOffset, uint64_t selrefsVA,
-                            uint64_t selectorIndex, uint64_t gotAddr,
-                            uint64_t msgSendIndex) const override;
+                            uint64_t &stubOffset, uint64_t selrefsVA,
+                            uint64_t selectorIndex,
+                            Symbol *objcMsgSend) const override;
   void populateThunk(InputSection *thunk, Symbol *funcSym) override;
   void applyOptimizationHints(uint8_t *, const ObjFile &) const override;
 };
@@ -117,13 +117,42 @@ static constexpr uint32_t objcStubsFastCode[] = {
     0xd4200020, // brk   #0x1
 };
 
+static constexpr uint32_t objcStubsSmallCode[] = {
+    0x90000001, // adrp  x1, __objc_selrefs@page
+    0xf9400021, // ldr   x1, [x1, @selector("foo")@pageoff]
+    0x14000000, // b     _objc_msgSend
+};
+
 void ARM64::writeObjCMsgSendStub(uint8_t *buf, Symbol *sym, uint64_t stubsAddr,
-                                 uint64_t stubOffset, uint64_t selrefsVA,
-                                 uint64_t selectorIndex, uint64_t gotAddr,
-                                 uint64_t msgSendIndex) const {
-  ::writeObjCMsgSendStub<LP64>(buf, objcStubsFastCode, sym, stubsAddr,
-                               stubOffset, selrefsVA, selectorIndex, gotAddr,
-                               msgSendIndex);
+                                 uint64_t &stubOffset, uint64_t selrefsVA,
+                                 uint64_t selectorIndex,
+                                 Symbol *objcMsgSend) const {
+  uint64_t objcMsgSendAddr;
+  uint64_t objcStubSize;
+  uint64_t objcMsgSendIndex;
+
+  if (config->objcStubsMode == ObjCStubsMode::fast) {
+    objcStubSize = target->objcStubsFastSize;
+    objcMsgSendAddr = in.got->addr;
+    objcMsgSendIndex = objcMsgSend->gotIndex;
+    ::writeObjCMsgSendFastStub<LP64>(buf, objcStubsFastCode, sym, stubsAddr,
+                                     stubOffset, selrefsVA, selectorIndex,
+                                     objcMsgSendAddr, objcMsgSendIndex);
+  } else {
+    assert(config->objcStubsMode == ObjCStubsMode::small);
+    objcStubSize = target->objcStubsSmallSize;
+    if (auto *d = dyn_cast<Defined>(objcMsgSend)) {
+      objcMsgSendAddr = d->getVA();
+      objcMsgSendIndex = 0;
+    } else {
+      objcMsgSendAddr = in.stubs->addr;
+      objcMsgSendIndex = objcMsgSend->stubsIndex;
+    }
+    ::writeObjCMsgSendSmallStub<LP64>(buf, objcStubsSmallCode, sym, stubsAddr,
+                                      stubOffset, selrefsVA, selectorIndex,
+                                      objcMsgSendAddr, objcMsgSendIndex);
+  }
+  stubOffset += objcStubSize;
 }
 
 // A thunk is the relaxed variation of stubCode. We don't need the
@@ -157,7 +186,9 @@ ARM64::ARM64() : ARM64Common(LP64()) {
   thunkSize = sizeof(thunkCode);
 
   objcStubsFastSize = sizeof(objcStubsFastCode);
-  objcStubsAlignment = 32;
+  objcStubsFastAlignment = 32;
+  objcStubsSmallSize = sizeof(objcStubsSmallCode);
+  objcStubsSmallAlignment = 4;
 
   // Branch immediate is two's complement 26 bits, which is implicitly
   // multiplied by 4 (since all functions are 4-aligned: The branch range
diff --git a/lld/MachO/Arch/ARM64Common.h b/lld/MachO/Arch/ARM64Common.h
@@ -154,10 +154,10 @@ inline void writeStubHelperEntry(uint8_t *buf8,
 
 template <class LP>
 inline void
-writeObjCMsgSendStub(uint8_t *buf, const uint32_t objcStubsFastCode[8],
-                     Symbol *sym, uint64_t stubsAddr, uint64_t stubOffset,
-                     uint64_t selrefsVA, uint64_t selectorIndex,
-                     uint64_t gotAddr, uint64_t msgSendIndex) {
+writeObjCMsgSendFastStub(uint8_t *buf, const uint32_t objcStubsFastCode[8],
+                         Symbol *sym, uint64_t stubsAddr, uint64_t stubOffset,
+                         uint64_t selrefsVA, uint64_t selectorIndex,
+                         uint64_t gotAddr, uint64_t msgSendIndex) {
   SymbolDiagnostic d = {sym, sym->getName()};
   auto *buf32 = reinterpret_cast<uint32_t *>(buf);
 
@@ -180,6 +180,30 @@ writeObjCMsgSendStub(uint8_t *buf, const uint32_t objcStubsFastCode[8],
   buf32[7] = objcStubsFastCode[7];
 }
 
+template <class LP>
+inline void
+writeObjCMsgSendSmallStub(uint8_t *buf, const uint32_t objcStubsSmallCode[3],
+                          Symbol *sym, uint64_t stubsAddr, uint64_t stubOffset,
+                          uint64_t selrefsVA, uint64_t selectorIndex,
+                          uint64_t msgSendAddr, uint64_t msgSendIndex) {
+  SymbolDiagnostic d = {sym, sym->getName()};
+  auto *buf32 = reinterpret_cast<uint32_t *>(buf);
+
+  auto pcPageBits = [stubsAddr, stubOffset](int i) {
+    return pageBits(stubsAddr + stubOffset + i * sizeof(uint32_t));
+  };
+
+  uint64_t selectorOffset = selectorIndex * LP::wordSize;
+  encodePage21(&buf32[0], d, objcStubsSmallCode[0],
+               pageBits(selrefsVA + selectorOffset) - pcPageBits(0));
+  encodePageOff12(&buf32[1], d, objcStubsSmallCode[1],
+                  selrefsVA + selectorOffset);
+  uint64_t msgSendStubVA = msgSendAddr + msgSendIndex * target->stubSize;
+  uint64_t pcVA = stubsAddr + stubOffset + 2 * sizeof(uint32_t);
+  encodeBranch26(&buf32[2], {nullptr, "objc_msgSend stub"},
+                 objcStubsSmallCode[2], msgSendStubVA - pcVA);
+}
+
 } // namespace lld::macho
 
 #endif
diff --git a/lld/MachO/Arch/ARM64_32.cpp b/lld/MachO/Arch/ARM64_32.cpp
@@ -34,9 +34,9 @@ struct ARM64_32 : ARM64Common {
   void writeStubHelperEntry(uint8_t *buf, const Symbol &,
                             uint64_t entryAddr) const override;
   void writeObjCMsgSendStub(uint8_t *buf, Symbol *sym, uint64_t stubsAddr,
-                            uint64_t stubOffset, uint64_t selrefsVA,
-                            uint64_t selectorIndex, uint64_t gotAddr,
-                            uint64_t msgSendIndex) const override;
+                            uint64_t &stubOffset, uint64_t selrefsVA,
+                            uint64_t selectorIndex,
+                            Symbol *objcMsgSend) const override;
 };
 
 } // namespace
@@ -100,10 +100,9 @@ void ARM64_32::writeStubHelperEntry(uint8_t *buf8, const Symbol &sym,
 }
 
 void ARM64_32::writeObjCMsgSendStub(uint8_t *buf, Symbol *sym,
-                                    uint64_t stubsAddr, uint64_t stubOffset,
+                                    uint64_t stubsAddr, uint64_t &stubOffset,
                                     uint64_t selrefsVA, uint64_t selectorIndex,
-                                    uint64_t gotAddr,
-                                    uint64_t msgSendIndex) const {
+                                    Symbol *objcMsgSend) const {
   fatal("TODO: implement this");
 }
 
diff --git a/lld/MachO/Arch/X86_64.cpp b/lld/MachO/Arch/X86_64.cpp
@@ -38,9 +38,9 @@ struct X86_64 : TargetInfo {
                             uint64_t entryAddr) const override;
 
   void writeObjCMsgSendStub(uint8_t *buf, Symbol *sym, uint64_t stubsAddr,
-                            uint64_t stubOffset, uint64_t selrefsVA,
-                            uint64_t selectorIndex, uint64_t gotAddr,
-                            uint64_t msgSendIndex) const override;
+                            uint64_t &stubOffset, uint64_t selrefsVA,
+                            uint64_t selectorIndex,
+                            Symbol *objcMsgSend) const override;
 
   void relaxGotLoad(uint8_t *loc, uint8_t type) const override;
   uint64_t getPageSize() const override { return 4 * 1024; }
@@ -182,16 +182,20 @@ static constexpr uint8_t objcStubsFastCode[] = {
 };
 
 void X86_64::writeObjCMsgSendStub(uint8_t *buf, Symbol *sym, uint64_t stubsAddr,
-                                  uint64_t stubOffset, uint64_t selrefsVA,
-                                  uint64_t selectorIndex, uint64_t gotAddr,
-                                  uint64_t msgSendIndex) const {
+                                  uint64_t &stubOffset, uint64_t selrefsVA,
+                                  uint64_t selectorIndex,
+                                  Symbol *objcMsgSend) const {
+  uint64_t objcMsgSendAddr = in.got->addr;
+  uint64_t objcMsgSendIndex = objcMsgSend->gotIndex;
+
   memcpy(buf, objcStubsFastCode, sizeof(objcStubsFastCode));
   SymbolDiagnostic d = {sym, sym->getName()};
   uint64_t stubAddr = stubsAddr + stubOffset;
   writeRipRelative(d, buf, stubAddr, 7,
                    selrefsVA + selectorIndex * LP64::wordSize);
   writeRipRelative(d, buf, stubAddr, 0xd,
-                   gotAddr + msgSendIndex * LP64::wordSize);
+                   objcMsgSendAddr + objcMsgSendIndex * LP64::wordSize);
+  stubOffset += target->objcStubsFastSize;
 }
 
 void X86_64::relaxGotLoad(uint8_t *loc, uint8_t type) const {
@@ -214,7 +218,7 @@ X86_64::X86_64() : TargetInfo(LP64()) {
   stubHelperEntrySize = sizeof(stubHelperEntry);
 
   objcStubsFastSize = sizeof(objcStubsFastCode);
-  objcStubsAlignment = 1;
+  objcStubsFastAlignment = 1;
 
   relocAttrs = {relocAttrsArray.data(), relocAttrsArray.size()};
 }
diff --git a/lld/MachO/Driver.cpp b/lld/MachO/Driver.cpp
@@ -827,9 +827,13 @@ static ObjCStubsMode getObjCStubsMode(const ArgList &args) {
   if (!arg)
     return ObjCStubsMode::fast;
 
-  if (arg->getOption().getID() == OPT_objc_stubs_small)
-    warn("-objc_stubs_small is not yet implemented, defaulting to "
-         "-objc_stubs_fast");
+  if (arg->getOption().getID() == OPT_objc_stubs_small) {
+    if (is_contained({AK_arm64e, AK_arm64}, config->arch()))
+      return ObjCStubsMode::small;
+    else
+      warn("-objc_stubs_small is not yet implemented, defaulting to "
+           "-objc_stubs_fast");
+  }
   return ObjCStubsMode::fast;
 }
 
diff --git a/lld/MachO/SyntheticSections.cpp b/lld/MachO/SyntheticSections.cpp
@@ -809,31 +809,49 @@ void StubHelperSection::setUp() {
 ObjCStubsSection::ObjCStubsSection()
     : SyntheticSection(segment_names::text, section_names::objcStubs) {
   flags = S_ATTR_SOME_INSTRUCTIONS | S_ATTR_PURE_INSTRUCTIONS;
-  align = target->objcStubsAlignment;
+  align = config->objcStubsMode == ObjCStubsMode::fast
+              ? target->objcStubsFastAlignment
+              : target->objcStubsSmallAlignment;
 }
 
 void ObjCStubsSection::addEntry(Symbol *sym) {
   assert(sym->getName().starts_with(symbolPrefix) && "not an objc stub");
   StringRef methname = sym->getName().drop_front(symbolPrefix.size());
   offsets.push_back(
       in.objcMethnameSection->getStringOffset(methname).outSecOff);
+
+  auto stubSize = config->objcStubsMode == ObjCStubsMode::fast
+                      ? target->objcStubsFastSize
+                      : target->objcStubsSmallSize;
   Defined *newSym = replaceSymbol<Defined>(
       sym, sym->getName(), nullptr, isec,
-      /*value=*/symbols.size() * target->objcStubsFastSize,
-      /*size=*/target->objcStubsFastSize,
+      /*value=*/symbols.size() * stubSize,
+      /*size=*/stubSize,
       /*isWeakDef=*/false, /*isExternal=*/true, /*isPrivateExtern=*/true,
       /*includeInSymtab=*/true, /*isReferencedDynamically=*/false,
       /*noDeadStrip=*/false);
   symbols.push_back(newSym);
 }
 
 void ObjCStubsSection::setUp() {
-  Symbol *objcMsgSend = symtab->addUndefined("_objc_msgSend", /*file=*/nullptr,
-                                             /*isWeakRef=*/false);
+  objcMsgSend = symtab->addUndefined("_objc_msgSend", /*file=*/nullptr,
+                                     /*isWeakRef=*/false);
+  if (auto *undefined = dyn_cast<Undefined>(objcMsgSend))
+    treatUndefinedSymbol(*undefined,
+                         "lazy binding (normally in libobjc.dylib)");
   objcMsgSend->used = true;
-  in.got->addEntry(objcMsgSend);
-  assert(objcMsgSend->isInGot());
-  objcMsgSendGotIndex = objcMsgSend->gotIndex;
+  if (config->objcStubsMode == ObjCStubsMode::fast) {
+    in.got->addEntry(objcMsgSend);
+    assert(objcMsgSend->isInGot());
+  } else {
+    assert(config->objcStubsMode == ObjCStubsMode::small);
+    // In line with ld64's behavior, when objc_msgSend is a direct symbol,
+    // we directly reference it.
+    // In other cases, typically when binding in libobjc.dylib,
+    // we generate a stub to invoke objc_msgSend.
+    if (!isa<Defined>(objcMsgSend))
+      in.stubs->addEntry(objcMsgSend);
+  }
 
   size_t size = offsets.size() * target->wordSize;
   uint8_t *selrefsData = bAlloc().Allocate<uint8_t>(size);
@@ -863,7 +881,10 @@ void ObjCStubsSection::setUp() {
 }
 
 uint64_t ObjCStubsSection::getSize() const {
-  return target->objcStubsFastSize * symbols.size();
+  auto stubSize = config->objcStubsMode == ObjCStubsMode::fast
+                      ? target->objcStubsFastSize
+                      : target->objcStubsSmallSize;
+  return stubSize * symbols.size();
 }
 
 void ObjCStubsSection::writeTo(uint8_t *buf) const {
@@ -875,8 +896,7 @@ void ObjCStubsSection::writeTo(uint8_t *buf) const {
     Defined *sym = symbols[i];
     target->writeObjCMsgSendStub(buf + stubOffset, sym, in.objcStubs->addr,
                                  stubOffset, in.objcSelrefs->getVA(), i,
-                                 in.got->addr, objcMsgSendGotIndex);
-    stubOffset += target->objcStubsFastSize;
+                                 objcMsgSend);
   }
 }
 
diff --git a/lld/MachO/SyntheticSections.h b/lld/MachO/SyntheticSections.h
@@ -336,7 +336,7 @@ class ObjCStubsSection final : public SyntheticSection {
 private:
   std::vector<Defined *> symbols;
   std::vector<uint32_t> offsets;
-  int objcMsgSendGotIndex = 0;
+  Symbol *objcMsgSend = nullptr;
 };
 
 // Note that this section may also be targeted by non-lazy bindings. In
diff --git a/lld/MachO/Target.h b/lld/MachO/Target.h
@@ -70,10 +70,9 @@ class TargetInfo {
                                     uint64_t entryAddr) const = 0;
 
   virtual void writeObjCMsgSendStub(uint8_t *buf, Symbol *sym,
-                                    uint64_t stubsAddr, uint64_t stubOffset,
+                                    uint64_t stubsAddr, uint64_t &stubOffset,
                                     uint64_t selrefsVA, uint64_t selectorIndex,
-                                    uint64_t gotAddr,
-                                    uint64_t msgSendIndex) const = 0;
+                                    Symbol *objcMsgSend) const = 0;
 
   // Symbols may be referenced via either the GOT or the stubs section,
   // depending on the relocation type. prepareSymbolRelocation() will set up the
@@ -121,7 +120,9 @@ class TargetInfo {
   size_t stubHelperHeaderSize;
   size_t stubHelperEntrySize;
   size_t objcStubsFastSize;
-  size_t objcStubsAlignment;
+  size_t objcStubsSmallSize;
+  size_t objcStubsFastAlignment;
+  size_t objcStubsSmallAlignment;
   uint8_t p2WordSize;
   size_t wordSize;
 
diff --git a/lld/test/MachO/arm64-objc-stubs-dyn.s b/lld/test/MachO/arm64-objc-stubs-dyn.s
@@ -0,0 +1,76 @@
+# REQUIRES: aarch64
+
+# RUN: llvm-mc -filetype=obj -triple=arm64-apple-darwin %s -o %t.o
+# RUN: %lld -arch arm64 -lSystem -U _objc_msgSend -o %t.out %t.o
+# RUN: llvm-otool -vs __TEXT __objc_stubs %t.out | FileCheck %s
+# RUN: %lld -arch arm64 -lSystem -U _objc_msgSend -o %t.out %t.o -dead_strip
+# RUN: llvm-otool -vs __TEXT __objc_stubs %t.out | FileCheck %s
+# RUN: %lld -arch arm64 -lSystem -U _objc_msgSend -o %t.out %t.o -objc_stubs_fast
+# RUN: llvm-otool -vs __TEXT __objc_stubs %t.out | FileCheck %s
+# RUN: %lld -arch arm64 -lSystem -U _objc_msgSend -o %t.out %t.o -objc_stubs_small
+# RUN: llvm-otool -vs __TEXT __stubs  %t.out | FileCheck %s --check-prefix=STUB
+# RUN: llvm-otool -vs __TEXT __objc_stubs %t.out | FileCheck %s --check-prefix=SMALL
+
+# Unlike arm64-objc-stubs.s, in this test, _objc_msgSend is not defined,
+#  which usually binds with libobjc.dylib.
+# 1. -objc_stubs_fast: No change as it uses GOT.
+# 2. -objc_stubs_small: Create a (shared) stub to invoke _objc_msgSend, to minimize the size.
+
+# CHECK: Contents of (__TEXT,__objc_stubs) section
+
+# CHECK-NEXT: _objc_msgSend$foo:
+# CHECK-NEXT: adrp    x1, 8 ; 0x100008000
+# CHECK-NEXT: ldr     x1, [x1, #0x10]
+# CHECK-NEXT: adrp    x16, 4 ; 0x100004000
+# CHECK-NEXT: ldr     x16, [x16]
+# CHECK-NEXT: br      x16
+# CHECK-NEXT: brk     #0x1
+# CHECK-NEXT: brk     #0x1
+# CHECK-NEXT: brk     #0x1
+
+# CHECK-NEXT: _objc_msgSend$length:
+# CHECK-NEXT: adrp    x1, 8 ; 0x100008000
+# CHECK-NEXT: ldr     x1, [x1, #0x18]
+# CHECK-NEXT: adrp    x16, 4 ; 0x100004000
+# CHECK-NEXT: ldr     x16, [x16]
+# CHECK-NEXT: br      x16
+# CHECK-NEXT: brk     #0x1
+# CHECK-NEXT: brk     #0x1
+# CHECK-NEXT: brk     #0x1
+
+# CHECK-EMPTY:
+
+# STUB: Contents of (__TEXT,__stubs) section
+# STUB-NEXT:  adrp    x16, 8 ; 0x100008000
+# STUB-NEXT:  ldr     x16, [x16]
+# STUB-NEXT:  br      x16
+
+# SMALL: Contents of (__TEXT,__objc_stubs) section
+# SMALL-NEXT: _objc_msgSend$foo:
+# SMALL-NEXT: adrp    x1, 8 ; 0x100008000
+# SMALL-NEXT: ldr     x1, [x1, #0x18]
+# SMALL-NEXT: b
+# SMALL-NEXT: _objc_msgSend$length:
+# SMALL-NEXT: adrp    x1, 8 ; 0x100008000
+# SMALL-NEXT: ldr     x1, [x1, #0x20]
+# SMALL-NEXT: b
+
+.section  __TEXT,__objc_methname,cstring_literals
+lselref1:
+  .asciz  "foo"
+lselref2:
+  .asciz  "bar"
+
+.section  __DATA,__objc_selrefs,literal_pointers,no_dead_strip
+.p2align  3
+.quad lselref1
+.quad lselref2
+
+.text
+
+.globl _main
+_main:
+  bl  _objc_msgSend$length
+  bl  _objc_msgSend$foo
+  bl  _objc_msgSend$foo
+  ret
diff --git a/lld/test/MachO/arm64-objc-stubs.s b/lld/test/MachO/arm64-objc-stubs.s
diff --git a/lld/test/MachO/x86-64-objc-stubs.s b/lld/test/MachO/x86-64-objc-stubs.s
