[ORC] Add a special 'main' JITDylib that is created on ExecutionSession

construction, a new convenience lookup method, and add-to layer methods.

ExecutionSession now creates a special 'main' JITDylib upon construction. All
subsequently created JITDylibs are added to the main JITDylib's search order by
default (controlled by the AddToMainDylibSearchOrder parameter to
ExecutionSession::createDylib). The main JITDylib's search order will be used in
the future to properly handle cross-JITDylib weak symbols, with the first
definition in this search order selected.

This commit also adds a new ExecutionSession::lookup convenience method that
performs a blocking lookup using the main JITDylib's search order, as this will
be a very common operation for clients.

Finally, new convenience overloads of IRLayer and ObjectLayer's add methods are
introduced that add the given program representations to the main dylib, which
is likely to be the common case.

llvm-svn: 342086

Author: lhames

llvm/include/llvm/ExecutionEngine/Orc/Core.h

@@ -761,14 +761,31 @@ class JITDylib {
 /// An ExecutionSession represents a running JIT program.
 class ExecutionSession : public ExecutionSessionBase {
 public:
+  using ExecutionSessionBase::lookup;
+
   /// Construct an ExecutionEngine.
   ///
   /// SymbolStringPools may be shared between ExecutionSessions.
-  ExecutionSession(std::shared_ptr<SymbolStringPool> SSP = nullptr)
-      : ExecutionSessionBase(std::move(SSP)) {}
+  ExecutionSession(std::shared_ptr<SymbolStringPool> SSP = nullptr);
+
+  /// Get the "main" JITDylib, which is created automatically on construction of
+  /// the ExecutionSession.
+  JITDylib &getMainJITDylib();
 
   /// Add a new JITDylib to this ExecutionSession.
-  JITDylib &createJITDylib(std::string Name);
+  JITDylib &createJITDylib(std::string Name,
+                           bool AddToMainDylibSearchOrder = true);
+
+  /// Convenience version of the blocking version of lookup in
+  /// ExecutionSessionBase. Uses the main JITDylib's search order as the lookup
+  /// order, and registers no dependencies.
+  Expected<SymbolMap> lookup(const SymbolNameSet &Symbols) {
+    return getMainJITDylib().withSearchOrderDo(
+        [&](const JITDylibList &SearchOrder) {
+          return ExecutionSessionBase::lookup(SearchOrder, Symbols,
+                                              NoDependenciesToRegister, true);
+        });
+  }
 
 private:
   std::vector<std::unique_ptr<JITDylib>> JDs;

llvm/include/llvm/ExecutionEngine/Orc/Layer.h

@@ -34,6 +34,12 @@ class IRLayer {
   /// JITDylib.
   virtual Error add(JITDylib &JD, VModuleKey K, std::unique_ptr<Module> M);
 
+  /// Adds a MaterializationUnit representing the given IR to the main
+  /// JITDylib.
+  Error add(VModuleKey K, std::unique_ptr<Module> M) {
+    return add(ES.getMainJITDylib(), K, std::move(M));
+  }
+
   /// Emit should materialize the given IR.
   virtual void emit(MaterializationResponsibility R, VModuleKey K,
                     std::unique_ptr<Module> M) = 0;
@@ -97,6 +103,12 @@ class ObjectLayer {
   /// JITDylib.
   virtual Error add(JITDylib &JD, VModuleKey K, std::unique_ptr<MemoryBuffer> O);
 
+  /// Adds a MaterializationUnit representing the given object to the main
+  /// JITDylib.
+  Error add(VModuleKey K, std::unique_ptr<MemoryBuffer> O) {
+    return add(ES.getMainJITDylib(), K, std::move(O));
+  }
+
   /// Emit should materialize the given IR.
   virtual void emit(MaterializationResponsibility R, VModuleKey K,
                     std::unique_ptr<MemoryBuffer> O) = 0;

llvm/lib/ExecutionEngine/Orc/Core.cpp

@@ -1677,10 +1677,23 @@ void JITDylib::transferEmittedNodeDependencies(
   }
 }
 
-JITDylib &ExecutionSession::createJITDylib(std::string Name) {
+ExecutionSession::ExecutionSession(std::shared_ptr<SymbolStringPool> SSP)
+    : ExecutionSessionBase(std::move(SSP)) {
+  // Construct the main dylib.
+  JDs.push_back(std::unique_ptr<JITDylib>(new JITDylib(*this, "<main>")));
+}
+
+JITDylib &ExecutionSession::getMainJITDylib() {
+  return runSessionLocked([this]() -> JITDylib & { return *JDs.front(); });
+}
+
+JITDylib &ExecutionSession::createJITDylib(std::string Name,
+                                           bool AddToMainDylibSearchOrder) {
   return runSessionLocked([&, this]() -> JITDylib & {
     JDs.push_back(
         std::unique_ptr<JITDylib>(new JITDylib(*this, std::move(Name))));
+    if (AddToMainDylibSearchOrder)
+      JDs.front()->addToSearchOrder(*JDs.back());
     return *JDs.back();
   });
 }

llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp

@@ -815,4 +815,14 @@ TEST_F(CoreAPIsStandardTest, TestMaterializeWeakSymbol) {
   FooResponsibility->emit();
 }
 
+TEST_F(CoreAPIsStandardTest, TestMainJITDylibAndDefaultLookupOrder) {
+  cantFail(ES.getMainJITDylib().define(absoluteSymbols({{Foo, FooSym}})));
+  auto Results = cantFail(ES.lookup({Foo}));
+
+  EXPECT_EQ(Results.size(), 1U) << "Incorrect number of results";
+  EXPECT_EQ(Results.count(Foo), 1U) << "Expected result for 'Foo'";
+  EXPECT_EQ(Results[Foo].getAddress(), FooSym.getAddress())
+      << "Expected result address to match Foo's address";
+}
+
 } // namespace

llvm/unittests/ExecutionEngine/Orc/OrcTestCommon.h

@@ -46,9 +46,9 @@ namespace orc {
 //     linkage and non-hidden visibility.
 // (5) V -- A JITDylib associated with ES.
 class CoreAPIsBasedStandardTest : public testing::Test {
-public:
 protected:
-  ExecutionSession ES;
+  std::shared_ptr<SymbolStringPool> SSP = std::make_shared<SymbolStringPool>();
+  ExecutionSession ES{SSP};
   JITDylib &JD = ES.createJITDylib("JD");
   SymbolStringPtr Foo = ES.getSymbolStringPool().intern("foo");
   SymbolStringPtr Bar = ES.getSymbolStringPool().intern("bar");