[ORC] Refactor definition-generation, add a generator for static libraries.

This patch replaces the JITDylib::DefinitionGenerator typedef with a class of
the same name, and adds support for attaching a sequence of DefinitionGeneration
objects to a JITDylib.

This patch also adds a new definition generator,
StaticLibraryDefinitionGenerator, that can be used to add symbols fom a static
library to a JITDylib. An object from the static library will be added (via
a supplied ObjectLayer reference) whenever a symbol from that object is
referenced.

To enable testing, lli is updated to add support for the --extra-archive option
when running in -jit-kind=orc-lazy mode.

llvm-svn: 368707

2 files changed, 134 insertions, 13 deletions

diff --git a/llvm/lib/ExecutionEngine/Orc/Core.cpp b/llvm/lib/ExecutionEngine/Orc/Core.cpp
index 8094d68047b..23dac7a1206 100644
--- a/llvm/lib/ExecutionEngine/Orc/Core.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/Core.cpp
@@ -694,7 +694,7 @@ ReexportsGenerator::ReexportsGenerator(JITDylib &SourceJD,
       Allow(std::move(Allow)) {}
 
 Expected<SymbolNameSet>
-ReexportsGenerator::operator()(JITDylib &JD, const SymbolNameSet &Names) {
+ReexportsGenerator::tryToGenerate(JITDylib &JD, const SymbolNameSet &Names) {
   orc::SymbolNameSet Added;
   orc::SymbolAliasMap AliasMap;
 
@@ -716,6 +716,19 @@ ReexportsGenerator::operator()(JITDylib &JD, const SymbolNameSet &Names) {
   return Added;
 }
 
+JITDylib::DefinitionGenerator::~DefinitionGenerator() {}
+
+void JITDylib::removeGenerator(DefinitionGenerator &G) {
+  ES.runSessionLocked([&]() {
+    auto I = std::find_if(DefGenerators.begin(), DefGenerators.end(),
+                          [&](const std::unique_ptr<DefinitionGenerator> &H) {
+                            return H.get() == &G;
+                          });
+    assert(I != DefGenerators.end() && "Generator not found");
+    DefGenerators.erase(I);
+  });
+}
+
 Error JITDylib::defineMaterializing(const SymbolFlagsMap &SymbolFlags) {
   return ES.runSessionLocked([&]() -> Error {
     std::vector<SymbolTable::iterator> AddedSyms;
@@ -1159,10 +1172,18 @@ Expected<SymbolFlagsMap> JITDylib::lookupFlags(const SymbolNameSet &Names) {
     if (!Unresolved)
       return Unresolved.takeError();
 
-    if (DefGenerator && !Unresolved->empty()) {
-      auto NewDefs = DefGenerator(*this, *Unresolved);
+    /// Run any definition generators.
+    for (auto &DG : DefGenerators) {
+
+      // Bail out early if we've resolved everything.
+      if (Unresolved->empty())
+        break;
+
+      // Run this generator.
+      auto NewDefs = DG->tryToGenerate(*this, *Unresolved);
       if (!NewDefs)
         return NewDefs.takeError();
+
       if (!NewDefs->empty()) {
         auto Unresolved2 = lookupFlagsImpl(Result, *NewDefs);
         if (!Unresolved2)
@@ -1171,7 +1192,10 @@ Expected<SymbolFlagsMap> JITDylib::lookupFlags(const SymbolNameSet &Names) {
         assert(Unresolved2->empty() &&
                "All fallback defs should have been found by lookupFlagsImpl");
       }
-    };
+
+      for (auto &Name : *NewDefs)
+        Unresolved->erase(Name);
+    }
     return Result;
   });
 }
@@ -1198,13 +1222,25 @@ Error JITDylib::lodgeQuery(std::shared_ptr<AsynchronousSymbolQuery> &Q,
   assert(Q && "Query can not be null");
 
   lodgeQueryImpl(Q, Unresolved, MatchNonExported, MUs);
-  if (DefGenerator && !Unresolved.empty()) {
-    auto NewDefs = DefGenerator(*this, Unresolved);
+
+  // Run any definition generators.
+  for (auto &DG : DefGenerators) {
+
+    // Bail out early if we have resolved everything.
+    if (Unresolved.empty())
+      break;
+
+    // Run the generator.
+    auto NewDefs = DG->tryToGenerate(*this, Unresolved);
+
     if (!NewDefs)
       return NewDefs.takeError();
+
+    llvm::dbgs() << "NewDefs is " << *NewDefs << "\n";
     if (!NewDefs->empty()) {
       for (auto &D : *NewDefs)
         Unresolved.erase(D);
+      llvm::dbgs() << "NewDefs is now " << *NewDefs << "\n";
       lodgeQueryImpl(Q, *NewDefs, MatchNonExported, MUs);
       assert(NewDefs->empty() &&
              "All fallback defs should have been found by lookupImpl");
@@ -1292,9 +1328,16 @@ JITDylib::legacyLookup(std::shared_ptr<AsynchronousSymbolQuery> Q,
   SymbolNameSet Unresolved = std::move(Names);
   auto Err = ES.runSessionLocked([&, this]() -> Error {
     QueryComplete = lookupImpl(Q, MUs, Unresolved);
-    if (DefGenerator && !Unresolved.empty()) {
+
+    // Run any definition generators.
+    for (auto &DG : DefGenerators) {
+
+      // Bail out early if we have resolved everything.
+      if (Unresolved.empty())
+        break;
+
       assert(!QueryComplete && "query complete but unresolved symbols remain?");
-      auto NewDefs = DefGenerator(*this, Unresolved);
+      auto NewDefs = DG->tryToGenerate(*this, Unresolved);
       if (!NewDefs)
         return NewDefs.takeError();
       if (!NewDefs->empty()) {
diff --git a/llvm/lib/ExecutionEngine/Orc/ExecutionUtils.cpp b/llvm/lib/ExecutionEngine/Orc/ExecutionUtils.cpp
index f7fc5f8f179..044e7eede08 100644
--- a/llvm/lib/ExecutionEngine/Orc/ExecutionUtils.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/ExecutionUtils.cpp
@@ -8,6 +8,7 @@
 
 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
 
+#include "llvm/ExecutionEngine/Orc/Layer.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/GlobalVariable.h"
@@ -178,20 +179,20 @@ DynamicLibrarySearchGenerator::DynamicLibrarySearchGenerator(
     : Dylib(std::move(Dylib)), Allow(std::move(Allow)),
       GlobalPrefix(GlobalPrefix) {}
 
-Expected<DynamicLibrarySearchGenerator>
+Expected<std::unique_ptr<DynamicLibrarySearchGenerator>>
 DynamicLibrarySearchGenerator::Load(const char *FileName, char GlobalPrefix,
                                     SymbolPredicate Allow) {
   std::string ErrMsg;
   auto Lib = sys::DynamicLibrary::getPermanentLibrary(FileName, &ErrMsg);
   if (!Lib.isValid())
     return make_error<StringError>(std::move(ErrMsg), inconvertibleErrorCode());
-  return DynamicLibrarySearchGenerator(std::move(Lib), GlobalPrefix,
-                                       std::move(Allow));
+  return llvm::make_unique<DynamicLibrarySearchGenerator>(
+      std::move(Lib), GlobalPrefix, std::move(Allow));
 }
 
 Expected<SymbolNameSet>
-DynamicLibrarySearchGenerator::operator()(JITDylib &JD,
-                                          const SymbolNameSet &Names) {
+DynamicLibrarySearchGenerator::tryToGenerate(JITDylib &JD,
+                                             const SymbolNameSet &Names) {
   orc::SymbolNameSet Added;
   orc::SymbolMap NewSymbols;
 
@@ -226,5 +227,82 @@ DynamicLibrarySearchGenerator::operator()(JITDylib &JD,
   return Added;
 }
 
+Expected<std::unique_ptr<StaticLibraryDefinitionGenerator>>
+StaticLibraryDefinitionGenerator::Load(ObjectLayer &L, const char *FileName) {
+  auto ArchiveBuffer = errorOrToExpected(MemoryBuffer::getFile(FileName));
+
+  if (!ArchiveBuffer)
+    return ArchiveBuffer.takeError();
+
+  return Create(L, std::move(*ArchiveBuffer));
+}
+
+Expected<std::unique_ptr<StaticLibraryDefinitionGenerator>>
+StaticLibraryDefinitionGenerator::Create(
+    ObjectLayer &L, std::unique_ptr<MemoryBuffer> ArchiveBuffer) {
+  Error Err = Error::success();
+
+  std::unique_ptr<StaticLibraryDefinitionGenerator> ADG(
+      new StaticLibraryDefinitionGenerator(L, std::move(ArchiveBuffer), Err));
+
+  if (Err)
+    return std::move(Err);
+
+  return std::move(ADG);
+}
+
+Expected<SymbolNameSet>
+StaticLibraryDefinitionGenerator::tryToGenerate(JITDylib &JD,
+                                                const SymbolNameSet &Names) {
+
+  DenseSet<std::pair<StringRef, StringRef>> ChildBufferInfos;
+  SymbolNameSet NewDefs;
+
+  for (const auto &Name : Names) {
+    auto Child = Archive.findSym(*Name);
+    if (!Child)
+      return Child.takeError();
+    if (*Child == None)
+      continue;
+    auto ChildBuffer = (*Child)->getMemoryBufferRef();
+    if (!ChildBuffer)
+      return ChildBuffer.takeError();
+    ChildBufferInfos.insert(
+        {ChildBuffer->getBuffer(), ChildBuffer->getBufferIdentifier()});
+    NewDefs.insert(Name);
+  }
+
+  for (auto ChildBufferInfo : ChildBufferInfos) {
+    MemoryBufferRef ChildBufferRef(ChildBufferInfo.first,
+                                   ChildBufferInfo.second);
+
+    if (auto Err =
+            L.add(JD, MemoryBuffer::getMemBuffer(ChildBufferRef), VModuleKey()))
+      return std::move(Err);
+
+    --UnrealizedObjects;
+  }
+
+  return NewDefs;
+}
+
+StaticLibraryDefinitionGenerator::StaticLibraryDefinitionGenerator(
+    ObjectLayer &L, std::unique_ptr<MemoryBuffer> ArchiveBuffer, Error &Err)
+    : L(L), ArchiveBuffer(std::move(ArchiveBuffer)),
+      Archive(*this->ArchiveBuffer, Err) {
+
+  if (Err)
+    return;
+
+  Error Err2 = Error::success();
+  for (auto _ : Archive.children(Err2)) {
+    (void)_;
+    ++UnrealizedObjects;
+  }
+
+  // No need to check this: We will leave it to the caller.
+  Err = std::move(Err2);
+}
+
 } // End namespace orc.
 } // End namespace llvm.