[lld-macho][nfc] Give non-text ConcatOutputSections order-independent finalization

[int3]

Summary

This diff is motivated by my work to add proper DWARF unwind support. As
detailed in PR50956 functions that need DWARF unwind need to have
compact unwind entries synthesized for them. These CU entries encode an
offset within __eh_frame that points to the corresponding DWARF FDE.

In order to encode this offset during
UnwindInfoSectionImpl::finalize(), we need to first assign values to
InputSection::outSecOff for each __eh_frame subsection. But
__eh_frame is ordered after __unwind_info (according to ld64 at
least), which puts us in a bit of a bind: outSecOff gets assigned
during finalization, but __eh_frame is being finalized after
__unwind_info.

But it occurred to me that there's no real need for most
ConcatOutputSections to be finalized sequentially. It's only necessary
for text-containing ConcatOutputSections that may contain branch relocs
which may need thunks. ConcatOutputSections containing other types of
data can be finalized in any order.

This diff moves the finalization logic for non-text sections into a
separate finalizeContents() method. This method is called before
section address assignment & unwind info finalization takes place. In
theory we could call these finalizeContents() methods in parallel, but
in practice it seems to be faster to do it all on the main thread.

[oontvoo]

lld/MachO/ConcatOutputSection.cpp

198–208

Why is the lambda needed here? (It's only got one caller ...)

Why not make it the body of the for loop?

410–418

nit: it seems a bit more intuitive to reverse the condition here. ie., "if segment == text", then make text-output-section, otherwise, concat-output-section.

lld/MachO/ConcatOutputSection.h

70–71

could we have some comments here on finalizae() vs finalizeContents() pls?

lld/MachO/Writer.cpp

976

less memory?
(citations needed 😄 )

[oontvoo]

thx!

lld/MachO/ConcatOutputSection.h

55

perhaps make this a private, too?

[int3]

lld/MachO/ConcatOutputSection.cpp

198–208

because I forgot to clean up my copypasta :p good catch

410–418

agreed

lld/MachO/Writer.cpp

976

oh, chromium_framework actually links a bit faster on my machine with the non-parallel setup. I'll make the comment more explicit :)

[int3]

lld/MachO/ConcatOutputSection.cpp

198–208

factoring it out into its own method so that TextOutputSection can share it too

lld/MachO/ConcatOutputSection.h

55

This does get used externally (in offsetToInputSection as well as sortSegmentsAndSections).

70–71

OutputSections::finalize() already has a comment block that talks about finalizeContents(). I'll add a comment here to reference it

[int3]

address comments

lld/MachO/ConcatOutputSection.h

 //===- ConcatOutputSection.h ------------------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
-#ifndef LLD_MACHO_MERGED_OUTPUT_SECTION_H
-#define LLD_MACHO_MERGED_OUTPUT_SECTION_H
+#ifndef LLD_MACHO_CONCAT_OUTPUT_SECTION_H
+#define LLD_MACHO_CONCAT_OUTPUT_SECTION_H
 
 #include "InputSection.h"
 #include "OutputSection.h"
 #include "lld/Common/LLVM.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/MapVector.h"
 
 namespace lld {
 namespace macho {
 
 class Defined;
 
 // Linking multiple files will inevitably mean resolving sections in different
 // files that are labeled with the same segment and section name. This class
 // contains all such sections and writes the data from each section sequentially
 // in the final binary.
-class ConcatOutputSection final : public OutputSection {
+class ConcatOutputSection : public OutputSection {
 public:
   explicit ConcatOutputSection(StringRef name)
       : OutputSection(ConcatKind, name) {}
 
   const ConcatInputSection *firstSection() const { return inputs.front(); }
   const ConcatInputSection *lastSection() const { return inputs.back(); }
   bool isNeeded() const override { return !inputs.empty(); }
 
   // These accessors will only be valid after finalizing the section
   uint64_t getSize() const override { return size; }
   uint64_t getFileSize() const override { return fileSize; }
 
-  void addInput(ConcatInputSection *input);
-  void finalize() override;
-  bool needsThunks() const;
-  uint64_t estimateStubsInRangeVA(size_t callIdx) const;
+  // Assign values to InputSection::outSecOff. In contrast to TextOutputSection,
+  // which does this in its implementation of `finalize()`, we can do this
+  // without `finalize()`'s sequential guarantees detailed in the block comment
+  // of `OutputSection::finalize()`.
+  virtual void finalizeContents();
 
+  void addInput(ConcatInputSection *input);
   void writeTo(uint8_t *buf) const override;
 
-  std::vector<ConcatInputSection *> inputs;
-  std::vector<ConcatInputSection *> thunks;
-
   static bool classof(const OutputSection *sec) {
     return sec->kind() == ConcatKind;
   }
 
   static ConcatOutputSection *getOrCreateForInput(const InputSection *);
 
-private:
-  void finalizeFlags(InputSection *input);
+  std::vector<ConcatInputSection *> inputs;
 
+protected:
   size_t size = 0;
   uint64_t fileSize = 0;
+  void finalizeOne(ConcatInputSection *);
+
+private:
+  void finalizeFlags(InputSection *input);
+};
+
+// ConcatOutputSections that contain code (text) require special handling to
+// support thunk insertion.
+class TextOutputSection : public ConcatOutputSection {
+public:
+  explicit TextOutputSection(StringRef name) : ConcatOutputSection(name) {}
+  void finalizeContents() override {}
+  void finalize() override;
+  bool needsThunks() const;
+  void writeTo(uint8_t *buf) const override;
+
+private:
+  uint64_t estimateStubsInRangeVA(size_t callIdx) const;
+
+  std::vector<ConcatInputSection *> thunks;
 };
 
 // We maintain one ThunkInfo per real function.
 //
 // The "active thunk" is represented by the sym/isec pair that
 // turns-over during finalize(): as the call-site address advances,
 // the active thunk goes out of branch-range, and we create a new
 // thunk to take its place.

lld/MachO/ConcatOutputSection.cpp

 
 DenseMap<Symbol *, ThunkInfo> lld::macho::thunkMap;
 
 // Determine whether we need thunks, which depends on the target arch -- RISC
 // (i.e., ARM) generally does because it has limited-range branch/call
 // instructions, whereas CISC (i.e., x86) generally doesn't. RISC only needs
 // thunks for programs so large that branch source & destination addresses
 // might differ more than the range of branch instruction(s).
-bool ConcatOutputSection::needsThunks() const {
+bool TextOutputSection::needsThunks() const {
   if (!target->usesThunks())
     return false;
   uint64_t isecAddr = addr;
   for (ConcatInputSection *isec : inputs)
     isecAddr = alignTo(isecAddr, isec->align) + isec->getSize();
   if (isecAddr - addr + in.stubs->getSize() <=
       std::min(target->backwardBranchRange, target->forwardBranchRange))
     return false;
   // Yes, this program is large enough to need thunks.
   for (ConcatInputSection *isec : inputs) {
     for (Reloc &r : isec->relocs) {
       if (!target->hasAttr(r.type, RelocAttrBits::BRANCH))
         continue;
       auto *sym = r.referent.get<Symbol *>();
       // Pre-populate the thunkMap and memoize call site counts for every
       // InputSection and ThunkInfo. We do this for the benefit of
-      // ConcatOutputSection::estimateStubsInRangeVA()
+      // estimateStubsInRangeVA().
       ThunkInfo &thunkInfo = thunkMap[sym];
       // Knowing ThunkInfo call site count will help us know whether or not we
       // might need to create more for this referent at the time we are
       // estimating distance to __stubs in estimateStubsInRangeVA().
       ++thunkInfo.callSiteCount;
       // We can avoid work on InputSections that have no BRANCH relocs.
       isec->hasCallSites = true;
     }
   }
   return true;
 }
 
 // Since __stubs is placed after __text, we must estimate the address
 // beyond which stubs are within range of a simple forward branch.
 // This is called exactly once, when the last input section has been finalized.
-uint64_t ConcatOutputSection::estimateStubsInRangeVA(size_t callIdx) const {
+uint64_t TextOutputSection::estimateStubsInRangeVA(size_t callIdx) const {
   // Tally the functions which still have call sites remaining to process,
   // which yields the maximum number of thunks we might yet place.
   size_t maxPotentialThunks = 0;
   for (auto &tp : thunkMap) {
     ThunkInfo &ti = tp.second;
     // This overcounts: Only sections that are in forward jump range from the
     // currently-active section get finalized, and all input sections are
     // finalized when estimateStubsInRangeVA() is called. So only backward
log("thunks = " + std::to_string(thunkMap.size()) +
       ", stubs = " + std::to_string(in.stubs->getSize()) + ", isecVA = " +
       to_hexString(isecVA) + ", threshold = " + to_hexString(stubsInRangeVA) +
       ", isecEnd = " + to_hexString(isecEnd) +
       ", tail = " + to_hexString(isecEnd - isecVA) +
       ", slop = " + to_hexString(forwardBranchRange - (isecEnd - isecVA)));
   return stubsInRangeVA;
 }
 
-void ConcatOutputSection::finalize() {
-  uint64_t isecAddr = addr;
-  uint64_t isecFileOff = fileOff;
-  auto finalizeOne = [&](ConcatInputSection *isec) {
-    isecAddr = alignTo(isecAddr, isec->align);
-    isecFileOff = alignTo(isecFileOff, isec->align);
-    isec->outSecOff = isecAddr - addr;
-    isec->isFinal = true;
-    isecAddr += isec->getSize();
-    isecFileOff += isec->getFileSize();
-  };
+void ConcatOutputSection::finalizeOne(ConcatInputSection *isec) {
+  size = alignTo(size, isec->align);
+  fileSize = alignTo(fileSize, isec->align);
+  isec->outSecOff = size;
+  size += isec->getSize();
+  fileSize += isec->getFileSize();
+}
+
+void ConcatOutputSection::finalizeContents() {
+  for (ConcatInputSection *isec : inputs)
+    finalizeOne(isec);
+}
 
+void TextOutputSection::finalize() {
   if (!needsThunks()) {
     for (ConcatInputSection *isec : inputs)
       finalizeOne(isec);
-    size = isecAddr - addr;
-    fileSize = isecFileOff - fileOff;
     return;
   }
 
   uint64_t forwardBranchRange = target->forwardBranchRange;
   uint64_t backwardBranchRange = target->backwardBranchRange;
   uint64_t stubsInRangeVA = 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.
-  // isecAddr is the start address of the first non-finalized 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 = 256 * thunkSize;
-    while (finalIdx < endIdx && isecAddr + inputs[finalIdx]->getSize() <
+    while (finalIdx < endIdx && addr + size + inputs[finalIdx]->getSize() <
                                     isecVA + forwardBranchRange - slop)
       finalizeOne(inputs[finalIdx++]);
 
     if (!isec->hasCallSites)
       continue;
 
     if (finalIdx == endIdx && stubsInRangeVA == TargetInfo::outOfRangeVA) {
       // When we have finalized all input sections, __stubs (destined
for (Reloc &r : reverse(relocs)) {
       if (thunkInfo.sym) {
         uint64_t thunkVA = thunkInfo.isec->getVA();
         if (lowVA <= thunkVA && thunkVA <= highVA) {
           r.referent = thunkInfo.sym;
           continue;
         }
       }
       // ... otherwise, create a new thunk.
-      if (isecAddr > highVA) {
+      if (addr + size > highVA) {
         // There were too many consecutive branch instructions for `slop`
         // above. 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");
       }
       thunkInfo.isec =
           makeSyntheticInputSection(isec->getSegName(), isec->getName());
for (Reloc &r : reverse(relocs)) {
             thunkSize, /*isWeakDef=*/false, /*isExternal=*/false,
             /*isPrivateExtern=*/true, /*isThumb=*/false,
             /*isReferencedDynamically=*/false, /*noDeadStrip=*/false,
             /*isWeakDefCanBeHidden=*/false);
       }
       thunkInfo.sym->used = true;
       target->populateThunk(thunkInfo.isec, funcSym);
       finalizeOne(thunkInfo.isec);
+      fprintf(stderr, "%llx\n", thunkInfo.isec->outSecOff);
       thunks.push_back(thunkInfo.isec);
       ++thunkCount;
     }
   }
-  size = isecAddr - addr;
-  fileSize = isecFileOff - fileOff;
 
   log("thunks for " + parent->name + "," + name +
       ": funcs = " + std::to_string(thunkMap.size()) +
       ", relocs = " + std::to_string(relocCount) +
       ", all calls = " + std::to_string(callSiteCount) +
       ", thunk calls = " + std::to_string(thunkCallCount) +
       ", thunks = " + std::to_string(thunkCount));
 }
 
 void ConcatOutputSection::writeTo(uint8_t *buf) const {
+  for (ConcatInputSection *isec : inputs)
+    isec->writeTo(buf + isec->outSecOff);
+}
+
+void TextOutputSection::writeTo(uint8_t *buf) const {
   // Merge input sections from thunk & ordinary vectors
   size_t i = 0, ie = inputs.size();
   size_t t = 0, te = thunks.size();
   while (i < ie || t < te) {
     while (i < ie && (t == te || inputs[i]->empty() ||
                       inputs[i]->outSecOff < thunks[t]->outSecOff)) {
       inputs[i]->writeTo(buf + inputs[i]->outSecOff);
       ++i;
case S_SYMBOL_STUBS:
     break;
   }
 }
 
 ConcatOutputSection *
 ConcatOutputSection::getOrCreateForInput(const InputSection *isec) {
   NamePair names = maybeRenameSection({isec->getSegName(), isec->getName()});
   ConcatOutputSection *&osec = concatOutputSections[names];
-  if (!osec)
+  if (!osec) {
+    if (isec->getSegName() == segment_names::text &&
+        isec->getName() != section_names::gccExceptTab &&
+        isec->getName() != section_names::ehFrame)
+      osec = make<TextOutputSection>(names.second);
+    else
       osec = make<ConcatOutputSection>(names.second);
+  }
   return osec;
 }
 
 NamePair macho::maybeRenameSection(NamePair key) {
   auto newNames = config->sectionRenameMap.find(key);
   if (newNames != config->sectionRenameMap.end())
     return newNames->second;
   return key;
 }

lld/MachO/Writer.cpp

template <class LP> void Writer::createOutputSections() {
 
   // dyld requires __LINKEDIT segment to always exist (even if empty).
   linkEditSegment = getOrCreateOutputSegment(segment_names::linkEdit);
 }
 
 void Writer::finalizeAddresses() {
   TimeTraceScope timeScope("Finalize addresses");
   uint64_t pageSize = target->getPageSize();
+
+  // We could parallelize this loop, but local benchmarking indicates it is
+  // faster to do it all in the main thread.
+  for (OutputSegment *seg : outputSegments) {
+    if (seg == linkEditSegment)
+      continue;
+    for (OutputSection *osec : seg->getSections()) {
+      if (!osec->isNeeded())
+        continue;
+      // Other kinds of OutputSections have already been finalized.
+      if (auto concatOsec = dyn_cast<ConcatOutputSection>(osec))
+          concatOsec->finalizeContents();
+    }
+  }
+
   // Ensure that segments (and the sections they contain) are allocated
   // addresses in ascending order, which dyld requires.
   //
   // Note that at this point, __LINKEDIT sections are empty, but we need to
   // determine addresses of other segments/sections before generating its
   // contents.
   for (OutputSegment *seg : outputSegments) {
     if (seg == linkEditSegment)