[JITLink][COFF] Use regular external symbol resolution for __ImageBase.

This fixes a concurrency bug in the COFF_x86_64 backend: lookupAsync was used
to find __ImageBase without blocking to wait for the result.

Rather than adding promises / futures, this patch just adds an external
__ImageBase symbol if needed. This is the canonical JITLink solution for
resolving external symbols, and is simpler and more concurrency friendly than
using promises / futures with lookupAsync.

Author: lhames

llvm/lib/ExecutionEngine/JITLink/COFFLinkGraphBuilder.cpp

@@ -642,5 +642,39 @@ COFFLinkGraphBuilder::exportCOMDATSymbol(COFFSymbolIndex SymIndex,
   return GSym;
 }
 
+Symbol *GetImageBaseSymbol::operator()(LinkGraph &G) {
+  if (ImageBase)
+    return *ImageBase;
+
+  auto IBN = G.intern(ImageBaseName);
+
+  // Check external symbols for image base.
+  for (auto *Sym : G.external_symbols()) {
+    if (Sym->getName() == IBN) {
+      ImageBase = Sym;
+      return Sym;
+    }
+  }
+
+  // Check absolute symbols (unexpected, but legal).
+  for (auto *Sym : G.absolute_symbols()) {
+    if (Sym->getName() == IBN) {
+      ImageBase = Sym;
+      return Sym;
+    }
+  }
+
+  // Finally, check defined symbols.
+  for (auto *Sym : G.defined_symbols()) {
+    if (Sym->hasName() && Sym->getName() == IBN) {
+      ImageBase = Sym;
+      return Sym;
+    }
+  }
+
+  ImageBase = nullptr;
+  return nullptr;
+}
+
 } // namespace jitlink
 } // namespace llvm

llvm/lib/ExecutionEngine/JITLink/COFFLinkGraphBuilder.h

@@ -79,6 +79,12 @@ class COFFLinkGraphBuilder {
     return GraphBlocks[SecIndex];
   }
 
+  Symbol &addImageBaseSymbol(StringRef Name = "__ImageBase") {
+    auto &ImageBase = G->addExternalSymbol(G->intern(Name), 0, true);
+    ImageBase.setLive(true);
+    return ImageBase;
+  }
+
   object::COFFObjectFile::section_iterator_range sections() const {
     return Obj.sections();
   }
@@ -216,6 +222,18 @@ Error COFFLinkGraphBuilder::forEachRelocation(const object::SectionRef &RelSec,
   return Error::success();
 }
 
+class GetImageBaseSymbol {
+public:
+  GetImageBaseSymbol(StringRef ImageBaseName = "__ImageBase")
+      : ImageBaseName(ImageBaseName) {}
+  Symbol *operator()(LinkGraph &G);
+  void reset() { ImageBase = std::nullopt; }
+
+private:
+  StringRef ImageBaseName;
+  std::optional<Symbol *> ImageBase;
+};
+
 } // end namespace jitlink
 } // end namespace llvm
 

llvm/lib/ExecutionEngine/JITLink/COFF_x86_64.cpp

@@ -93,9 +93,12 @@ class COFFLinkGraphBuilder_x86_64 : public COFFLinkGraphBuilder {
 
     Edge::Kind Kind = Edge::Invalid;
     const char *FixupPtr = BlockToFix.getContent().data() + Offset;
+    Symbol *ImageBase = GetImageBaseSymbol()(getGraph());
 
     switch (Rel.getType()) {
     case COFF::RelocationTypeAMD64::IMAGE_REL_AMD64_ADDR32NB: {
+      if (!ImageBase)
+        ImageBase = &addImageBaseSymbol();
       Kind = EdgeKind_coff_x86_64::Pointer32NB;
       Addend = *reinterpret_cast<const support::little32_t *>(FixupPtr);
       break;
@@ -192,20 +195,15 @@ class COFFLinkGraphBuilder_x86_64 : public COFFLinkGraphBuilder {
 
 class COFFLinkGraphLowering_x86_64 {
 public:
-  COFFLinkGraphLowering_x86_64(std::shared_ptr<orc::SymbolStringPool> SSP)
-      : SSP(std::move(SSP)) {
-    ImageBaseName = this->SSP->intern("__ImageBase");
-  }
   // Lowers COFF x86_64 specific edges to generic x86_64 edges.
-  Error lowerCOFFRelocationEdges(LinkGraph &G, JITLinkContext &Ctx) {
+  Error operator()(LinkGraph &G) {
     for (auto *B : G.blocks()) {
       for (auto &E : B->edges()) {
         switch (E.getKind()) {
         case EdgeKind_coff_x86_64::Pointer32NB: {
-          auto ImageBase = getImageBaseAddress(G, Ctx);
-          if (!ImageBase)
-            return ImageBase.takeError();
-          E.setAddend(E.getAddend() - ImageBase->getValue());
+          auto ImageBase = GetImageBase(G);
+          assert(ImageBase && "__ImageBase symbol must be defined");
+          E.setAddend(E.getAddend() - ImageBase->getAddress().getValue());
           E.setKind(x86_64::Pointer32);
           break;
         }
@@ -237,61 +235,18 @@ class COFFLinkGraphLowering_x86_64 {
   }
 
 private:
-  const orc::SymbolStringPtr &getImageBaseSymbolName() const {
-    return this->ImageBaseName;
-  }
-
   orc::ExecutorAddr getSectionStart(Section &Sec) {
     if (!SectionStartCache.count(&Sec)) {
       SectionRange Range(Sec);
       SectionStartCache[&Sec] = Range.getStart();
+      return Range.getStart();
     }
     return SectionStartCache[&Sec];
   }
 
-  Expected<orc::ExecutorAddr> getImageBaseAddress(LinkGraph &G,
-                                                  JITLinkContext &Ctx) {
-    if (this->ImageBase)
-      return this->ImageBase;
-    for (auto *S : G.defined_symbols())
-      if (S->getName() == getImageBaseSymbolName()) {
-        this->ImageBase = S->getAddress();
-        return this->ImageBase;
-      }
-
-    JITLinkContext::LookupMap Symbols;
-    Symbols[getImageBaseSymbolName()] = SymbolLookupFlags::RequiredSymbol;
-    orc::ExecutorAddr ImageBase;
-    Error Err = Error::success();
-    Ctx.lookup(Symbols,
-               createLookupContinuation([&](Expected<AsyncLookupResult> LR) {
-                 ErrorAsOutParameter _(Err);
-                 if (!LR) {
-                   Err = LR.takeError();
-                   return;
-                 }
-                 ImageBase = LR->begin()->second.getAddress();
-               }));
-    if (Err)
-      return std::move(Err);
-    this->ImageBase = ImageBase;
-    return ImageBase;
-  }
-
+  GetImageBaseSymbol GetImageBase;
   DenseMap<Section *, orc::ExecutorAddr> SectionStartCache;
-  orc::ExecutorAddr ImageBase;
-  std::shared_ptr<orc::SymbolStringPool> SSP;
-  orc::SymbolStringPtr ImageBaseName = nullptr;
 };
-
-Error lowerEdges_COFF_x86_64(LinkGraph &G, JITLinkContext *Ctx) {
-  LLVM_DEBUG(dbgs() << "Lowering COFF x86_64 edges:\n");
-  COFFLinkGraphLowering_x86_64 GraphLowering(G.getSymbolStringPool());
-  if (auto Err = GraphLowering.lowerCOFFRelocationEdges(G, *Ctx))
-    return Err;
-
-  return Error::success();
-}
 } // namespace
 
 namespace llvm {
@@ -349,9 +304,7 @@ void link_COFF_x86_64(std::unique_ptr<LinkGraph> G,
       Config.PrePrunePasses.push_back(markAllSymbolsLive);
 
     // Add COFF edge lowering passes.
-    JITLinkContext *CtxPtr = Ctx.get();
-    Config.PreFixupPasses.push_back(
-        [CtxPtr](LinkGraph &G) { return lowerEdges_COFF_x86_64(G, CtxPtr); });
+    Config.PreFixupPasses.push_back(COFFLinkGraphLowering_x86_64());
   }
 
   if (auto Err = Ctx->modifyPassConfig(*G, Config))