[ExecutionEngine] Add weak symbol support to RuntimeDyld

Support weak symbols by first looking up if there is an externally visible symbol we can find,
and only if that fails using the one in the object file we're loading.

Reviewed By: lhames
Subscribers: llvm-commits
Differential Revision: http://reviews.llvm.org/D6950

llvm-svn: 227228

7 files changed, 93 insertions, 2 deletions

diff --git a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyld.cpp b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyld.cpp
index fb53ba96055..81b8eb2238f 100644
--- a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyld.cpp
+++ b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyld.cpp
@@ -168,6 +168,7 @@ RuntimeDyldImpl::loadObjectImpl(const object::ObjectFile &Obj) {
     uint32_t Flags = I->getFlags();
 
     bool IsCommon = Flags & SymbolRef::SF_Common;
+    bool IsWeak = Flags & SymbolRef::SF_Weak;
     if (IsCommon)
       CommonSymbols.push_back(*I);
     else {
@@ -188,6 +189,11 @@ RuntimeDyldImpl::loadObjectImpl(const object::ObjectFile &Obj) {
           continue;
         StringRef SectionData;
         Check(SI->getContents(SectionData));
+        // TODO: It make make sense to delay emitting the section for weak
+        // symbols until they are actually required, but that's not possible
+        // currently, because we only know whether we will need the symbol
+        // in resolveRelocations, which happens after we have already finalized
+        // the Load.
         bool IsCode = SI->isText();
         unsigned SectionID =
             findOrEmitSection(Obj, *SI, IsCode, LocalSections);
@@ -198,7 +204,11 @@ RuntimeDyldImpl::loadObjectImpl(const object::ObjectFile &Obj) {
         SymbolInfo::Visibility Vis =
           (Flags & SymbolRef::SF_Exported) ?
             SymbolInfo::Default : SymbolInfo::Hidden;
-        GlobalSymbolTable[Name] = SymbolInfo(SectionID, SectOffset, Vis);
+        if (!IsWeak) {
+          GlobalSymbolTable[Name] = SymbolInfo(SectionID, SectOffset, Vis);
+        } else {
+          WeakSymbolTable[Name] = SymbolInfo(SectionID, SectOffset, Vis);
+        }
       }
     }
   }
@@ -768,6 +778,21 @@ void RuntimeDyldImpl::resolveExternalSymbols() {
                SymInfo.getOffset();
       }
 
+      // If we didn't find the symbol yet, and it is present in the weak symbol
+      // table, the definition from this object file needs to be used, so emit
+      // it now
+      if (!Addr) {
+        RTDyldSymbolTable::const_iterator Loc = WeakSymbolTable.find(Name);
+        if (Loc != WeakSymbolTable.end()) {
+          SymbolInfo SymInfo = Loc->second;
+          Addr = getSectionLoadAddress(SymInfo.getSectionID()) + SymInfo.getOffset();
+          // Since the weak symbol is now, materialized, add it to the
+          // GlobalSymbolTable. If somebody later asks the ExecutionEngine
+          // for the address of this symbol that's where it'll look
+          GlobalSymbolTable[Name] = SymInfo;
+        }
+      }
+
       // FIXME: Implement error handling that doesn't kill the host program!
       if (!Addr)
         report_fatal_error("Program used external function '" + Name +
@@ -784,6 +809,7 @@ void RuntimeDyldImpl::resolveExternalSymbols() {
 
     ExternalSymbolRelocations.erase(i);
   }
+  WeakSymbolTable.clear();
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldELF.cpp b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldELF.cpp
index 0f3ca0f2f39..a01b063831a 100644
--- a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldELF.cpp
+++ b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldELF.cpp
@@ -947,6 +947,7 @@ relocation_iterator RuntimeDyldELF::processRelocationRef(
       break;
     }
     case SymbolRef::ST_Data:
+    case SymbolRef::ST_Function:
     case SymbolRef::ST_Unknown: {
       Value.SymbolName = TargetName.data();
       Value.Addend = Addend;
diff --git a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldImpl.h b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldImpl.h
index f37a9a768a2..69d6cf53034 100644
--- a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldImpl.h
+++ b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldImpl.h
@@ -203,6 +203,9 @@ protected:
   // A global symbol table for symbols from all loaded modules.
   RTDyldSymbolTable GlobalSymbolTable;
 
+  // Like the global symbol table but for weak symbols
+  RTDyldSymbolTable WeakSymbolTable;
+
   // Keep a map of common symbols to their info pairs
   typedef std::vector<SymbolRef> CommonSymbolList;
 
diff --git a/llvm/unittests/ExecutionEngine/MCJIT/CMakeLists.txt b/llvm/unittests/ExecutionEngine/MCJIT/CMakeLists.txt
index b10cbb4c9ea..dcf8264203e 100644
--- a/llvm/unittests/ExecutionEngine/MCJIT/CMakeLists.txt
+++ b/llvm/unittests/ExecutionEngine/MCJIT/CMakeLists.txt
@@ -1,5 +1,6 @@
 set(LLVM_LINK_COMPONENTS
   Analysis
+  AsmParser
   Core
   ExecutionEngine
   IPO
diff --git a/llvm/unittests/ExecutionEngine/MCJIT/MCJITTest.cpp b/llvm/unittests/ExecutionEngine/MCJIT/MCJITTest.cpp
index 64d8c2fca04..f9ca0fc8717 100644
--- a/llvm/unittests/ExecutionEngine/MCJIT/MCJITTest.cpp
+++ b/llvm/unittests/ExecutionEngine/MCJIT/MCJITTest.cpp
@@ -13,6 +13,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ExecutionEngine/MCJIT.h"
+#include "llvm/Support/DynamicLibrary.h"
 #include "MCJITTestBase.h"
 #include "gtest/gtest.h"
 
@@ -199,4 +200,45 @@ TEST_F(MCJITTest, multiple_decl_lookups) {
   EXPECT_EQ(A, B) << "Repeat calls to getPointerToFunction fail.";
 }
 
+// Test weak symbol linking when the weak symbol is present in a shared
+// library
+TEST_F(MCJITTest, weak_symbol_present) {
+  SKIP_UNSUPPORTED_PLATFORM;
+
+  int FakeWeakSymbol;
+  llvm::sys::DynamicLibrary::AddSymbol("FakeWeakSymbol", &FakeWeakSymbol);
+  createJITFromAssembly(
+  "$FakeWeakSymbol = comdat any\n"
+  "@FakeWeakSymbol = linkonce_odr global i32 42, comdat, align 4\n"
+  "define i32 @weak_test(i32* %arg) {\n"
+  "    %r = icmp eq i32* %arg, @FakeWeakSymbol\n"
+  "    %ret = zext i1 %r to i32\n"
+  "    ret i32 %ret\n"
+  "  }");
+
+  uint64_t Addr = TheJIT->getFunctionAddress("weak_test");;
+  EXPECT_TRUE(Addr != 0);
+  int32_t(*FuncPtr)(int32_t *) = (int32_t(*)(int32_t *))Addr;
+  EXPECT_EQ(FuncPtr(&FakeWeakSymbol),1);
+  EXPECT_TRUE(TheJIT->getGlobalValueAddress("FakeWeakSymbol") == 0);
+}
+
+// Test weak symbol linking when the weak symbol is not present in a
+// shared library
+TEST_F(MCJITTest, weak_symbol_absent) {
+  SKIP_UNSUPPORTED_PLATFORM;
+
+  SMDiagnostic Error;
+  createJITFromAssembly(
+  "  $FakeWeakSymbol2 = comdat any\n"
+  "  @FakeWeakSymbol2 = linkonce_odr global i32 42, comdat, align 4\n"
+  "  define i32* @get_weak() {\n"
+  "    ret i32* @FakeWeakSymbol2\n"
+  "  }\n");
+  void*(*FuncPtr)() =
+    (void*(*)(void))TheJIT->getFunctionAddress("get_weak");
+  EXPECT_EQ(FuncPtr(),(void*)TheJIT->getGlobalValueAddress("FakeWeakSymbol2"));
+}
+
+
 }
diff --git a/llvm/unittests/ExecutionEngine/MCJIT/MCJITTestBase.h b/llvm/unittests/ExecutionEngine/MCJIT/MCJITTestBase.h
index 35af417bf14..b1943a763d6 100644
--- a/llvm/unittests/ExecutionEngine/MCJIT/MCJITTestBase.h
+++ b/llvm/unittests/ExecutionEngine/MCJIT/MCJITTestBase.h
@@ -18,6 +18,7 @@
 #define LLVM_UNITTESTS_EXECUTIONENGINE_MCJIT_MCJITTESTBASE_H
 
 #include "MCJITTestAPICommon.h"
+#include "llvm/AsmParser/Parser.h"
 #include "llvm/Config/config.h"
 #include "llvm/ExecutionEngine/ExecutionEngine.h"
 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
@@ -27,6 +28,8 @@
 #include "llvm/IR/Module.h"
 #include "llvm/IR/TypeBuilder.h"
 #include "llvm/Support/CodeGen.h"
+#include "llvm/Support/SourceMgr.h"
+#include "llvm/Support/raw_ostream.h"
 
 namespace llvm {
 
@@ -338,6 +341,21 @@ protected:
     assert(TheJIT.get() != NULL && "error creating MCJIT with EngineBuilder");
   }
 
+  void createJITFromAssembly(const char *Test) {
+    SMDiagnostic Error;
+    M = parseAssemblyString(Test, Error, Context);
+
+    std::string errMsg;
+    raw_string_ostream os(errMsg);
+    Error.print("", os);
+
+    // A failure here means that the test itself is buggy.
+    if (!M)
+      report_fatal_error(os.str().c_str());
+
+    createJIT(std::move(M));
+  }
+
   CodeGenOpt::Level OptLevel;
   Reloc::Model RelocModel;
   CodeModel::Model CodeModel;
diff --git a/llvm/unittests/ExecutionEngine/MCJIT/Makefile b/llvm/unittests/ExecutionEngine/MCJIT/Makefile
index 2822b20cdda..522ef750df7 100644
--- a/llvm/unittests/ExecutionEngine/MCJIT/Makefile
+++ b/llvm/unittests/ExecutionEngine/MCJIT/Makefile
@@ -9,7 +9,7 @@
 
 LEVEL = ../../..
 TESTNAME = MCJIT
-LINK_COMPONENTS := core ipo mcjit native support
+LINK_COMPONENTS := core asmparser ipo mcjit native support
 
 include $(LEVEL)/Makefile.config
 include $(LLVM_SRC_ROOT)/unittests/Makefile.unittest