1 files changed, 24 insertions, 8 deletions

diff --git a/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.cpp b/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.cpp
index d4270b5aa796..7b594fd2c0ea 100644
--- a/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.cpp
+++ b/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.cpp
@@ -151,9 +151,12 @@ JITLinkerBase::SegmentLayoutMap JITLinkerBase::layOutBlocks() {
   for (auto &KV : Layout) {
 
     auto CompareBlocks = [](const Block *LHS, const Block *RHS) {
+      // Sort by section, address and size
       if (LHS->getSection().getOrdinal() != RHS->getSection().getOrdinal())
         return LHS->getSection().getOrdinal() < RHS->getSection().getOrdinal();
-      return LHS->getOrdinal() < RHS->getOrdinal();
+      if (LHS->getAddress() != RHS->getAddress())
+        return LHS->getAddress() < RHS->getAddress();
+      return LHS->getSize() < RHS->getSize();
     };
 
     auto &SegLists = KV.second;
@@ -254,25 +257,35 @@ Error JITLinkerBase::allocateSegments(const SegmentLayoutMap &Layout) {
   return Error::success();
 }
 
-DenseSet<StringRef> JITLinkerBase::getExternalSymbolNames() const {
+JITLinkContext::LookupMap JITLinkerBase::getExternalSymbolNames() const {
   // Identify unresolved external symbols.
-  DenseSet<StringRef> UnresolvedExternals;
+  JITLinkContext::LookupMap UnresolvedExternals;
   for (auto *Sym : G->external_symbols()) {
     assert(Sym->getAddress() == 0 &&
            "External has already been assigned an address");
     assert(Sym->getName() != StringRef() && Sym->getName() != "" &&
            "Externals must be named");
-    UnresolvedExternals.insert(Sym->getName());
+    SymbolLookupFlags LookupFlags =
+        Sym->getLinkage() == Linkage::Weak
+            ? SymbolLookupFlags::WeaklyReferencedSymbol
+            : SymbolLookupFlags::RequiredSymbol;
+    UnresolvedExternals[Sym->getName()] = LookupFlags;
   }
   return UnresolvedExternals;
 }
 
 void JITLinkerBase::applyLookupResult(AsyncLookupResult Result) {
   for (auto *Sym : G->external_symbols()) {
+    assert(Sym->getOffset() == 0 &&
+           "External symbol is not at the start of its addressable block");
     assert(Sym->getAddress() == 0 && "Symbol already resolved");
     assert(!Sym->isDefined() && "Symbol being resolved is already defined");
-    assert(Result.count(Sym->getName()) && "Missing resolution for symbol");
-    Sym->getAddressable().setAddress(Result[Sym->getName()].getAddress());
+    auto ResultI = Result.find(Sym->getName());
+    if (ResultI != Result.end())
+      Sym->getAddressable().setAddress(ResultI->second.getAddress());
+    else
+      assert(Sym->getLinkage() == Linkage::Weak &&
+             "Failed to resolve non-weak reference");
   }
 
   LLVM_DEBUG({
@@ -282,8 +295,11 @@ void JITLinkerBase::applyLookupResult(AsyncLookupResult Result) {
              << formatv("{0:x16}", Sym->getAddress()) << "\n";
   });
   assert(llvm::all_of(G->external_symbols(),
-                      [](Symbol *Sym) { return Sym->getAddress() != 0; }) &&
-         "All symbols should have been resolved by this point");
+                      [](Symbol *Sym) {
+                        return Sym->getAddress() != 0 ||
+                               Sym->getLinkage() == Linkage::Weak;
+                      }) &&
+         "All strong external symbols should have been resolved by now");
 }
 
 void JITLinkerBase::deallocateAndBailOut(Error Err) {