[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

8 files changed, 277 insertions, 57 deletions

diff --git a/llvm/include/llvm/ExecutionEngine/Orc/Core.h b/llvm/include/llvm/ExecutionEngine/Orc/Core.h
index 9a4c4455fd4..1cd2f299c7a 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/Core.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/Core.h
@@ -411,26 +411,6 @@ reexports(JITDylib &SourceJD, SymbolAliasMap Aliases,
 Expected<SymbolAliasMap>
 buildSimpleReexportsAliasMap(JITDylib &SourceJD, const SymbolNameSet &Symbols);
 
-/// ReexportsGenerator can be used with JITDylib::setGenerator to automatically
-/// re-export a subset of the source JITDylib's symbols in the target.
-class ReexportsGenerator {
-public:
-  using SymbolPredicate = std::function<bool(SymbolStringPtr)>;
-
-  /// Create a reexports generator. If an Allow predicate is passed, only
-  /// symbols for which the predicate returns true will be reexported. If no
-  /// Allow predicate is passed, all symbols will be exported.
-  ReexportsGenerator(JITDylib &SourceJD, bool MatchNonExported = false,
-                     SymbolPredicate Allow = SymbolPredicate());
-
-  Expected<SymbolNameSet> operator()(JITDylib &JD, const SymbolNameSet &Names);
-
-private:
-  JITDylib &SourceJD;
-  bool MatchNonExported = false;
-  SymbolPredicate Allow;
-};
-
 /// Represents the state that a symbol has reached during materialization.
 enum class SymbolState : uint8_t {
   Invalid,       /// No symbol should be in this state.
@@ -502,8 +482,12 @@ class JITDylib {
   friend class ExecutionSession;
   friend class MaterializationResponsibility;
 public:
-  using GeneratorFunction = std::function<Expected<SymbolNameSet>(
-      JITDylib &Parent, const SymbolNameSet &Names)>;
+  class DefinitionGenerator {
+  public:
+    virtual ~DefinitionGenerator();
+    virtual Expected<SymbolNameSet>
+    tryToGenerate(JITDylib &Parent, const SymbolNameSet &Names) = 0;
+  };
 
   using AsynchronousSymbolQuerySet =
     std::set<std::shared_ptr<AsynchronousSymbolQuery>>;
@@ -519,13 +503,20 @@ public:
   /// Get a reference to the ExecutionSession for this JITDylib.
   ExecutionSession &getExecutionSession() const { return ES; }
 
-  /// Set a definition generator. If set, whenever a symbol fails to resolve
-  /// within this JITDylib, lookup and lookupFlags will pass the unresolved
-  /// symbols set to the definition generator. The generator can optionally
-  /// add a definition for the unresolved symbols to the dylib.
-  void setGenerator(GeneratorFunction DefGenerator) {
-    this->DefGenerator = std::move(DefGenerator);
-  }
+  /// Adds a definition generator to this JITDylib and returns a referenece to
+  /// it.
+  ///
+  /// When JITDylibs are searched during lookup, if no existing definition of
+  /// a symbol is found, then any generators that have been added are run (in
+  /// the order that they were added) to potentially generate a definition.
+  template <typename GeneratorT>
+  GeneratorT &addGenerator(std::unique_ptr<GeneratorT> DefGenerator);
+
+  /// Remove a definition generator from this JITDylib.
+  ///
+  /// The given generator must exist in this JITDylib's generators list (i.e.
+  /// have been added and not yet removed).
+  void removeGenerator(DefinitionGenerator &G);
 
   /// Set the search order to be used when fixing up definitions in JITDylib.
   /// This will replace the previous search order, and apply to any symbol
@@ -744,7 +735,7 @@ private:
   SymbolTable Symbols;
   UnmaterializedInfosMap UnmaterializedInfos;
   MaterializingInfosMap MaterializingInfos;
-  GeneratorFunction DefGenerator;
+  std::vector<std::unique_ptr<DefinitionGenerator>> DefGenerators;
   JITDylibSearchList SearchOrder;
 };
 
@@ -932,6 +923,14 @@ private:
       OutstandingMUs;
 };
 
+template <typename GeneratorT>
+GeneratorT &JITDylib::addGenerator(std::unique_ptr<GeneratorT> DefGenerator) {
+  auto &G = *DefGenerator;
+  ES.runSessionLocked(
+      [&]() { DefGenerators.push_back(std::move(DefGenerator)); });
+  return G;
+}
+
 template <typename Func>
 auto JITDylib::withSearchOrderDo(Func &&F)
     -> decltype(F(std::declval<const JITDylibSearchList &>())) {
@@ -971,6 +970,27 @@ Error JITDylib::define(std::unique_ptr<MaterializationUnitType> &MU) {
   });
 }
 
+/// ReexportsGenerator can be used with JITDylib::setGenerator to automatically
+/// re-export a subset of the source JITDylib's symbols in the target.
+class ReexportsGenerator : public JITDylib::DefinitionGenerator {
+public:
+  using SymbolPredicate = std::function<bool(SymbolStringPtr)>;
+
+  /// Create a reexports generator. If an Allow predicate is passed, only
+  /// symbols for which the predicate returns true will be reexported. If no
+  /// Allow predicate is passed, all symbols will be exported.
+  ReexportsGenerator(JITDylib &SourceJD, bool MatchNonExported = false,
+                     SymbolPredicate Allow = SymbolPredicate());
+
+  Expected<SymbolNameSet> tryToGenerate(JITDylib &JD,
+                                        const SymbolNameSet &Names) override;
+
+private:
+  JITDylib &SourceJD;
+  bool MatchNonExported = false;
+  SymbolPredicate Allow;
+};
+
 /// Mangles symbol names then uniques them in the context of an
 /// ExecutionSession.
 class MangleAndInterner {
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/ExecutionUtils.h b/llvm/include/llvm/ExecutionEngine/Orc/ExecutionUtils.h
index 75865920c74..cf0a428662e 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/ExecutionUtils.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/ExecutionUtils.h
@@ -19,6 +19,7 @@
 #include "llvm/ExecutionEngine/Orc/Core.h"
 #include "llvm/ExecutionEngine/Orc/OrcError.h"
 #include "llvm/ExecutionEngine/RuntimeDyld.h"
+#include "llvm/Object/Archive.h"
 #include "llvm/Support/DynamicLibrary.h"
 #include <algorithm>
 #include <cstdint>
@@ -37,6 +38,8 @@ class Value;
 
 namespace orc {
 
+class ObjectLayer;
+
 /// This iterator provides a convenient way to iterate over the elements
 ///        of an llvm.global_ctors/llvm.global_dtors instance.
 ///
@@ -237,7 +240,7 @@ public:
 /// If an instance of this class is attached to a JITDylib as a fallback
 /// definition generator, then any symbol found in the given DynamicLibrary that
 /// passes the 'Allow' predicate will be added to the JITDylib.
-class DynamicLibrarySearchGenerator {
+class DynamicLibrarySearchGenerator : public JITDylib::DefinitionGenerator {
 public:
   using SymbolPredicate = std::function<bool(SymbolStringPtr)>;
 
@@ -253,19 +256,20 @@ public:
   /// Permanently loads the library at the given path and, on success, returns
   /// a DynamicLibrarySearchGenerator that will search it for symbol definitions
   /// in the library. On failure returns the reason the library failed to load.
-  static Expected<DynamicLibrarySearchGenerator>
+  static Expected<std::unique_ptr<DynamicLibrarySearchGenerator>>
   Load(const char *FileName, char GlobalPrefix,
        SymbolPredicate Allow = SymbolPredicate());
 
   /// Creates a DynamicLibrarySearchGenerator that searches for symbols in
   /// the current process.
-  static Expected<DynamicLibrarySearchGenerator>
+  static Expected<std::unique_ptr<DynamicLibrarySearchGenerator>>
   GetForCurrentProcess(char GlobalPrefix,
                        SymbolPredicate Allow = SymbolPredicate()) {
     return Load(nullptr, GlobalPrefix, std::move(Allow));
   }
 
-  Expected<SymbolNameSet> operator()(JITDylib &JD, const SymbolNameSet &Names);
+  Expected<SymbolNameSet> tryToGenerate(JITDylib &JD,
+                                        const SymbolNameSet &Names) override;
 
 private:
   sys::DynamicLibrary Dylib;
@@ -273,6 +277,40 @@ private:
   char GlobalPrefix;
 };
 
+/// A utility class to expose symbols from a static library.
+///
+/// If an instance of this class is attached to a JITDylib as a fallback
+/// definition generator, then any symbol found in the archive will result in
+/// the containing object being added to the JITDylib.
+class StaticLibraryDefinitionGenerator : public JITDylib::DefinitionGenerator {
+public:
+  /// Try to create a StaticLibraryDefinitionGenerator from the given path.
+  ///
+  /// This call will succeed if the file at the given path is a static library
+  /// is a valid archive, otherwise it will return an error.
+  static Expected<std::unique_ptr<StaticLibraryDefinitionGenerator>>
+  Load(ObjectLayer &L, const char *FileName);
+
+  /// Try to create a StaticLibrarySearchGenerator from the given memory buffer.
+  /// Thhis call will succeed if the buffer contains a valid archive, otherwise
+  /// it will return an error.
+  static Expected<std::unique_ptr<StaticLibraryDefinitionGenerator>>
+  Create(ObjectLayer &L, std::unique_ptr<MemoryBuffer> ArchiveBuffer);
+
+  Expected<SymbolNameSet> tryToGenerate(JITDylib &JD,
+                                        const SymbolNameSet &Names) override;
+
+private:
+  StaticLibraryDefinitionGenerator(ObjectLayer &L,
+                                   std::unique_ptr<MemoryBuffer> ArchiveBuffer,
+                                   Error &Err);
+
+  ObjectLayer &L;
+  std::unique_ptr<MemoryBuffer> ArchiveBuffer;
+  object::Archive Archive;
+  size_t UnrealizedObjects = 0;
+};
+
 } // end namespace orc
 } // end namespace llvm
 
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.
diff --git a/llvm/test/ExecutionEngine/OrcLazy/static-library-support.ll b/llvm/test/ExecutionEngine/OrcLazy/static-library-support.ll
new file mode 100644
index 00000000000..304160c7f78
--- /dev/null
+++ b/llvm/test/ExecutionEngine/OrcLazy/static-library-support.ll
@@ -0,0 +1,11 @@
+; This first line will generate the .o files for the next run line
+; RUN: llc -filetype=obj -o %t.o %p/Inputs/basic-object-source.ll
+; RUN: llvm-ar r %t.a %t.o
+; RUN: lli -jit-kind=orc-lazy -extra-archive %t.a %s
+
+declare i32 @foo()
+
+define i32 @main() {
+  %r = call i32 @foo( )   ; <i32> [#uses=1]
+  ret i32 %r
+}
diff --git a/llvm/tools/lli/lli.cpp b/llvm/tools/lli/lli.cpp
index fc76a6c5ad2..84b0f6c08c9 100644
--- a/llvm/tools/lli/lli.cpp
+++ b/llvm/tools/lli/lli.cpp
@@ -792,7 +792,7 @@ int runOrcLazyJIT(const char *ProgName) {
     });
     return TSM;
   });
-  J->getMainJITDylib().setGenerator(
+  J->getMainJITDylib().addGenerator(
       ExitOnErr(orc::DynamicLibrarySearchGenerator::GetForCurrentProcess(
           J->getDataLayout().getGlobalPrefix())));
 
@@ -832,6 +832,16 @@ int runOrcLazyJIT(const char *ProgName) {
       ExitOnErr(
           J->addLazyIRModule(JD, orc::ThreadSafeModule(std::move(M), TSCtx)));
     }
+
+    for (auto EAItr = ExtraArchives.begin(), EAEnd = ExtraArchives.end();
+         EAItr != EAEnd; ++EAItr) {
+      auto EAIdx = ExtraArchives.getPosition(EAItr - ExtraArchives.begin());
+      assert(EAIdx != 0 && "ExtraArchive should have index > 0");
+      auto JDItr = std::prev(IdxToDylib.lower_bound(EAIdx));
+      auto &JD = *JDItr->second;
+      JD.addGenerator(ExitOnErr(orc::StaticLibraryDefinitionGenerator::Load(
+          J->getObjLinkingLayer(), EAItr->c_str())));
+    }
   }
 
   // Add the objects.
diff --git a/llvm/tools/llvm-jitlink/llvm-jitlink.cpp b/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
index eff2ec18bc5..79c185a926e 100644
--- a/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
+++ b/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
@@ -380,7 +380,7 @@ Error loadProcessSymbols(Session &S) {
   auto FilterMainEntryPoint = [InternedEntryPointName](SymbolStringPtr Name) {
     return Name != InternedEntryPointName;
   };
-  S.ES.getMainJITDylib().setGenerator(
+  S.ES.getMainJITDylib().addGenerator(
       ExitOnErr(orc::DynamicLibrarySearchGenerator::GetForCurrentProcess(
           GlobalPrefix, FilterMainEntryPoint)));
 
diff --git a/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp b/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
index adc4f52b8e4..375318c323f 100644
--- a/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
@@ -243,14 +243,15 @@ TEST_F(CoreAPIsStandardTest, LookupFlagsTest) {
 
 TEST_F(CoreAPIsStandardTest, LookupWithGeneratorFailure) {
 
-  class BadGenerator {
+  class BadGenerator : public JITDylib::DefinitionGenerator {
   public:
-    Expected<SymbolNameSet> operator()(JITDylib &, const SymbolNameSet &) {
+    Expected<SymbolNameSet> tryToGenerate(JITDylib &,
+                                          const SymbolNameSet &) override {
       return make_error<StringError>("BadGenerator", inconvertibleErrorCode());
     }
   };
 
-  JD.setGenerator(BadGenerator());
+  JD.addGenerator(llvm::make_unique<BadGenerator>());
 
   EXPECT_THAT_ERROR(JD.lookupFlags({Foo}).takeError(), Failed<StringError>())
       << "Generator failure did not propagate through lookupFlags";
@@ -343,7 +344,7 @@ TEST_F(CoreAPIsStandardTest, TestReexportsGenerator) {
 
   auto Filter = [this](SymbolStringPtr Name) { return Name != Bar; };
 
-  JD.setGenerator(ReexportsGenerator(JD2, false, Filter));
+  JD.addGenerator(llvm::make_unique<ReexportsGenerator>(JD2, false, Filter));
 
   auto Flags = cantFail(JD.lookupFlags({Foo, Bar, Baz}));
   EXPECT_EQ(Flags.size(), 1U) << "Unexpected number of results";
@@ -667,10 +668,29 @@ TEST_F(CoreAPIsStandardTest, DefineMaterializingSymbol) {
 TEST_F(CoreAPIsStandardTest, GeneratorTest) {
   cantFail(JD.define(absoluteSymbols({{Foo, FooSym}})));
 
-  JD.setGenerator([&](JITDylib &JD2, const SymbolNameSet &Names) {
-    cantFail(JD2.define(absoluteSymbols({{Bar, BarSym}})));
-    return SymbolNameSet({Bar});
-  });
+  class TestGenerator : public JITDylib::DefinitionGenerator {
+  public:
+    TestGenerator(SymbolMap Symbols) : Symbols(std::move(Symbols)) {}
+    Expected<SymbolNameSet> tryToGenerate(JITDylib &JD,
+                                          const SymbolNameSet &Names) {
+      SymbolMap NewDefs;
+      SymbolNameSet NewNames;
+
+      for (auto &Name : Names) {
+        if (Symbols.count(Name)) {
+          NewDefs[Name] = Symbols[Name];
+          NewNames.insert(Name);
+        }
+      }
+      cantFail(JD.define(absoluteSymbols(std::move(NewDefs))));
+      return NewNames;
+    };
+
+  private:
+    SymbolMap Symbols;
+  };
+
+  JD.addGenerator(llvm::make_unique<TestGenerator>(SymbolMap({{Bar, BarSym}})));
 
   auto Result =
       cantFail(ES.lookup(JITDylibSearchList({{&JD, false}}), {Foo, Bar}));