[mach-o] Fix lazy binding offsets

The way lazy binding works in mach-o is that the linker generates a helper
function and has the stub (PLT) initially jump to it.  The helper function
pushes an extra parameter then jumps into dyld.  The extra parameter is an
offset into the lazy binding info where dyld will find the information about
which symbol to bind and way lazy binding pointer to update.

llvm-svn: 221654

7 files changed, 114 insertions, 123 deletions

diff --git a/lld/lib/ReaderWriter/MachO/ArchHandler_arm.cpp b/lld/lib/ReaderWriter/MachO/ArchHandler_arm.cpp
index e30dd4dede0..65a40802c54 100644
--- a/lld/lib/ReaderWriter/MachO/ArchHandler_arm.cpp
+++ b/lld/lib/ReaderWriter/MachO/ArchHandler_arm.cpp
@@ -998,9 +998,11 @@ void ArchHandler_arm::applyFixupFinal(const Reference &ref, uint8_t *loc,
       *loc32 = targetAddress - fixupAddress + ref.addend();
     break;
   case lazyPointer:
-  case lazyImmediateLocation:
     // do nothing
     break;
+  case lazyImmediateLocation:
+    *loc32 = ref.addend();
+    break;
   case invalid:
     llvm_unreachable("invalid ARM Reference Kind");
     break;
diff --git a/lld/lib/ReaderWriter/MachO/ArchHandler_arm64.cpp b/lld/lib/ReaderWriter/MachO/ArchHandler_arm64.cpp
index c39e4c1cdaf..c4456a896aa 100644
--- a/lld/lib/ReaderWriter/MachO/ArchHandler_arm64.cpp
+++ b/lld/lib/ReaderWriter/MachO/ArchHandler_arm64.cpp
@@ -573,9 +573,11 @@ void ArchHandler_arm64::applyFixupFinal(const Reference &ref, uint8_t *loc,
     *loc32 = (targetAddress - fixupAddress) + ref.addend();
     return;
   case lazyPointer:
-  case lazyImmediateLocation:
     // Do nothing
     return;
+  case lazyImmediateLocation:
+    *loc32 = ref.addend();
+    return;
   case invalid:
     // Fall into llvm_unreachable().
     break;
diff --git a/lld/lib/ReaderWriter/MachO/ArchHandler_x86.cpp b/lld/lib/ReaderWriter/MachO/ArchHandler_x86.cpp
index c278042cf8d..a2b6067efa6 100644
--- a/lld/lib/ReaderWriter/MachO/ArchHandler_x86.cpp
+++ b/lld/lib/ReaderWriter/MachO/ArchHandler_x86.cpp
@@ -461,9 +461,11 @@ void ArchHandler_x86::applyFixupFinal(const Reference &ref, uint8_t *loc,
   case modeCode:
   case modeData:
   case lazyPointer:
-  case lazyImmediateLocation:
     // do nothing
     break;
+  case lazyImmediateLocation:
+    *loc32 = ref.addend();
+    break;
   default:
     llvm_unreachable("invalid x86 Reference Kind");
     break;
diff --git a/lld/lib/ReaderWriter/MachO/ArchHandler_x86_64.cpp b/lld/lib/ReaderWriter/MachO/ArchHandler_x86_64.cpp
index b0c47af8259..ccab82419f7 100644
--- a/lld/lib/ReaderWriter/MachO/ArchHandler_x86_64.cpp
+++ b/lld/lib/ReaderWriter/MachO/ArchHandler_x86_64.cpp
@@ -521,8 +521,10 @@ void ArchHandler_x86_64::applyFixupFinal(
     *loc32 = fixupAddress - targetAddress + ref.addend();
     return;
   case lazyPointer:
+    // Do nothing
+    return;
   case lazyImmediateLocation:
-    // do nothing
+    *loc32 = ref.addend();
     return;
   case imageOffset:
   case imageOffsetGot:
diff --git a/lld/lib/ReaderWriter/MachO/MachONormalizedFileBinaryWriter.cpp b/lld/lib/ReaderWriter/MachO/MachONormalizedFileBinaryWriter.cpp
index 543b60fe8bc..102b185df52 100644
--- a/lld/lib/ReaderWriter/MachO/MachONormalizedFileBinaryWriter.cpp
+++ b/lld/lib/ReaderWriter/MachO/MachONormalizedFileBinaryWriter.cpp
@@ -125,6 +125,12 @@ private:
     void append_uleb128(uint64_t value) {
       llvm::encodeULEB128(value, _ostream);
     }
+    void append_uleb128Fixed(uint64_t value, unsigned byteCount) {
+      unsigned min = llvm::getULEB128Size(value);
+      assert(min <= byteCount);
+      unsigned pad = byteCount - min;
+      llvm::encodeULEB128(value, _ostream, pad);
+    }
     void append_sleb128(int64_t value) {
       llvm::encodeSLEB128(value, _ostream);
     }
@@ -999,6 +1005,7 @@ void MachOFileLayout::buildRebaseInfo() {
 
 void MachOFileLayout::buildBindInfo() {
   // TODO: compress bind info.
+  uint64_t lastAddend = 0;
   for (const BindLocation& entry : _file.bindingInfo) {
     _bindingInfo.append_byte(BIND_OPCODE_SET_TYPE_IMM | entry.kind);
     _bindingInfo.append_byte(BIND_OPCODE_SET_SEGMENT_AND_OFFSET_ULEB
@@ -1007,9 +1014,10 @@ void MachOFileLayout::buildBindInfo() {
     _bindingInfo.append_byte(BIND_OPCODE_SET_DYLIB_ORDINAL_IMM | entry.ordinal);
     _bindingInfo.append_byte(BIND_OPCODE_SET_SYMBOL_TRAILING_FLAGS_IMM);
     _bindingInfo.append_string(entry.symbolName);
-    if (entry.addend != 0) {
+    if (entry.addend != lastAddend) {
       _bindingInfo.append_byte(BIND_OPCODE_SET_ADDEND_SLEB);
       _bindingInfo.append_sleb128(entry.addend);
+      lastAddend = entry.addend;
     }
     _bindingInfo.append_byte(BIND_OPCODE_DO_BIND);
   }
@@ -1022,11 +1030,12 @@ void MachOFileLayout::buildLazyBindInfo() {
     _lazyBindingInfo.append_byte(BIND_OPCODE_SET_TYPE_IMM | entry.kind);
     _lazyBindingInfo.append_byte(BIND_OPCODE_SET_SEGMENT_AND_OFFSET_ULEB
                             | entry.segIndex);
-    _lazyBindingInfo.append_uleb128(entry.segOffset);
+    _lazyBindingInfo.append_uleb128Fixed(entry.segOffset, 5);
     _lazyBindingInfo.append_byte(BIND_OPCODE_SET_DYLIB_ORDINAL_IMM | entry.ordinal);
     _lazyBindingInfo.append_byte(BIND_OPCODE_SET_SYMBOL_TRAILING_FLAGS_IMM);
     _lazyBindingInfo.append_string(entry.symbolName);
     _lazyBindingInfo.append_byte(BIND_OPCODE_DO_BIND);
+    _lazyBindingInfo.append_byte(BIND_OPCODE_DONE);
   }
   _lazyBindingInfo.append_byte(BIND_OPCODE_DONE);
   _lazyBindingInfo.align(_is64 ? 8 : 4);
diff --git a/lld/lib/ReaderWriter/MachO/MachONormalizedFileFromAtoms.cpp b/lld/lib/ReaderWriter/MachO/MachONormalizedFileFromAtoms.cpp
index 6d0236bcc91..8333d099ccc 100644
--- a/lld/lib/ReaderWriter/MachO/MachONormalizedFileFromAtoms.cpp
+++ b/lld/lib/ReaderWriter/MachO/MachONormalizedFileFromAtoms.cpp
@@ -244,6 +244,7 @@ SectionInfo *Util::getFinalSection(DefinedAtom::ContentType atomType) {
     switch (atomType) {
     case DefinedAtom::typeCode:
     case DefinedAtom::typeStub:
+    case DefinedAtom::typeStubHelper:
       sectionAttrs = S_ATTR_PURE_INSTRUCTIONS;
       break;
     default:
diff --git a/lld/lib/ReaderWriter/MachO/StubsPass.cpp b/lld/lib/ReaderWriter/MachO/StubsPass.cpp
index 6f58b9954f9..38af5597618 100644
--- a/lld/lib/ReaderWriter/MachO/StubsPass.cpp
+++ b/lld/lib/ReaderWriter/MachO/StubsPass.cpp
@@ -24,6 +24,7 @@
 #include "lld/Core/Simple.h"
 #include "lld/ReaderWriter/MachOLinkingContext.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallVector.h"
 
 
 namespace lld {
@@ -107,7 +108,7 @@ private:
 class StubAtom : public SimpleDefinedAtom {
 public:
   StubAtom(const File &file, const ArchHandler::StubInfo &stubInfo)
-      : SimpleDefinedAtom(file), _stubInfo(stubInfo) { }
+      : SimpleDefinedAtom(file), _stubInfo(stubInfo){ }
 
   ContentType contentType() const override {
     return DefinedAtom::typeStub;
@@ -205,14 +206,8 @@ private:
 class StubsPass : public Pass {
 public:
   StubsPass(const MachOLinkingContext &context)
-    : _context(context)
-    , _archHandler(_context.archHandler())
-    , _stubInfo(_archHandler.stubInfo())
-    , _file("<mach-o Stubs pass>")
-    , _helperCommonAtom(nullptr)
-    , _helperCacheNLPAtom(nullptr)
-    , _helperBinderNLPAtom(nullptr) {
-  }
+    : _context(context), _archHandler(_context.archHandler()),
+    _stubInfo(_archHandler.stubInfo()), _file("<mach-o Stubs pass>") { }
 
 
   void perform(std::unique_ptr<MutableFile> &mergedFile) override {
@@ -230,7 +225,7 @@ public:
         assert(target != nullptr);
         if (isa<SharedLibraryAtom>(target)) {
           // Calls to shared libraries go through stubs.
-          replaceCalleeWithStub(target, ref);
+          _targetToUses[target].push_back(ref);
           continue;
         }
         const DefinedAtom *defTarget = dyn_cast<DefinedAtom>(target);
@@ -238,47 +233,93 @@ public:
           // Calls to interposable functions in same linkage unit must also go
           // through a stub.
           assert(defTarget->scope() != DefinedAtom::scopeTranslationUnit);
-          replaceCalleeWithStub(target, ref);
+          _targetToUses[target].push_back(ref);
         }
       }
     }
+
     // Exit early if no stubs needed.
-    if (_targetToStub.empty())
+    if (_targetToUses.empty())
       return;
-    
+
+    // First add help-common and GOT slots used by lazy binding.
+    SimpleDefinedAtom *helperCommonAtom =
+        new (_file.allocator()) StubHelperCommonAtom(_file, _stubInfo);
+    SimpleDefinedAtom *helperCacheNLPAtom =
+        new (_file.allocator()) NonLazyPointerAtom(_file, _context.is64Bit());
+    SimpleDefinedAtom *helperBinderNLPAtom =
+        new (_file.allocator()) NonLazyPointerAtom(_file, _context.is64Bit());
+    addReference(helperCommonAtom, _stubInfo.stubHelperCommonReferenceToCache,
+                 helperCacheNLPAtom);
+    addOptReference(
+        helperCommonAtom, _stubInfo.stubHelperCommonReferenceToCache,
+        _stubInfo.optStubHelperCommonReferenceToCache, helperCacheNLPAtom);
+    addReference(helperCommonAtom, _stubInfo.stubHelperCommonReferenceToBinder,
+                 helperBinderNLPAtom);
+    addOptReference(
+        helperCommonAtom, _stubInfo.stubHelperCommonReferenceToBinder,
+        _stubInfo.optStubHelperCommonReferenceToBinder, helperBinderNLPAtom);
+    mergedFile->addAtom(*helperCommonAtom);
+    mergedFile->addAtom(*helperBinderNLPAtom);
+    mergedFile->addAtom(*helperCacheNLPAtom);
+
     // Add reference to dyld_stub_binder in libSystem.dylib
-    if (_helperBinderNLPAtom) {
-      bool found = false;
-      for (const SharedLibraryAtom *atom : mergedFile->sharedLibrary()) {
-        if (atom->name().equals(_stubInfo.binderSymbolName)) {
-          addReference(_helperBinderNLPAtom,  
-                       _stubInfo.nonLazyPointerReferenceToBinder, atom);
-          found = true;
-          break;
-        }
+    bool binderFound = false;
+    for (const SharedLibraryAtom *atom : mergedFile->sharedLibrary()) {
+      if (atom->name().equals(_stubInfo.binderSymbolName)) {
+        addReference(helperBinderNLPAtom,
+                     _stubInfo.nonLazyPointerReferenceToBinder, atom);
+        binderFound = true;
+        break;
       }
-      assert(found && "dyld_stub_binder not found");
     }
-    
-    // Add all stubs to master file.
-    for (auto it : _targetToStub) {
-      mergedFile->addAtom(*it.second);
-    }
-    // Add helper code atoms.
-    mergedFile->addAtom(*_helperCommonAtom);
-    for (const DefinedAtom *lp : _stubHelperAtoms) {
-      mergedFile->addAtom(*lp);
-    }
-    // Add GOT slots used for lazy binding.
-    mergedFile->addAtom(*_helperBinderNLPAtom);
-    mergedFile->addAtom(*_helperCacheNLPAtom);
-    // Add all lazy pointers to master file.
-    for (const DefinedAtom *lp : _lazyPointers) {
+    assert(binderFound && "dyld_stub_binder not found");
+
+    // Sort targets by name, so stubs and lazy pointers are consistent
+    std::vector<const Atom *> targetsNeedingStubs;
+    for (auto it : _targetToUses)
+      targetsNeedingStubs.push_back(it.first);
+    std::sort(targetsNeedingStubs.begin(), targetsNeedingStubs.end(),
+              [](const Atom * left, const Atom * right) {
+      return (left->name().compare(right->name()) < 0);
+    });
+
+    // Make and append stubs, lazy pointers, and helpers in alphabetical order.
+    unsigned lazyOffset = 0;
+    for (const Atom *target : targetsNeedingStubs) {
+      StubAtom *stub = new (_file.allocator()) StubAtom(_file, _stubInfo);
+      LazyPointerAtom *lp =
+          new (_file.allocator()) LazyPointerAtom(_file, _context.is64Bit());
+      StubHelperAtom *helper =
+          new (_file.allocator()) StubHelperAtom(_file, _stubInfo);
+
+      addReference(stub, _stubInfo.stubReferenceToLP, lp);
+      addOptReference(stub, _stubInfo.stubReferenceToLP,
+                      _stubInfo.optStubReferenceToLP, lp);
+      addReference(lp, _stubInfo.lazyPointerReferenceToHelper, helper);
+      addReference(lp, _stubInfo.lazyPointerReferenceToFinal, target);
+      addReference(helper, _stubInfo.stubHelperReferenceToImm, helper);
+      addReferenceAddend(helper, _stubInfo.stubHelperReferenceToImm, helper,
+                         lazyOffset);
+      addReference(helper, _stubInfo.stubHelperReferenceToHelperCommon,
+                   helperCommonAtom);
+
+      mergedFile->addAtom(*stub);
       mergedFile->addAtom(*lp);
+      mergedFile->addAtom(*helper);
+
+      // Update each reference to use stub.
+      for (const Reference *ref : _targetToUses[target]) {
+        assert(ref->target() == target);
+        // Switch call site to reference stub atom instead.
+        const_cast<Reference *>(ref)->setTarget(stub);
+      }
+
+      // Calculate new offset
+      lazyOffset += target->name().size() + 12;
     }
   }
 
-
 private:
 
   bool noTextRelocs() {
@@ -288,51 +329,6 @@ private:
   bool isCallSite(const Reference &ref) {
     return _archHandler.isCallSite(ref);
   }
-
-  void replaceCalleeWithStub(const Atom *target, const Reference *ref) {
-    // Make file-format specific stub and other support atoms.
-    const DefinedAtom *stub = this->getStub(*target);
-    assert(stub != nullptr);
-    // Switch call site to reference stub atom instead.
-    const_cast<Reference *>(ref)->setTarget(stub);
-  }
-
-  const DefinedAtom* getStub(const Atom& target) {
-    auto pos = _targetToStub.find(&target);
-    if ( pos != _targetToStub.end() ) {
-      // Reuse an existing stub.
-      assert(pos->second != nullptr);
-      return pos->second;
-    }
-    else {
-      // There is no existing stub, so create a new one.
-      return this->makeStub(target);
-    }
-  }
-
-  const DefinedAtom* makeStub(const Atom &target) {
-    SimpleDefinedAtom* stub   = new (_file.allocator()) 
-                                                    StubAtom(_file, _stubInfo);
-    SimpleDefinedAtom* lp     = new (_file.allocator()) 
-                                    LazyPointerAtom(_file, _context.is64Bit());
-    SimpleDefinedAtom* helper = new (_file.allocator()) 
-                                              StubHelperAtom(_file, _stubInfo);
-   
-    addReference(stub, _stubInfo.stubReferenceToLP, lp);
-    addOptReference(stub, _stubInfo.stubReferenceToLP,
-                    _stubInfo.optStubReferenceToLP, lp);
-    addReference(lp, _stubInfo.lazyPointerReferenceToHelper, helper);
-    addReference(lp, _stubInfo.lazyPointerReferenceToFinal, &target);
-    addReference(helper, _stubInfo.stubHelperReferenceToImm, helper);
-    addReference(helper, _stubInfo.stubHelperReferenceToHelperCommon, 
-                 helperCommon());
-    
-    _stubHelperAtoms.push_back(helper);
-    _targetToStub[&target] = stub;
-    _lazyPointers.push_back(lp);
-    
-    return stub;
-  }
  
   void addReference(SimpleDefinedAtom* atom,
                     const ArchHandler::ReferenceInfo &refInfo,
@@ -342,6 +338,13 @@ private:
                       target, refInfo.addend);
   }
 
+  void addReferenceAddend(SimpleDefinedAtom *atom,
+                          const ArchHandler::ReferenceInfo &refInfo,
+                          const lld::Atom *target, uint64_t addend) {
+    atom->addReference(Reference::KindNamespace::mach_o, refInfo.arch,
+                       refInfo.kind, refInfo.offset, target, addend);
+  }
+
    void addOptReference(SimpleDefinedAtom* atom,
                     const ArchHandler::ReferenceInfo &refInfo,
                     const ArchHandler::OptionalRefInfo &optRef,
@@ -353,44 +356,14 @@ private:
                       target, optRef.addend);
   }
 
-  const DefinedAtom* helperCommon() {
-    if ( !_helperCommonAtom ) {
-      // Lazily create common helper code and data.
-      _helperCommonAtom    = new (_file.allocator()) 
-                                         StubHelperCommonAtom(_file, _stubInfo);
-      _helperCacheNLPAtom  = new (_file.allocator()) 
-                                  NonLazyPointerAtom(_file, _context.is64Bit());
-      _helperBinderNLPAtom = new (_file.allocator()) 
-                                  NonLazyPointerAtom(_file, _context.is64Bit());
-      addReference(_helperCommonAtom, 
-                   _stubInfo.stubHelperCommonReferenceToCache,
-                   _helperCacheNLPAtom);
-      addOptReference(_helperCommonAtom,
-                      _stubInfo.stubHelperCommonReferenceToCache,
-                      _stubInfo.optStubHelperCommonReferenceToCache,
-                      _helperCacheNLPAtom);
-      addReference(_helperCommonAtom,
-                   _stubInfo.stubHelperCommonReferenceToBinder,
-                   _helperBinderNLPAtom);
-      addOptReference(_helperCommonAtom,
-                      _stubInfo.stubHelperCommonReferenceToBinder,
-                      _stubInfo.optStubHelperCommonReferenceToBinder,
-                      _helperBinderNLPAtom);
-    }
-    return _helperCommonAtom;
-  }
-
+  typedef llvm::DenseMap<const Atom*,
+                         llvm::SmallVector<const Reference *, 8>> TargetToUses;
 
   const MachOLinkingContext                      &_context;
   mach_o::ArchHandler                            &_archHandler;
   const ArchHandler::StubInfo                    &_stubInfo;
   MachOFile                                       _file;
-  llvm::DenseMap<const Atom*, const DefinedAtom*> _targetToStub;
-  std::vector<const DefinedAtom*>                 _lazyPointers;
-  std::vector<const DefinedAtom*>                 _stubHelperAtoms;
-  SimpleDefinedAtom                              *_helperCommonAtom;
-  SimpleDefinedAtom                              *_helperCacheNLPAtom;
-  SimpleDefinedAtom                              *_helperBinderNLPAtom;
+  TargetToUses                                    _targetToUses;
 };