[lld-macho] Move adding bindings for stub targets out of Writer (NFC)

[BertalanD]

Summary

We now re-use the existing needsBinding() helper to determine if a
branch has to go through a stub. The logic for determining which type of
binding is needed can now be done inside StubsSection::addEntry().

This commit prepares us for adding chained fixups support, which will
complicate the logic for adding binding entries.

[int3]

Would be nice to review this together with the stacked chained fixups diff, for context :)

lld/MachO/SyntheticSections.cpp

680

I think we can add an else { llvm_unreachable(...) } here

[thakis]

In D132476#3743012, @int3 wrote:

Would be nice to review this together with the stacked chained fixups diff, for context :)

See here: https://github.com/BertalanD/llvm-project/commits/chained-fixups

[int3]

Kinda wish we didn't have to repeat the checks for symbol type / symbol properties in prepareBranchTarget and the stubs addEntry method, but I can't see a way around it

lld/MachO/Writer.cpp

577–578

how do you feel about replacing this with a bool branchTargetNeedsStubs() method, then have its callers call stubs->addEntry themselves? IMO it makes it clearer that the logic in this function is now solely about deciding whether to do something, rather than how it should be done

[int3]

lld/MachO/Writer.cpp

1133–1135

if we go the branchTargetNeedsStubs() route, I would remove the llvm_unreachable() assert in it and then we wouldn't have to check for isa<Undefined> here

[BertalanD]

lld/MachO/Writer.cpp

577–578

That's a good idea, I agree that it would make the logic clearer.

Actually, it looks like the logic of branchTargetNeedsStubs() would be identical to needsBinding(), so we can do away with this function altogether.

[BertalanD]

Call in.stubs->addEntry() directly, and use needsBinding() to check if a stub is needed.

[int3]

Oh yeah I forgot about needsBinding(). Thanks!

[BertalanD]

Updated StubsSection::addEntry() comment.

[thakis]

Nice!

lld/MachO/SyntheticSections.h

public:
   StubsSection();
   uint64_t getSize() const override;
   bool isNeeded() const override { return !entries.empty(); }
   void finalize() override;
   void writeTo(uint8_t *buf) const override;
   const llvm::SetVector<Symbol *> &getEntries() const { return entries; }
   // Returns whether the symbol was added. Note that every stubs entry will
   // have a corresponding entry in the LazyPointerSection.
-  bool addEntry(Symbol *);
+  void addEntry(Symbol *);
   uint64_t getVA(uint32_t stubsIndex) const {
     assert(isFinal || target->usesThunks());
     // ConcatOutputSection::finalize() can seek the address of a
     // stub before its address is assigned. Before __stubs is
     // finalized, return a contrived out-of-range address.
     return isFinal ? addr + stubsIndex * target->stubSize
                    : TargetInfo::outOfRangeVA;
   }

lld/MachO/SyntheticSections.cpp

void StubsSection::writeTo(uint8_t *buf) const {
   for (const Symbol *sym : entries) {
     target->writeStub(buf + off, *sym);
     off += target->stubSize;
   }
 }
 
 void StubsSection::finalize() { isFinal = true; }
 
-bool StubsSection::addEntry(Symbol *sym) {
+static void addBindingsForStub(Symbol *sym) {
+  if (auto *dysym = dyn_cast<DylibSymbol>(sym)) {
+    if (sym->isWeakDef()) {
+      in.binding->addEntry(dysym, in.lazyPointers->isec,
+                           sym->stubsIndex * target->wordSize);
+      in.weakBinding->addEntry(sym, in.lazyPointers->isec,
+                               sym->stubsIndex * target->wordSize);
+    } else {
+      in.lazyBinding->addEntry(dysym);
+    }
+  } else if (auto *defined = dyn_cast<Defined>(sym)) {
+    if (defined->isExternalWeakDef()) {
+      in.rebase->addEntry(in.lazyPointers->isec,
+                          sym->stubsIndex * target->wordSize);
+      in.weakBinding->addEntry(sym, in.lazyPointers->isec,
+                               sym->stubsIndex * target->wordSize);
+    } else if (defined->interposable) {
+      in.lazyBinding->addEntry(sym);
+    } else {
+      llvm_unreachable("invalid stub target");
+    }
+  } else {
+    llvm_unreachable("invalid stub target symbol type");
+  }
+}
+
+void StubsSection::addEntry(Symbol *sym) {
   bool inserted = entries.insert(sym);
-  if (inserted)
+  if (inserted) {
     sym->stubsIndex = entries.size() - 1;
-  return inserted;
+    addBindingsForStub(sym);
+  }
 }
 
 StubHelperSection::StubHelperSection()
     : SyntheticSection(segment_names::text, section_names::stubHelper) {
   flags = S_ATTR_SOME_INSTRUCTIONS | S_ATTR_PURE_INSTRUCTIONS;
   align = 4; // This section comprises machine instructions
 }
 

lld/MachO/Writer.cpp

void Writer::treatSpecialUndefineds() {
   // patterns need not match.
   for (const CachedHashStringRef &cachedName :
        config->exportedSymbols.literals) {
     if (const Symbol *sym = symtab->find(cachedName))
       if (const auto *undefined = dyn_cast<Undefined>(sym))
         treatUndefinedSymbol(*undefined, "-exported_symbol(s_list)");
   }
 }
 
-// Add stubs and bindings where necessary (e.g. if the symbol is a
-// DylibSymbol.)
-static void prepareBranchTarget(Symbol *sym) {
-  if (auto *dysym = dyn_cast<DylibSymbol>(sym)) {
-    if (in.stubs->addEntry(dysym)) {
-      if (sym->isWeakDef()) {
-        in.binding->addEntry(dysym, in.lazyPointers->isec,
-                             sym->stubsIndex * target->wordSize);
-        in.weakBinding->addEntry(sym, in.lazyPointers->isec,
-                                 sym->stubsIndex * target->wordSize);
-      } else {
-        in.lazyBinding->addEntry(dysym);
-      }
-    }
-  } else if (auto *defined = dyn_cast<Defined>(sym)) {
-    if (defined->isExternalWeakDef()) {
-      if (in.stubs->addEntry(sym)) {
-        in.rebase->addEntry(in.lazyPointers->isec,
-                            sym->stubsIndex * target->wordSize);
-        in.weakBinding->addEntry(sym, in.lazyPointers->isec,
-                                 sym->stubsIndex * target->wordSize);
-      }
-    } else if (defined->interposable) {
-      if (in.stubs->addEntry(sym))
-        in.lazyBinding->addEntry(sym);
-    }
-  } else {
-    llvm_unreachable("invalid branch target symbol type");
-  }
-}
-
 // Can a symbol's address can only be resolved at runtime?
 static bool needsBinding(const Symbol *sym) {
   if (isa<DylibSymbol>(sym))
     return true;
   if (const auto *defined = dyn_cast<Defined>(sym))
     return defined->isExternalWeakDef() || defined->interposable;
   return false;
 }
 
 static void prepareSymbolRelocation(Symbol *sym, const InputSection *isec,
                                     const lld::macho::Reloc &r) {
   assert(sym->isLive());
   const RelocAttrs &relocAttrs = target->getRelocAttrs(r.type);
 
   if (relocAttrs.hasAttr(RelocAttrBits::BRANCH)) {
-    prepareBranchTarget(sym);
+    if (needsBinding(sym))
+      in.stubs->addEntry(sym);
   } else if (relocAttrs.hasAttr(RelocAttrBits::GOT)) {
     if (relocAttrs.hasAttr(RelocAttrBits::POINTER) || needsBinding(sym))
       in.got->addEntry(sym);
   } else if (relocAttrs.hasAttr(RelocAttrBits::TLV)) {
     if (needsBinding(sym))
       in.tlvPointers->addEntry(sym);
   } else if (relocAttrs.hasAttr(RelocAttrBits::UNSIGNED)) {
     // References from thread-local variable sections are treated as offsets
void Writer::writeOutputFile() {
   writeUuid();
   writeCodeSignature();
 
   if (auto e = buffer->commit())
     error("failed to write to the output file: " + toString(std::move(e)));
 }
 
 template <class LP> void Writer::run() {
   treatSpecialUndefineds();
-  if (config->entry && !isa<Undefined>(config->entry))
-    prepareBranchTarget(config->entry);
+  if (config->entry && needsBinding(config->entry))
+    in.stubs->addEntry(config->entry);
 
   // Canonicalization of all pointers to InputSections should be handled by
   // these two scan* methods. I.e. from this point onward, for all live
   // InputSections, we should have `isec->canonical() == isec`.
   scanSymbols();
   if (in.objcStubs->isNeeded())
     in.objcStubs->setup();
   scanRelocations();