Add weak symbol support to RuntimeDyld

[loladiro]

Summary

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. I'd like some feedback if these semantics seem correct, since the semantics of weak symbols during dynamic linking are always tricky.

[lhames]

This looks good to me. Thanks Keno!

lib/ExecutionEngine/RuntimeDyld/RuntimeDyld.cpp

for (symbol_iterator I = Obj.symbol_begin(), E = Obj.symbol_end(); I != E;
     object::SymbolRef::Type SymType;
     StringRef Name;
     Check(I->getType(SymType));
     Check(I->getName(Name));
 
     uint32_t Flags = I->getFlags();
 
     bool IsCommon = Flags & SymbolRef::SF_Common;
+    bool IsWeak = Flags & SymbolRef::SF_Weak;
     if (IsCommon) {
       // Add the common symbols to a list.  We'll allocate them all below.
       if (!GlobalSymbolTable.count(Name)) {
         uint32_t Align;
         Check(I->getAlignment(Align));
         uint64_t Size = 0;
         Check(I->getSize(Size));
         CommonSize += Size + Align;
         CommonSymbols[*I] = CommonSymbolInfo(Size, Align);
       }
     } else {
       if (SymType == object::SymbolRef::ST_Function ||
           SymType == object::SymbolRef::ST_Data ||
           SymType == object::SymbolRef::ST_Unknown) {
         uint64_t SectOffset;
         StringRef SectionData;
         section_iterator SI = Obj.section_end();
         Check(getOffset(*I, SectOffset));
         Check(I->getSection(SI));
         if (SI == Obj.section_end())
           continue;
         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);
         DEBUG(dbgs() << "\tOffset: " << format("%p", (uintptr_t)SectOffset)
                      << " flags: " << Flags << " SID: " << SectionID);
+        if (!IsWeak)
+        {
           GlobalSymbolTable[Name] = SymbolLoc(SectionID, SectOffset);
+        } else {
+          WeakSymbolTable[Name] = SymbolLoc(SectionID,SectOffset);
+        }
       }
     }
     DEBUG(dbgs() << "\tType: " << SymType << " Name: " << Name << "\n");
   }
 
   // Allocate common symbols
   if (CommonSize != 0)
     emitCommonSymbols(Obj, CommonSymbols, CommonSize, GlobalSymbolTable);
if (Name.size() == 0) {
         i = ExternalSymbolRelocations.find(Name);
       } else {
         // We found the symbol in our global table.  It was probably in a
         // Module that we loaded previously.
         SymbolLoc SymLoc = Loc->second;
         Addr = getSectionLoadAddress(SymLoc.first) + SymLoc.second;
       }
 
+      // 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) {
+        SymbolTableMap::const_iterator Loc = WeakSymbolTable.find(Name);
+        if (Loc != WeakSymbolTable.end()) {
+          SymbolLoc SymLoc = Loc->second;
+          Addr = getSectionLoadAddress(SymLoc.first) + SymLoc.second;
+          // 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] = SymLoc;
+        }
+      }
+
       // FIXME: Implement error handling that doesn't kill the host program!
       if (!Addr)
         report_fatal_error("Program used external function '" + Name +
                            "' which could not be resolved!");
 
       updateGOTEntries(Name, Addr);
       DEBUG(dbgs() << "Resolving relocations Name: " << Name << "\t"
                    << format("0x%lx", Addr) << "\n");
       // This list may have been updated when we called getSymbolAddress, so
       // don't change this code to get the list earlier.
       RelocationList &Relocs = i->second;
       resolveRelocationList(Relocs, Addr);
     }
 
     ExternalSymbolRelocations.erase(i);
   }
+  WeakSymbolTable.clear();
 }
 
 //===----------------------------------------------------------------------===//
 // RuntimeDyld class implementation
 
 uint64_t RuntimeDyld::LoadedObjectInfo::getSectionLoadAddress(
                                                   StringRef SectionName) const {
   for (unsigned I = BeginIdx; I != EndIdx; ++I)

lib/ExecutionEngine/RuntimeDyld/RuntimeDyldELF.cpp

case SymbolRef::ST_Debug: {
         llvm_unreachable("Symbol section not found, bad object file format!");
       DEBUG(dbgs() << "\t\tThis is section symbol\n");
       bool isCode = si->isText();
       Value.SectionID = findOrEmitSection(Obj, (*si), isCode, ObjSectionToID);
       Value.Addend = Addend;
       break;
     }
     case SymbolRef::ST_Data:
+    case SymbolRef::ST_Function:
     case SymbolRef::ST_Unknown: {
       Value.SymbolName = TargetName.data();
       Value.Addend = Addend;
 
       // Absolute relocations will have a zero symbol ID (STN_UNDEF), which
       // will manifest here as a NULL symbol name.
       // We can set this as a valid (but empty) symbol name, and rely
       // on addRelocationForSymbol to handle this.

lib/ExecutionEngine/RuntimeDyld/RuntimeDyldImpl.h

#define RTDYLD_INVALID_SECTION_ID ((SID)(-1))
   typedef std::map<SectionRef, unsigned> ObjSectionToIDMap;
 
   // A global symbol table for symbols from all loaded modules.  Maps the
   // symbol name to a (SectionID, offset in section) pair.
   typedef std::pair<unsigned, uintptr_t> SymbolLoc;
   typedef StringMap<SymbolLoc> SymbolTableMap;
   SymbolTableMap GlobalSymbolTable;
 
+  // Like the global symbol table but for weak symbols
+  SymbolTableMap WeakSymbolTable;
+
   // Pair representing the size and alignment requirement for a common symbol.
   typedef std::pair<unsigned, unsigned> CommonSymbolInfo;
   // Keep a map of common symbols to their info pairs
   typedef std::map<SymbolRef, CommonSymbolInfo> CommonSymbolMap;
 
   // For each symbol, keep a list of relocations based on it. Anytime
   // its address is reassigned (the JIT re-compiled the function, e.g.),
   // the relocations get re-resolved.

unittests/ExecutionEngine/MCJIT/CMakeLists.txt

 set(LLVM_LINK_COMPONENTS
   Analysis
+  AsmParser
   Core
   ExecutionEngine
   IPO
   MC
   MCJIT
   ScalarOpts
   Support
   Target

unittests/ExecutionEngine/MCJIT/MCJITTest.cpp

 //===- MCJITTest.cpp - Unit tests for the MCJIT ---------------------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // This test suite verifies basic MCJIT functionality such as making function
 // calls, using global variables, and compiling multpile modules.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ExecutionEngine/MCJIT.h"
+#include "llvm/Support/DynamicLibrary.h"
 #include "MCJITTestBase.h"
 #include "gtest/gtest.h"
 
 using namespace llvm;
 
 namespace {
 
 class MCJITTest : public testing::Test, public MCJITTestBase {
TEST_F(MCJITTest, multiple_decl_lookups) {
   createJIT(std::move(M));
   void *A = TheJIT->getPointerToFunction(Foo);
   void *B = TheJIT->getPointerToFunction(Foo);
 
   EXPECT_TRUE(A != 0) << "Failed lookup - test not correctly configured.";
   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"));
+}
+
+
 }

unittests/ExecutionEngine/MCJIT/MCJITTestBase.h

 //
 //===----------------------------------------------------------------------===//
 
 
 #ifndef LLVM_UNITTESTS_EXECUTIONENGINE_MCJIT_MCJITTESTBASE_H
 #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"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/LLVMContext.h"
 #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 {
 
 /// Helper class that can build very simple Modules
 class TrivialModuleBuilder {
 protected:
   LLVMContext Context;
   IRBuilder<> Builder;
TheJIT.reset(EB.setEngineKind(EngineKind::JIT)
                  .setMArch(MArch)
                  .setMCPU(sys::getHostCPUName())
                  //.setMAttrs(MAttrs)
                  .create());
     // At this point, we cannot modify the module any more.
     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;
   StringRef MArch;
   SmallVector<std::string, 1> MAttrs;
   std::unique_ptr<ExecutionEngine> TheJIT;
   std::unique_ptr<RTDyldMemoryManager> MM;
 
   std::unique_ptr<Module> M;
 };
 
 } // namespace llvm
 
 #endif

unittests/ExecutionEngine/MCJIT/Makefile

 ##===- unittests/ExecutionEngine/MCJIT/Makefile ------------*- Makefile -*-===##
 #
 #                     The LLVM Compiler Infrastructure
 #
 # This file is distributed under the University of Illinois Open Source
 # License. See LICENSE.TXT for details.
 #
 ##===----------------------------------------------------------------------===##
 
 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
 
 # Permit these tests to use the MCJIT's symbolic lookup.
 LD.Flags += $(RDYNAMIC)