5 files changed, 58 insertions, 6 deletions

diff --git a/llvm/include/llvm/ExecutionEngine/Orc/Core.h b/llvm/include/llvm/ExecutionEngine/Orc/Core.h
index e6fe93751e4..e9eb1727596 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/Core.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/Core.h
@@ -761,14 +761,31 @@ private:
 /// 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;
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/Layer.h b/llvm/include/llvm/ExecutionEngine/Orc/Layer.h
index 93e1e52dcee..013350f6233 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/Layer.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/Layer.h
@@ -34,6 +34,12 @@ public:
   /// 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 @@ public:
   /// 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;
diff --git a/llvm/lib/ExecutionEngine/Orc/Core.cpp b/llvm/lib/ExecutionEngine/Orc/Core.cpp
index 1a1fd6384fc..469f849b1c2 100644
--- a/llvm/lib/ExecutionEngine/Orc/Core.cpp
+++ b/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();
   });
 }
diff --git a/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp b/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
index 796736e86bd..b3cd4b517ca 100644
--- a/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
+++ b/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
diff --git a/llvm/unittests/ExecutionEngine/Orc/OrcTestCommon.h b/llvm/unittests/ExecutionEngine/Orc/OrcTestCommon.h
index 05a0b11b866..acc0e5b04ab 100644
--- a/llvm/unittests/ExecutionEngine/Orc/OrcTestCommon.h
+++ b/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");