[lld][macho] Fix thunks with multiple text sections (#199747)

When there are multiple `__text` sections, LLD might not generate thunks
to stubs sections when they are required, leading to relocation errors.
```
ld64.lld: error: a.o:(symbol _foo+0x0): relocation BRANCH26 is out of range: 134217744 is not in [-134217728, 134217727]; references _extern_sym
```

Create `TextOutputSection::estimateStubsEndVA()` to correctly estimate
the end VA of the last stubs section so we can tell when branches to
stub symbols will be in range.

Technically this could cause lld to generate more thunks in some cases.
If a binary requires thunks (the `__TEXT` segment is >128MiB) and has
multiple `__text` sections like `__text_cold` or `__lcxx_override`, then
all branches to stub symbols will require thunks. Without this change we
could get the relocation errors above. We might be able to workaround
the problem by placing the `__text` section last.

Fixes https://github.com/llvm/llvm-project/issues/195387.

lld/MachO/ConcatOutputSection.cpp

@@ -65,6 +65,9 @@ DenseMap<ThunkKey, ThunkInfo, ThunkMapKeyInfo> lld::macho::thunkMap;
 bool TextOutputSection::needsThunks() const {
   if (!target->usesThunks())
     return false;
+  // FIXME: It is not enough to just estimate the size of this section. We
+  // should compute parent->needsThunks by estimating the size of all __text
+  // sections. See https://github.com/llvm/llvm-project/issues/195387
   uint64_t isecAddr = addr;
   for (ConcatInputSection *isec : inputs)
     isecAddr = alignToPowerOf2(isecAddr, isec->align) + isec->getSize();
@@ -191,9 +194,47 @@ void TextOutputSection::createThunk(const ConcatInputSection &isec,
   thunks.push_back(thunkInfo.isec);
 }
 
+std::optional<uint64_t>
+TextOutputSection::estimateStubsEndVA(unsigned numPotentialThunks) const {
+  if (!parent)
+    return std::nullopt;
+
+  auto sections =
+      ArrayRef(parent->getSections())
+          .drop_until([&](const OutputSection *osec) { return osec == this; });
+
+  // Walk backwards to find the last stubs section
+  while (!sections.empty()) {
+    auto *osec = sections.back();
+    if (osec->isNeeded() && (osec == in.stubs || osec == in.objcStubs))
+      break;
+    sections.consume_back();
+  }
+  if (sections.empty())
+    return std::nullopt;
+
+  assert(inputs.empty() || inputs.back()->isFinal);
+  uint64_t estimatedStubsEnd =
+      addr + size + numPotentialThunks * target->thunkSize;
+  for (auto *osec : sections) {
+    if (osec == this)
+      continue;
+    if (!osec->isNeeded())
+      continue;
+    // Check if we will emit any more sections before the last stubs section
+    if (osec != in.stubs && osec != in.stubHelper && osec != in.objcStubs)
+      return std::nullopt;
+    estimatedStubsEnd =
+        alignToPowerOf2(estimatedStubsEnd, osec->align) + osec->getSize();
+  }
+  return estimatedStubsEnd;
+}
+
 bool TextOutputSection::isTargetStubsAndInRange(
     const ConcatInputSection &isec, const Relocation &r,
-    uint64_t estimatedStubsEnd) const {
+    std::optional<uint64_t> estimatedStubsEnd) const {
+  if (!estimatedStubsEnd.has_value())
+    return false;
   auto *funcSym = cast<Symbol *>(r.referent);
   if (!funcSym->isInStubs() && !(in.objcStubs && in.objcStubs->isNeeded() &&
                                  ObjCStubsSection::isObjCStubSymbol(funcSym)))
@@ -201,7 +242,7 @@ bool TextOutputSection::isTargetStubsAndInRange(
   if (r.addend)
     return false;
   uint64_t highVA = isec.getVA() + r.offset + target->forwardBranchRange;
-  return estimatedStubsEnd <= highVA;
+  return *estimatedStubsEnd <= highVA;
 }
 
 void TextOutputSection::finalize() {
@@ -294,15 +335,7 @@ void TextOutputSection::finalize() {
       branchTargets.insert(thunkKey);
   }
 
-  uint64_t estimatedTextEnd =
-      addr + size + branchTargets.size() * target->thunkSize;
-  uint64_t estimatedStubsEnd =
-      alignToPowerOf2(estimatedTextEnd, in.stubs->align) + in.stubs->getSize();
-  if (in.objcStubs && in.objcStubs->isNeeded())
-    estimatedStubsEnd =
-        alignToPowerOf2(estimatedStubsEnd, in.objcStubs->align) +
-        in.objcStubs->getSize();
-
+  auto estimatedStubsEnd = estimateStubsEndVA(branchTargets.size());
   for (auto [isec, r, thunk] : deferredBranchRedirects) {
     if (isTargetKnownInRange(*isec, *r))
       continue;

lld/MachO/ConcatOutputSection.h

@@ -112,12 +112,16 @@ private:
   /// Create a new thunk and update \p r to target the new thunk.
   void createThunk(const ConcatInputSection &isec, Relocation &r,
                    ThunkInfo &thunkInfo);
+  /// \return the largest possible stub section end VA or \p std::nullopt if we
+  /// can't estimate this yet. Used to determine if stub symbol targets are in
+  /// range.
+  std::optional<uint64_t> estimateStubsEndVA(unsigned numPotentialThunks) const;
   /// \return true if the target in \p r is in __stubs or __objc_stubs and in
   /// range from the location in \p isec. \p estimatedStubsEnd is the estimated
   /// VA of the end of the last stubs section.
   bool isTargetStubsAndInRange(const ConcatInputSection &isec,
                                const Relocation &r,
-                               uint64_t estimatedStubsEnd) const;
+                               std::optional<uint64_t> estimatedStubsEnd) const;
   /// The number of relocations updated to point to thunks.
   size_t thunkCallCount = 0;
 };

lld/test/MachO/arm64-thunk-stubs-multi-text.s

@@ -0,0 +1,50 @@
+# REQUIRES: aarch64
+
+# RUN: llvm-mc -filetype=obj -triple=arm64-apple-darwin %s -o %t.o
+# RUN: %lld -arch arm64 -U _extern_sym -o %t %t.o
+# RUN: llvm-objdump --no-print-imm-hex -d --no-show-raw-insn %t | FileCheck %s --implicit-check-not=.thunk.
+
+# CHECK-LABEL: Disassembly of section __TEXT,__text:
+
+# CHECK-LABEL: <_main>:
+# CHECK-NEXT:    bl
+# CHECK-NEXT:    bl 0x[[#%x,THUNK:]] <_extern_sym.thunk.0>
+# CHECK-NEXT:    ret
+
+# CHECK-LABEL: <_foo>:
+# CHECK-NEXT:    bl 0x[[#%x,THUNK:]] <_extern_sym.thunk.0>
+# CHECK-NEXT:    ret
+
+# CHECK: [[#THUNK]] <_extern_sym.thunk.0>:
+
+# CHECK-LABEL: Disassembly of section __TEXT,__lcxx_override:
+# CHECK-LABEL: Disassembly of section __TEXT,__stubs:
+
+.text
+
+.globl _main
+_main:
+  bl _foo
+  bl _extern_sym
+  ret
+
+_spacer0:
+.space 0x4000000-8
+
+.globl _foo
+_foo:
+  bl _extern_sym
+  ret
+
+_spacer1:
+.space 0x4000000
+
+.section __TEXT,__lcxx_override,regular,pure_instructions
+_bar:
+  bl _extern_sym
+  ret
+
+_spacer2:
+.space 0x4000000
+
+.subsections_via_symbols

lld/test/MachO/arm64-thunk-stubs.s

@@ -0,0 +1,75 @@
+# REQUIRES: aarch64
+
+# RUN: llvm-mc -filetype=obj -triple=arm64-apple-darwin %s -o %t.o
+# RUN: %lld -arch arm64 -U _extern_sym -o %t %t.o
+# RUN: llvm-objdump --no-print-imm-hex -d --no-show-raw-insn %t | FileCheck %s --implicit-check-not=.thunk.
+
+# CHECK-LABEL: Disassembly of section __TEXT,__text:
+
+# CHECK-LABEL: <_main>:
+# CHECK-NEXT:    bl
+# CHECK-NEXT:    bl 0x[[#%x,THUNK:]] <_extern_sym.thunk.0>
+# CHECK-NEXT:    bl 0x[[#%x,THUNK_FOO:]] <_objc_msgSend$foo.thunk.0>
+# CHECK-NEXT:    bl 0x[[#%x,THUNK_BAR:]] <_objc_msgSend$bar.thunk.0>
+# CHECK-NEXT:    ret
+
+# CHECK-LABEL: <_foo>:
+# CHECK-NEXT:    bl 0x[[#%x,THUNK:]] <_extern_sym.thunk.0>
+# CHECK-NEXT:    bl 0x[[#%x,THUNK_FOO:]] <_objc_msgSend$foo.thunk.0>
+# CHECK-NEXT:    bl 0x[[#%x,THUNK_BAR:]] <_objc_msgSend$bar.thunk.0>
+# CHECK-NEXT:    ret
+
+# CHECK: [[#THUNK]] <_extern_sym.thunk.0>:
+# CHECK: [[#THUNK_FOO]] <_objc_msgSend$foo.thunk.0>:
+# CHECK: [[#THUNK_BAR]] <_objc_msgSend$bar.thunk.0>:
+
+# CHECK-LABEL: Disassembly of section __TEXT,__stubs:
+# CHECK-LABEL: Disassembly of section __TEXT,__objc_stubs:
+
+.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 _objc_msgSend
+_objc_msgSend:
+  ret
+
+.text
+
+.globl _main
+_main:
+  bl _foo
+  bl _extern_sym
+  bl _objc_msgSend$foo
+  bl _objc_msgSend$bar
+  ret
+
+_spacer0:
+.space 0x4000000-8
+
+.globl _foo
+_foo:
+  bl _extern_sym
+  bl _objc_msgSend$foo
+  bl _objc_msgSend$bar
+  ret
+
+_spacer1:
+.space 0x8000000
+
+.globl _goo
+_goo:
+  bl _extern_sym
+  bl _objc_msgSend$foo
+  bl _objc_msgSend$bar
+  ret
+
+.subsections_via_symbols

lld/test/MachO/arm64-thunks.s

@@ -27,7 +27,6 @@
 
 # OBJC: Sections:
 # OBJC: __text
-# OBJC-NEXT: __lcxx_override
 # OBJC-NEXT: __stubs
 # OBJC-NEXT: __stub_helper
 # OBJC-NEXT: __objc_stubs
@@ -64,7 +63,6 @@
 # MAP-NEXT: 0x{{[[:xdigit:]]+}} {{.*}} _e.thunk.1
 # MAP-NEXT: 0x{{[[:xdigit:]]+}} {{.*}} _f.thunk.1
 # MAP-NEXT: 0x{{[[:xdigit:]]+}} {{.*}} _fold_func_low_addr.thunk.0
-# MAP-NEXT: 0x{{[[:xdigit:]]+}} {{.*}} _z
 
 
 # CHECK: Disassembly of section __TEXT,__text:
@@ -229,10 +227,6 @@
 # CHECK:  adrp x16, 0x[[#%x, FOLD_LOW_PAGE]]000
 # CHECK:  add  x16, x16, #[[#%d, FOLD_LOW_OFFSET]]
 
-# CHECK: Disassembly of section __TEXT,__lcxx_override:
-# CHECK: <_z>:
-# CHECK:  bl 0x[[#%x, A_THUNK_0]] <_a.thunk.0>
-
 # CHECK: Disassembly of section __TEXT,__stubs:
 
 # CHECK: [[#%x, NAN_PAGE + NAN_OFFSET]] <__stubs>:
@@ -423,15 +417,3 @@ _fold_func_high_addr:
   # dramatic memory usage and a huge linker map file
   .space 0x4000000, 'A'
   .byte 0
-
-
-.section __TEXT,__lcxx_override,regular,pure_instructions
-
-.globl _z
-.no_dead_strip _z
-.p2align 2
-_z:
-  bl _a
-  ## Ensure calling into stubs works
-  bl _extern_sym
-  ret