[RuntimeDyld] Decouple RuntimeDyldChecker from RuntimeDyld.

This will allow RuntimeDyldChecker (and rtdyld-check tests) to test a new JIT
linker: JITLink (https://reviews.llvm.org/D58704).

llvm-svn: 357947

1 files changed, 212 insertions, 48 deletions

diff --git a/llvm/tools/llvm-rtdyld/llvm-rtdyld.cpp b/llvm/tools/llvm-rtdyld/llvm-rtdyld.cpp
index 69d70e94808..5a5d035d27d 100644
--- a/llvm/tools/llvm-rtdyld/llvm-rtdyld.cpp
+++ b/llvm/tools/llvm-rtdyld/llvm-rtdyld.cpp
@@ -29,9 +29,12 @@
 #include "llvm/Support/InitLLVM.h"
 #include "llvm/Support/Memory.h"
 #include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/Path.h"
 #include "llvm/Support/TargetRegistry.h"
 #include "llvm/Support/TargetSelect.h"
 #include "llvm/Support/raw_ostream.h"
+
+#include <future>
 #include <list>
 
 using namespace llvm;
@@ -137,14 +140,50 @@ PrintAllocationRequests("print-alloc-requests",
                                  "manager by RuntimeDyld"),
                         cl::Hidden);
 
+ExitOnError ExitOnErr;
+
 /* *** */
 
+using SectionIDMap = StringMap<unsigned>;
+using FileToSectionIDMap = StringMap<SectionIDMap>;
+
+void dumpFileToSectionIDMap(const FileToSectionIDMap &FileToSecIDMap) {
+  for (const auto &KV : FileToSecIDMap) {
+    llvm::dbgs() << "In " << KV.first() << "\n";
+    for (auto &KV2 : KV.second)
+      llvm::dbgs() << "  \"" << KV2.first() << "\" -> " << KV2.second << "\n";
+  }
+}
+
+Expected<unsigned> getSectionId(const FileToSectionIDMap &FileToSecIDMap,
+                                StringRef FileName, StringRef SectionName) {
+  auto I = FileToSecIDMap.find(FileName);
+  if (I == FileToSecIDMap.end())
+    return make_error<StringError>("No file named " + FileName,
+                                   inconvertibleErrorCode());
+  auto &SectionIDs = I->second;
+  auto J = SectionIDs.find(SectionName);
+  if (J == SectionIDs.end())
+    return make_error<StringError>("No section named \"" + SectionName +
+                                   "\" in file " + FileName,
+                                   inconvertibleErrorCode());
+  return J->second;
+}
+
 // A trivial memory manager that doesn't do anything fancy, just uses the
 // support library allocation routines directly.
 class TrivialMemoryManager : public RTDyldMemoryManager {
 public:
-  SmallVector<sys::MemoryBlock, 16> FunctionMemory;
-  SmallVector<sys::MemoryBlock, 16> DataMemory;
+  struct SectionInfo {
+    SectionInfo(StringRef Name, sys::MemoryBlock MB, unsigned SectionID)
+      : Name(Name), MB(std::move(MB)), SectionID(SectionID) {}
+    std::string Name;
+    sys::MemoryBlock MB;
+    unsigned SectionID = ~0U;
+  };
+
+  SmallVector<SectionInfo, 16> FunctionMemory;
+  SmallVector<SectionInfo, 16> DataMemory;
 
   uint8_t *allocateCodeSection(uintptr_t Size, unsigned Alignment,
                                unsigned SectionID,
@@ -153,6 +192,11 @@ public:
                                unsigned SectionID, StringRef SectionName,
                                bool IsReadOnly) override;
 
+  /// If non null, records subsequent Name -> SectionID mappings.
+  void setSectionIDsMap(SectionIDMap *SecIDMap) {
+    this->SecIDMap = SecIDMap;
+  }
+
   void *getPointerToNamedFunction(const std::string &Name,
                                   bool AbortOnFailure = true) override {
     return nullptr;
@@ -192,7 +236,8 @@ public:
     SlabSize = Size;
   }
 
-  uint8_t *allocateFromSlab(uintptr_t Size, unsigned Alignment, bool isCode) {
+  uint8_t *allocateFromSlab(uintptr_t Size, unsigned Alignment, bool isCode,
+                            StringRef SectionName, unsigned SectionID) {
     Size = alignTo(Size, Alignment);
     if (CurrentSlabOffset + Size > SlabSize)
       report_fatal_error("Can't allocate enough memory. Tune --preallocate");
@@ -200,9 +245,9 @@ public:
     uintptr_t OldSlabOffset = CurrentSlabOffset;
     sys::MemoryBlock MB((void *)OldSlabOffset, Size);
     if (isCode)
-      FunctionMemory.push_back(MB);
+      FunctionMemory.push_back(SectionInfo(SectionName, MB, SectionID));
     else
-      DataMemory.push_back(MB);
+      DataMemory.push_back(SectionInfo(SectionName, MB, SectionID));
     CurrentSlabOffset += Size;
     return (uint8_t*)OldSlabOffset;
   }
@@ -213,6 +258,7 @@ private:
   bool UsePreallocation = false;
   uintptr_t SlabSize = 0;
   uintptr_t CurrentSlabOffset = 0;
+  SectionIDMap *SecIDMap = nullptr;
 };
 
 uint8_t *TrivialMemoryManager::allocateCodeSection(uintptr_t Size,
@@ -223,8 +269,13 @@ uint8_t *TrivialMemoryManager::allocateCodeSection(uintptr_t Size,
     outs() << "allocateCodeSection(Size = " << Size << ", Alignment = "
            << Alignment << ", SectionName = " << SectionName << ")\n";
 
+  dbgs() << "  Registering code section \"" << SectionName << "\"\n";
+  if (SecIDMap)
+    (*SecIDMap)[SectionName] = SectionID;
+
   if (UsePreallocation)
-    return allocateFromSlab(Size, Alignment, true /* isCode */);
+    return allocateFromSlab(Size, Alignment, true /* isCode */,
+                            SectionName, SectionID);
 
   std::error_code EC;
   sys::MemoryBlock MB =
@@ -234,7 +285,7 @@ uint8_t *TrivialMemoryManager::allocateCodeSection(uintptr_t Size,
                                       EC);
   if (!MB.base())
     report_fatal_error("MemoryManager allocation failed: " + EC.message());
-  FunctionMemory.push_back(MB);
+  FunctionMemory.push_back(SectionInfo(SectionName, MB, SectionID));
   return (uint8_t*)MB.base();
 }
 
@@ -247,8 +298,13 @@ uint8_t *TrivialMemoryManager::allocateDataSection(uintptr_t Size,
     outs() << "allocateDataSection(Size = " << Size << ", Alignment = "
            << Alignment << ", SectionName = " << SectionName << ")\n";
 
+  dbgs() << "  Registering code section \"" << SectionName << "\"\n";
+  if (SecIDMap)
+    (*SecIDMap)[SectionName] = SectionID;
+
   if (UsePreallocation)
-    return allocateFromSlab(Size, Alignment, false /* isCode */);
+    return allocateFromSlab(Size, Alignment, false /* isCode */, SectionName,
+                            SectionID);
 
   std::error_code EC;
   sys::MemoryBlock MB =
@@ -258,7 +314,7 @@ uint8_t *TrivialMemoryManager::allocateDataSection(uintptr_t Size,
                                       EC);
   if (!MB.base())
     report_fatal_error("MemoryManager allocation failed: " + EC.message());
-  DataMemory.push_back(MB);
+  DataMemory.push_back(SectionInfo(SectionName, MB, SectionID));
   return (uint8_t*)MB.base();
 }
 
@@ -469,9 +525,11 @@ static int executeInput() {
   // Invalidate the instruction cache for each loaded function.
   for (auto &FM : MemMgr.FunctionMemory) {
 
+    auto &FM_MB = FM.MB;
+
     // Make sure the memory is executable.
     // setExecutable will call InvalidateInstructionCache.
-    if (auto EC = sys::Memory::protectMappedMemory(FM,
+    if (auto EC = sys::Memory::protectMappedMemory(FM_MB,
                                                    sys::Memory::MF_READ |
                                                    sys::Memory::MF_EXEC))
       ErrorAndExit("unable to mark function executable: '" + EC.message() +
@@ -505,10 +563,10 @@ static int checkAllExpressions(RuntimeDyldChecker &Checker) {
   return 0;
 }
 
-void applySpecificSectionMappings(RuntimeDyldChecker &Checker) {
+void applySpecificSectionMappings(RuntimeDyld &Dyld,
+                                  const FileToSectionIDMap &FileToSecIDMap) {
 
   for (StringRef Mapping : SpecificSectionMappings) {
-
     size_t EqualsIdx = Mapping.find_first_of("=");
     std::string SectionIDStr = Mapping.substr(0, EqualsIdx);
     size_t ComaIdx = Mapping.find_first_of(",");
@@ -519,17 +577,10 @@ void applySpecificSectionMappings(RuntimeDyldChecker &Checker) {
 
     std::string FileName = SectionIDStr.substr(0, ComaIdx);
     std::string SectionName = SectionIDStr.substr(ComaIdx + 1);
+    unsigned SectionID =
+      ExitOnErr(getSectionId(FileToSecIDMap, FileName, SectionName));
 
-    uint64_t OldAddrInt;
-    std::string ErrorMsg;
-    std::tie(OldAddrInt, ErrorMsg) =
-      Checker.getSectionAddr(FileName, SectionName, true);
-
-    if (ErrorMsg != "")
-      report_fatal_error(ErrorMsg);
-
-    void* OldAddr = reinterpret_cast<void*>(static_cast<uintptr_t>(OldAddrInt));
-
+    auto* OldAddr = Dyld.getSectionContent(SectionID).data();
     std::string NewAddrStr = Mapping.substr(EqualsIdx + 1);
     uint64_t NewAddr;
 
@@ -537,7 +588,7 @@ void applySpecificSectionMappings(RuntimeDyldChecker &Checker) {
       report_fatal_error("Invalid section address in mapping '" + Mapping +
                          "'.");
 
-    Checker.getRTDyld().mapSectionAddress(OldAddr, NewAddr);
+    Dyld.mapSectionAddress(OldAddr, NewAddr);
   }
 }
 
@@ -553,21 +604,17 @@ void applySpecificSectionMappings(RuntimeDyldChecker &Checker) {
 //                            (e.g. 1 << 32) to stress-test stubs, GOTs, etc.
 //
 static void remapSectionsAndSymbols(const llvm::Triple &TargetTriple,
-                                    TrivialMemoryManager &MemMgr,
-                                    RuntimeDyldChecker &Checker) {
+                                    RuntimeDyld &Dyld,
+                                    TrivialMemoryManager &MemMgr) {
 
   // Set up a work list (section addr/size pairs).
-  typedef std::list<std::pair<void*, uint64_t>> WorklistT;
+  typedef std::list<const TrivialMemoryManager::SectionInfo*> WorklistT;
   WorklistT Worklist;
 
   for (const auto& CodeSection : MemMgr.FunctionMemory)
-    Worklist.push_back(std::make_pair(CodeSection.base(), CodeSection.size()));
+    Worklist.push_back(&CodeSection);
   for (const auto& DataSection : MemMgr.DataMemory)
-    Worklist.push_back(std::make_pair(DataSection.base(), DataSection.size()));
-
-  // Apply any section-specific mappings that were requested on the command
-  // line.
-  applySpecificSectionMappings(Checker);
+    Worklist.push_back(&DataSection);
 
   // Keep an "already allocated" mapping of section target addresses to sizes.
   // Sections whose address mappings aren't specified on the command line will
@@ -582,16 +629,16 @@ static void remapSectionsAndSymbols(const llvm::Triple &TargetTriple,
        I != E;) {
     WorklistT::iterator Tmp = I;
     ++I;
-    auto LoadAddr = Checker.getSectionLoadAddress(Tmp->first);
 
-    if (LoadAddr &&
-        *LoadAddr != static_cast<uint64_t>(
-                       reinterpret_cast<uintptr_t>(Tmp->first))) {
+    auto LoadAddr = Dyld.getSectionLoadAddress((*Tmp)->SectionID);
+
+    if (LoadAddr != static_cast<uint64_t>(
+          reinterpret_cast<uintptr_t>((*Tmp)->MB.base()))) {
       // A section will have a LoadAddr of 0 if it wasn't loaded for whatever
       // reason (e.g. zero byte COFF sections). Don't include those sections in
       // the allocation map.
-      if (*LoadAddr != 0)
-        AlreadyAllocated[*LoadAddr] = Tmp->second;
+      if (LoadAddr != 0)
+        AlreadyAllocated[LoadAddr] = (*Tmp)->MB.size();
       Worklist.erase(Tmp);
     }
   }
@@ -609,19 +656,19 @@ static void remapSectionsAndSymbols(const llvm::Triple &TargetTriple,
 
   // Process any elements remaining in the worklist.
   while (!Worklist.empty()) {
-    std::pair<void*, uint64_t> CurEntry = Worklist.front();
+    auto *CurEntry = Worklist.front();
     Worklist.pop_front();
 
     uint64_t NextSectionAddr = TargetAddrStart;
 
     for (const auto &Alloc : AlreadyAllocated)
-      if (NextSectionAddr + CurEntry.second + TargetSectionSep <= Alloc.first)
+      if (NextSectionAddr + CurEntry->MB.size() + TargetSectionSep <= Alloc.first)
         break;
       else
         NextSectionAddr = Alloc.first + Alloc.second + TargetSectionSep;
 
-    AlreadyAllocated[NextSectionAddr] = CurEntry.second;
-    Checker.getRTDyld().mapSectionAddress(CurEntry.first, NextSectionAddr);
+    Dyld.mapSectionAddress(CurEntry->MB.base(), NextSectionAddr);
+    AlreadyAllocated[NextSectionAddr] = CurEntry->MB.size();
   }
 
   // Add dummy symbols to the memory manager.
@@ -693,18 +740,116 @@ static int linkAndVerify() {
   // Instantiate a dynamic linker.
   TrivialMemoryManager MemMgr;
   doPreallocation(MemMgr);
+
+  using StubOffsets = StringMap<uint32_t>;
+  using SectionStubs = StringMap<StubOffsets>;
+  using FileStubs = StringMap<SectionStubs>;
+
+  FileStubs StubMap;
   RuntimeDyld Dyld(MemMgr, MemMgr);
   Dyld.setProcessAllSections(true);
-  RuntimeDyldChecker Checker(Dyld, Disassembler.get(), InstPrinter.get(),
-                             llvm::dbgs());
+
+  Dyld.setNotifyStubEmitted(
+      [&StubMap](StringRef FilePath, StringRef SectionName,
+                 StringRef SymbolName, uint32_t StubOffset) {
+        StubMap[sys::path::filename(FilePath)][SectionName][SymbolName] =
+          StubOffset;
+      });
+
+  auto GetSymbolAddress =
+    [&Dyld, &MemMgr](StringRef Symbol) -> Expected<JITTargetAddress> {
+      if (auto InternalSymbol = Dyld.getSymbol(Symbol))
+        return InternalSymbol.getAddress();
+
+#ifdef _MSC_VER
+      using ExpectedLookupResult = MSVCPExpected<JITSymbolResolver::LookupResult>;
+#else
+      using ExpectedLookupResult = Expected<JITSymbolResolver::LookupResult>;
+#endif
+
+      auto ResultP = std::make_shared<std::promise<ExpectedLookupResult>>();
+      auto ResultF = ResultP->get_future();
+
+      MemMgr.lookup(
+        JITSymbolResolver::LookupSet({Symbol}),
+        [=](Expected<JITSymbolResolver::LookupResult> Result) {
+          ResultP->set_value(std::move(Result));
+        });
+
+      auto Result = ResultF.get();
+      if (!Result)
+        return Result.takeError();
+
+      auto I = Result->find(Symbol);
+      assert(I != Result->end() &&
+             "Expected symbol address if no error occurred");
+      return I->second.getAddress();
+    };
+
+  auto IsSymbolValid =
+    [&Dyld, GetSymbolAddress](StringRef Symbol) {
+      if (Dyld.getSymbol(Symbol))
+        return true;
+      auto Addr = GetSymbolAddress(Symbol);
+      if (!Addr) {
+        logAllUnhandledErrors(Addr.takeError(), errs(), "RTDyldChecker: ");
+        return false;
+      }
+      return *Addr != 0;
+    };
+
+  FileToSectionIDMap FileToSecIDMap;
+
+  auto GetSectionAddress =
+    [&Dyld, &FileToSecIDMap](StringRef FileName, StringRef SectionName) {
+      unsigned SectionID =
+        ExitOnErr(getSectionId(FileToSecIDMap, FileName, SectionName));
+      return Dyld.getSectionLoadAddress(SectionID);
+    };
+
+  auto GetSectionContent =
+    [&Dyld, &FileToSecIDMap](StringRef FileName, StringRef SectionName) {
+      unsigned SectionID =
+        ExitOnErr(getSectionId(FileToSecIDMap, FileName, SectionName));
+      return Dyld.getSectionContent(SectionID);
+    };
+
+
+  auto GetSymbolContents =
+    [&Dyld](StringRef Symbol) {
+      auto *SymAddr = static_cast<char*>(Dyld.getSymbolLocalAddress(Symbol));
+      if (!SymAddr)
+        return StringRef();
+      unsigned SectionID = Dyld.getSymbolSectionID(Symbol);
+      if (SectionID == ~0U)
+        return StringRef();
+      StringRef SecContent = Dyld.getSectionContent(SectionID);
+      uint64_t SymSize = SecContent.size() - (SymAddr - SecContent.data());
+      return StringRef(SymAddr, SymSize);
+    };
+
+  auto GetStubOffset =
+    [&StubMap](StringRef FileName, StringRef SectionName, StringRef SymbolName) -> Expected<uint32_t> {
+      if (!StubMap.count(FileName))
+        return make_error<StringError>("File name not found", inconvertibleErrorCode());
+      if (!StubMap[FileName].count(SectionName))
+        return make_error<StringError>("Section name not found", inconvertibleErrorCode());
+      if (!StubMap[FileName][SectionName].count(SymbolName))
+        return make_error<StringError>("Symbol name not found", inconvertibleErrorCode());
+      return StubMap[FileName][SectionName][SymbolName];
+    };
+
+  // We will initialize this below once we have the first object file and can
+  // know the endianness.
+  std::unique_ptr<RuntimeDyldChecker> Checker;
 
   // If we don't have any input files, read from stdin.
   if (!InputFileList.size())
     InputFileList.push_back("-");
-  for (auto &Filename : InputFileList) {
+  for (auto &InputFile : InputFileList) {
     // Load the input memory buffer.
     ErrorOr<std::unique_ptr<MemoryBuffer>> InputBuffer =
-        MemoryBuffer::getFileOrSTDIN(Filename);
+        MemoryBuffer::getFileOrSTDIN(InputFile);
 
     if (std::error_code EC = InputBuffer.getError())
       ErrorAndExit("unable to read input: '" + EC.message() + "'");
@@ -722,6 +867,22 @@ static int linkAndVerify() {
 
     ObjectFile &Obj = **MaybeObj;
 
+    if (!Checker)
+      Checker =
+        llvm::make_unique<RuntimeDyldChecker>(IsSymbolValid, GetSymbolAddress,
+                                              GetSymbolContents,
+                                              GetSectionAddress,
+                                              GetSectionContent, GetStubOffset,
+                                              Obj.isLittleEndian()
+                                                ? support::little
+                                                : support::big,
+                                              Disassembler.get(),
+                                              InstPrinter.get(), dbgs());
+
+    auto FileName = sys::path::filename(InputFile);
+    dbgs() << "In " << FileName << ":\n";
+    MemMgr.setSectionIDsMap(&FileToSecIDMap[FileName]);
+
     // Load the object file
     Dyld.loadObject(Obj);
     if (Dyld.hasError()) {
@@ -731,7 +892,8 @@ static int linkAndVerify() {
 
   // Re-map the section addresses into the phony target address space and add
   // dummy symbols.
-  remapSectionsAndSymbols(TheTriple, MemMgr, Checker);
+  applySpecificSectionMappings(Dyld, FileToSecIDMap);
+  remapSectionsAndSymbols(TheTriple, Dyld, MemMgr);
 
   // Resolve all the relocations we can.
   Dyld.resolveRelocations();
@@ -739,7 +901,7 @@ static int linkAndVerify() {
   // Register EH frames.
   Dyld.registerEHFrames();
 
-  int ErrorCode = checkAllExpressions(Checker);
+  int ErrorCode = checkAllExpressions(*Checker);
   if (Dyld.hasError())
     ErrorAndExit("RTDyld reported an error applying relocations:\n  " +
                  Dyld.getErrorString());
@@ -757,6 +919,8 @@ int main(int argc, char **argv) {
 
   cl::ParseCommandLineOptions(argc, argv, "llvm MC-JIT tool\n");
 
+  ExitOnErr.setBanner(std::string(argv[0]) + ": ");
+
   switch (Action) {
   case AC_Execute:
     return executeInput();