[lld] [lld][macho] Restructure thunk generation algorithm (PR #193367)

[Zhaoxuan Jiang via llvm-commits]

================
@@ -276,171 +115,196 @@ void TextOutputSection::finalize() {
 
   uint64_t forwardBranchRange = target->forwardBranchRange;
   uint64_t backwardBranchRange = target->backwardBranchRange;
-  uint64_t branchTargetThresholdVA = TargetInfo::outOfRangeVA;
   size_t thunkSize = target->thunkSize;
-  size_t relocCount = 0;
-  size_t callSiteCount = 0;
   size_t thunkCallCount = 0;
   size_t thunkCount = 0;
 
-  // Walk all sections in order. Finalize all sections that are less than
-  // forwardBranchRange in front of it.
-  // isecVA is the address of the current section.
-  // addr + size is the start address of the first non-finalized section.
-
-  // inputs[finalIdx] is for finalization (address-assignment)
-  size_t finalIdx = 0;
-  // Kick-off by ensuring that the first input section has an address
-  for (size_t callIdx = 0, endIdx = inputs.size(); callIdx < endIdx;
-       ++callIdx) {
-    if (finalIdx == callIdx)
-      finalizeOne(inputs[finalIdx++]);
-    ConcatInputSection *isec = inputs[callIdx];
-    assert(isec->isFinal);
-    uint64_t isecVA = isec->getVA();
-
-    // Assign addresses up-to the forward branch-range limit.
-    // Every call instruction needs a small number of bytes (on Arm64: 4),
-    // and each inserted thunk needs a slightly larger number of bytes
-    // (on Arm64: 12). If a section starts with a branch instruction and
-    // contains several branch instructions in succession, then the distance
-    // from the current position to the position where the thunks are inserted
-    // grows. So leave room for a bunch of thunks.
-    unsigned slop = config->slopScale * thunkSize;
-    while (finalIdx < endIdx) {
-      uint64_t expectedNewSize =
-          alignToPowerOf2(addr + size, inputs[finalIdx]->align) +
-          inputs[finalIdx]->getSize();
-      if (expectedNewSize >= isecVA + forwardBranchRange - slop)
+  auto isTargetInRange = [&](const ConcatInputSection &isec,
+                             const Relocation &r) -> bool {
+    uint64_t callVA = isec.getVA() + r.offset;
+    uint64_t lowVA =
+        backwardBranchRange < callVA ? callVA - backwardBranchRange : 0;
+    uint64_t highVA = callVA + forwardBranchRange;
+    auto *funcSym = cast<Symbol *>(r.referent);
+    uint64_t funcVA = resolveSymbolOffsetVA(funcSym, r.type, r.addend);
+    // Check if the referent is reachable with a simple call instruction.
+    return lowVA <= funcVA && funcVA <= highVA;
+  };
+  auto getThunkInRange = [&](const ConcatInputSection &isec,
+                             const Relocation &r) -> Defined * {
+    assert(!isTargetInRange(isec, r));
+    uint64_t callVA = isec.getVA() + r.offset;
+    uint64_t lowVA =
+        backwardBranchRange < callVA ? callVA - backwardBranchRange : 0;
+    uint64_t highVA = callVA + forwardBranchRange;
+    auto *funcSym = cast<Symbol *>(r.referent);
+    ThunkInfo &thunkInfo = thunkMap[ThunkKey{funcSym, r.addend}];
+    if (thunkInfo.sym) {
+      uint64_t thunkVA = thunkInfo.isec->getVA();
+      if (lowVA <= thunkVA && thunkVA <= highVA)
+        return thunkInfo.sym;
+    }
+    return nullptr;
+  };
+  auto createThunk = [&](const ConcatInputSection &isec, Relocation &r) {
+    assert(getThunkInRange(isec, r) == nullptr);
+    assert(isec.isFinal);
+    uint64_t highVA = isec.getVA() + r.offset + forwardBranchRange;
+    if (addr + size > highVA) {
+      // There were too many consecutive branch instructions for `slop`
+      // below. If you hit this: For the current algorithm, just bumping up
+      // slop above and trying again is probably simplest. (See also PR51578
+      // comment 5).
+      fatal(Twine(__FUNCTION__) +
+            ": FIXME: thunk range overrun. Consider increasing the "
+            "slop-scale with `--slop-scale=<unsigned_int>`.");
+    }
+    auto *funcSym = cast<Symbol *>(r.referent);
+    ThunkInfo &thunkInfo = thunkMap[ThunkKey{funcSym, r.addend}];
+    thunkInfo.isec =
+        makeSyntheticInputSection(isec.getSegName(), isec.getName());
+    thunkInfo.isec->parent = this;
+    assert(thunkInfo.isec->live);
+
+    std::string addendSuffix;
+    if (r.addend != 0)
+      addendSuffix = "+" + std::to_string(r.addend);
+    StringRef thunkName =
+        saver().save(funcSym->getName() + addendSuffix + ".thunk." +
+                     std::to_string(thunkInfo.sequence++));
+    if (!isa<Defined>(funcSym) || cast<Defined>(funcSym)->isExternal()) {
+      r.referent = thunkInfo.sym = symtab->addDefined(
+          thunkName, /*file=*/nullptr, thunkInfo.isec, /*value=*/0, thunkSize,
+          /*isWeakDef=*/false, /*isPrivateExtern=*/true,
+          /*isReferencedDynamically=*/false, /*noDeadStrip=*/false,
+          /*isWeakDefCanBeHidden=*/false);
+    } else {
+      r.referent = thunkInfo.sym = make<Defined>(
+          thunkName, /*file=*/nullptr, thunkInfo.isec, /*value=*/0, thunkSize,
+          /*isWeakDef=*/false, /*isExternal=*/false, /*isPrivateExtern=*/true,
+          /*includeInSymtab=*/true, /*isReferencedDynamically=*/false,
+          /*noDeadStrip=*/false, /*isWeakDefCanBeHidden=*/false);
+    }
+    thunkInfo.sym->used = true;
+    target->populateThunk(thunkInfo.isec, funcSym, r.addend);
+    // The thunk itself bakes in the addend, so the call-site reloc must
+    // branch to the thunk start with no extra offset.
+    r.addend = 0;
+    ++thunkCallCount;
+    finalizeOne(thunkInfo.isec);
+    thunks.push_back(thunkInfo.isec);
+    ++thunkCount;
+  };
+
+  std::deque<std::tuple<ConcatInputSection *, Relocation *, ThunkKey>>
+      branchesToProcess;
+  SmallVector<std::tuple<ConcatInputSection *, Relocation *, Defined *>>
+      deferredBranchRedirects;
+
+  const uint64_t slop = config->slopScale * thunkSize;
+  for (auto *isec : inputs) {
+    while (!branchesToProcess.empty()) {
+      auto &[callerIsec, r, thunkKey] = branchesToProcess.front();
+      assert(callerIsec->isFinal);
+      if (isTargetInRange(*callerIsec, *r)) {
+        branchesToProcess.pop_front();
+        continue;
+      }
+      if (auto *sym = getThunkInRange(*callerIsec, *r)) {
+        deferredBranchRedirects.emplace_back(callerIsec, r, sym);
+        branchesToProcess.pop_front();
+        continue;
+      }
+      uint64_t highVA = callerIsec->getVA() + r->offset + forwardBranchRange;
+      uint64_t nextEnd =
+          alignToPowerOf2(addr + size, isec->align) + isec->getSize();
+      if (nextEnd + slop <= highVA)
         break;
-      finalizeOne(inputs[finalIdx++]);
+
+      createThunk(*callerIsec, *r);
+      branchesToProcess.pop_front();
     }
+    finalizeOne(isec);
 
-    if (!isec->hasCallSites)
+    bool hasCallsite = llvm::any_of(isec->relocs, [](Relocation &r) {
+      return target->hasAttr(r.type, RelocAttrBits::BRANCH);
+    });
+    if (!hasCallsite)
----------------
nocchijiang wrote:

I don't think @drodriguez was asking to remove the assert in this PR. I agree we should not try to change the assert semantics in this PR, but I do want this assert to go away eventually. We have `ld64 -r` inputs whose reloc order does not satisfy LLD’s current expectation, and our local workaround is quite ugly: we build a pointer copy of the branch relocs and sort that copy just to satisfy the assumption here.

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