LLD patch: Fix the ASM code generated for __stub_helpers section

This applies 93ca847862af07632197dcf2d8a68b9b27a26d7a
from the llvm-project git monorepo to the embedded LLD.

deps/lld/lib/ReaderWriter/MachO/ArchHandler.h

@@ -112,6 +112,10 @@ public:
   /// info in final executables.
   virtual bool isLazyPointer(const Reference &);
 
+  /// Reference from an __stub_helper entry to the required offset of the
+  /// lazy bind commands.
+  virtual Reference::KindValue lazyImmediateLocationKind() = 0;
+
   /// Returns true if the specified relocation is paired to the next relocation.
   virtual bool isPairedReloc(const normalized::Relocation &) = 0;
 

deps/lld/lib/ReaderWriter/MachO/ArchHandler_arm.cpp

@@ -67,6 +67,10 @@ public:
     return invalid;
   }
 
+  Reference::KindValue lazyImmediateLocationKind() override {
+    return lazyImmediateLocation;
+  }
+
   Reference::KindValue pointerKind() override {
     return invalid;
   }

deps/lld/lib/ReaderWriter/MachO/ArchHandler_arm64.cpp

@@ -127,6 +127,10 @@ public:
     return pointer64;
   }
 
+  Reference::KindValue lazyImmediateLocationKind() override {
+    return lazyImmediateLocation;
+  }
+
   uint32_t dwarfCompactUnwindType() override {
     return 0x03000000;
   }

deps/lld/lib/ReaderWriter/MachO/ArchHandler_x86.cpp

@@ -70,6 +70,10 @@ public:
     return delta32;
   }
 
+  Reference::KindValue lazyImmediateLocationKind() override {
+    return lazyImmediateLocation;
+  }
+
   Reference::KindValue unwindRefToEhFrameKind() override {
     return invalid;
   }

deps/lld/lib/ReaderWriter/MachO/ArchHandler_x86_64.cpp

@@ -116,6 +116,10 @@ public:
     return unwindFDEToFunction;
   }
 
+  Reference::KindValue lazyImmediateLocationKind() override {
+    return lazyImmediateLocation;
+  }
+
   Reference::KindValue unwindRefToEhFrameKind() override {
     return unwindInfoToEhFrame;
   }

deps/lld/lib/ReaderWriter/MachO/MachONormalizedFileFromAtoms.cpp

@@ -172,6 +172,8 @@ private:
                                    SymbolScope &symbolScope);
   void         appendSection(SectionInfo *si, NormalizedFile &file);
   uint32_t     sectionIndexForAtom(const Atom *atom);
+  void fixLazyReferenceImm(const DefinedAtom *atom, uint32_t offset,
+                           NormalizedFile &file);
 
   typedef llvm::DenseMap<const Atom*, uint32_t> AtomToIndex;
   struct AtomAndIndex { const Atom *atom; uint32_t index; SymbolScope scope; };
@@ -1423,6 +1425,8 @@ void Util::addRebaseAndBindingInfo(const lld::File &atomFile,
 
   uint8_t segmentIndex;
   uint64_t segmentStartAddr;
+  uint32_t offsetInBindInfo = 0;
+
   for (SectionInfo *sect : _sectionInfos) {
     segIndexForSection(sect, segmentIndex, segmentStartAddr);
     for (const AtomInfo &info : sect->atomsAndOffsets) {
@@ -1467,6 +1471,59 @@ void Util::addRebaseAndBindingInfo(const lld::File &atomFile,
           bind.symbolName = targ->name();
           bind.addend = ref->addend();
           nFile.lazyBindingInfo.push_back(bind);
+
+          // Now that we know the segmentOffset and the ordinal attribute,
+          // we can fix the helper's code
+
+          fixLazyReferenceImm(atom, offsetInBindInfo, nFile);
+
+          // 5 bytes for opcodes + variable sizes (target name + \0 and offset
+          // encode's size)
+          offsetInBindInfo +=
+              6 + targ->name().size() + llvm::getULEB128Size(bind.segOffset);
+          if (bind.ordinal > BIND_IMMEDIATE_MASK)
+            offsetInBindInfo += llvm::getULEB128Size(bind.ordinal);
+        }
+      }
+    }
+  }
+}
+
+void Util::fixLazyReferenceImm(const DefinedAtom *atom, uint32_t offset,
+                               NormalizedFile &file) {
+  for (const auto &ref : *atom) {
+    const DefinedAtom *da = dyn_cast<DefinedAtom>(ref->target());
+    if (da == nullptr)
+      return;
+
+    const Reference *helperRef = nullptr;
+    for (const Reference *hr : *da) {
+      if (hr->kindValue() == _archHandler.lazyImmediateLocationKind()) {
+        helperRef = hr;
+        break;
+      }
+    }
+    if (helperRef == nullptr)
+      continue;
+
+    // TODO: maybe get the fixed atom content from _archHandler ?
+    for (SectionInfo *sectInfo : _sectionInfos) {
+      for (const AtomInfo &atomInfo : sectInfo->atomsAndOffsets) {
+        if (atomInfo.atom == helperRef->target()) {
+          auto sectionContent =
+              file.sections[sectInfo->normalizedSectionIndex].content;
+          uint8_t *rawb =
+              file.ownedAllocations.Allocate<uint8_t>(sectionContent.size());
+          llvm::MutableArrayRef<uint8_t> newContent{rawb,
+                                                    sectionContent.size()};
+          std::copy(sectionContent.begin(), sectionContent.end(),
+                    newContent.begin());
+          llvm::support::ulittle32_t *loc =
+              reinterpret_cast<llvm::support::ulittle32_t *>(
+                  &newContent[atomInfo.offsetInSection +
+                              helperRef->offsetInAtom()]);
+          *loc = offset;
+          file.sections[sectInfo->normalizedSectionIndex].content = newContent;
         }
       }
     }

deps/lld/test/mach-o/lazy-bind-x86_64.yaml

@@ -80,8 +80,8 @@ undefined-symbols:
 
 # CHECK-HELPERS:Disassembly of section __TEXT,__stub_helper:
 # CHECK-HELPERS: 	68 00 00 00 00            pushq	$0
-# CHECK-HELPERS: 	68 10 00 00 00            pushq	$16
-# CHECK-HELPERS: 	68 20 00 00 00            pushq	$32
+# CHECK-HELPERS: 	68 0b 00 00 00            pushq	$11
+# CHECK-HELPERS: 	68 16 00 00 00            pushq	$22
 
 # Make sure the stub helper is correctly aligned
 # CHECK-DYLIBS:   sectname __stub_helper