[ORC] Merge redundant jitlink::Symbol -> JITSymbolFlags mappings.

Adds a getJITSymbolFlagsForSymbol function that returns the JITSymbolFlags
for a given jitlink::Symbol, and replaces severalredundant copies of that
mapping with calls to the new function. This fixes a bug in
LinkGraphMaterializationUnit::scanLinkGraph where we were failing to set the
JITSymbolFlags::Weak flag for weak symbols, and a bug in
ObjectLinkingLayer::claimOrExternalizeWeakAndCommonSymbols where we were
failing to set the JITSymbolFlags::Callable flag for callable symbols.

diff --git a/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp b/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp
index 2b11c47..b889992 100644 (file)
--- a/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp
@@ -22,6 +22,21 @@ using namespace llvm::orc;
 
 namespace {
 
+JITSymbolFlags getJITSymbolFlagsForSymbol(Symbol &Sym) {
+  JITSymbolFlags Flags;
+
+  if (Sym.getLinkage() == Linkage::Weak)
+    Flags |= JITSymbolFlags::Weak;
+
+  if (Sym.getScope() == Scope::Default)
+    Flags |= JITSymbolFlags::Exported;
+
+  if (Sym.isCallable())
+    Flags |= JITSymbolFlags::Callable;
+
+  return Flags;
+}
+
 class LinkGraphMaterializationUnit : public MaterializationUnit {
 public:
   static std::unique_ptr<LinkGraphMaterializationUnit>
@@ -48,14 +63,8 @@ private:
         continue;
       assert(Sym->hasName() && "Anonymous non-local symbol?");
 
-      JITSymbolFlags Flags;
-      if (Sym->getScope() == Scope::Default)
-        Flags |= JITSymbolFlags::Exported;
-
-      if (Sym->isCallable())
-        Flags |= JITSymbolFlags::Callable;
-
-      LGI.SymbolFlags[ES.intern(Sym->getName())] = Flags;
+      LGI.SymbolFlags[ES.intern(Sym->getName())] =
+          getJITSymbolFlagsForSymbol(*Sym);
     }
 
     if (hasInitializerSection(G))
@@ -189,14 +198,7 @@ public:
     for (auto *Sym : G.defined_symbols())
       if (Sym->hasName() && Sym->getScope() != Scope::Local) {
         auto InternedName = ES.intern(Sym->getName());
-        JITSymbolFlags Flags;
-
-        if (Sym->isCallable())
-          Flags |= JITSymbolFlags::Callable;
-        if (Sym->getScope() == Scope::Default)
-          Flags |= JITSymbolFlags::Exported;
-        if (Sym->getLinkage() == Linkage::Weak)
-          Flags |= JITSymbolFlags::Weak;
+        auto Flags = getJITSymbolFlagsForSymbol(*Sym);
 
         InternedResult[InternedName] =
             JITEvaluatedSymbol(Sym->getAddress().getValue(), Flags);
@@ -210,13 +212,7 @@ public:
     for (auto *Sym : G.absolute_symbols())
       if (Sym->hasName() && Sym->getScope() != Scope::Local) {
         auto InternedName = ES.intern(Sym->getName());
-        JITSymbolFlags Flags;
-        if (Sym->isCallable())
-          Flags |= JITSymbolFlags::Callable;
-        if (Sym->getScope() == Scope::Default)
-          Flags |= JITSymbolFlags::Exported;
-        if (Sym->getLinkage() == Linkage::Weak)
-          Flags |= JITSymbolFlags::Weak;
+        auto Flags = getJITSymbolFlagsForSymbol(*Sym);
         InternedResult[InternedName] =
             JITEvaluatedSymbol(Sym->getAddress().getValue(), Flags);
         if (AutoClaim && !MR->getSymbols().count(InternedName)) {
@@ -407,10 +403,8 @@ private:
           Sym->getScope() != Scope::Local) {
         auto Name = ES.intern(Sym->getName());
         if (!MR->getSymbols().count(ES.intern(Sym->getName()))) {
-          JITSymbolFlags SF = JITSymbolFlags::Weak;
-          if (Sym->getScope() == Scope::Default)
-            SF |= JITSymbolFlags::Exported;
-          NewSymbolsToClaim[Name] = SF;
+          NewSymbolsToClaim[Name] =
+              getJITSymbolFlagsForSymbol(*Sym) | JITSymbolFlags::Weak;
           NameToSym.push_back(std::make_pair(std::move(Name), Sym));
         }
       }

diff --git a/llvm/unittests/ExecutionEngine/Orc/ObjectLinkingLayerTest.cpp b/llvm/unittests/ExecutionEngine/Orc/ObjectLinkingLayerTest.cpp
index 1f638f4..605fb25 100644 (file)
--- a/llvm/unittests/ExecutionEngine/Orc/ObjectLinkingLayerTest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/ObjectLinkingLayerTest.cpp
@@ -48,10 +48,75 @@ TEST_F(ObjectLinkingLayerTest, AddLinkGraph) {
                                    orc::ExecutorAddr(0x1000), 8, 0);
   G->addDefinedSymbol(B1, 4, "_X", 4, Linkage::Strong, Scope::Default, false,
                       false);
+  G->addDefinedSymbol(B1, 4, "_Y", 4, Linkage::Weak, Scope::Default, false,
+                      false);
+  G->addDefinedSymbol(B1, 4, "_Z", 4, Linkage::Strong, Scope::Hidden, false,
+                      false);
+  G->addDefinedSymbol(B1, 4, "_W", 4, Linkage::Strong, Scope::Default, true,
+                      false);
 
   EXPECT_THAT_ERROR(ObjLinkingLayer.add(JD, std::move(G)), Succeeded());
 
   EXPECT_THAT_EXPECTED(ES.lookup(&JD, "_X"), Succeeded());
 }
 
+TEST_F(ObjectLinkingLayerTest, ClaimLateDefinedWeakSymbols) {
+  // Check that claiming weak symbols works as expected.
+  //
+  // To do this we'll need a custom plugin to inject some new symbols during
+  // the link.
+  class TestPlugin : public ObjectLinkingLayer::Plugin {
+  public:
+    void modifyPassConfig(MaterializationResponsibility &MR,
+                          jitlink::LinkGraph &G,
+                          jitlink::PassConfiguration &Config) override {
+      Config.PrePrunePasses.insert(
+          Config.PrePrunePasses.begin(), [](LinkGraph &G) {
+            auto *DataSec = G.findSectionByName("__data");
+            auto &DataBlock = G.createContentBlock(
+                *DataSec, BlockContent, orc::ExecutorAddr(0x2000), 8, 0);
+            G.addDefinedSymbol(DataBlock, 4, "_x", 4, Linkage::Weak,
+                               Scope::Default, false, false);
+
+            auto &TextSec =
+                G.createSection("__text", MemProt::Read | MemProt::Write);
+            auto &FuncBlock = G.createContentBlock(
+                TextSec, BlockContent, orc::ExecutorAddr(0x3000), 8, 0);
+            G.addDefinedSymbol(FuncBlock, 4, "_f", 4, Linkage::Weak,
+                               Scope::Default, true, false);
+
+            return Error::success();
+          });
+    }
+
+    Error notifyFailed(MaterializationResponsibility &MR) override {
+      llvm_unreachable("unexpected error");
+    }
+
+    Error notifyRemovingResources(JITDylib &JD, ResourceKey K) override {
+      return Error::success();
+    }
+    void notifyTransferringResources(JITDylib &JD, ResourceKey DstKey,
+                                     ResourceKey SrcKey) override {
+      llvm_unreachable("unexpected resource transfer");
+    }
+  };
+
+  ObjLinkingLayer.addPlugin(std::make_unique<TestPlugin>());
+
+  auto G =
+      std::make_unique<LinkGraph>("foo", Triple("x86_64-apple-darwin"), 8,
+                                  support::little, x86_64::getEdgeKindName);
+
+  auto &DataSec = G->createSection("__data", MemProt::Read | MemProt::Write);
+  auto &DataBlock = G->createContentBlock(DataSec, BlockContent,
+                                          orc::ExecutorAddr(0x1000), 8, 0);
+  G->addDefinedSymbol(DataBlock, 4, "_anchor", 4, Linkage::Weak, Scope::Default,
+                      false, true);
+
+  EXPECT_THAT_ERROR(ObjLinkingLayer.add(JD, std::move(G)), Succeeded());
+
+  EXPECT_THAT_EXPECTED(ES.lookup(&JD, "_anchor"), Succeeded());
+}
+
 } // end anonymous namespace