2 files changed, 37 insertions, 25 deletions

diff --git a/llvm/include/llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h b/llvm/include/llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h
index cd9ec36da0d..479658b11e9 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h
@@ -43,22 +43,34 @@ public:
                          const RuntimeDyld::LoadedObjectInfo &)>;
 
   /// Functor for receiving finalization notifications.
-  using NotifyEmittedFunction = std::function<void(VModuleKey)>;
+  using NotifyEmittedFunction =
+      std::function<void(VModuleKey, std::unique_ptr<MemoryBuffer>)>;
 
   using GetMemoryManagerFunction =
       std::function<std::unique_ptr<RuntimeDyld::MemoryManager>()>;
 
   /// Construct an ObjectLinkingLayer with the given NotifyLoaded,
   ///        and NotifyEmitted functors.
-  RTDyldObjectLinkingLayer(
-      ExecutionSession &ES, GetMemoryManagerFunction GetMemoryManager,
-      NotifyLoadedFunction NotifyLoaded = NotifyLoadedFunction(),
-      NotifyEmittedFunction NotifyEmitted = NotifyEmittedFunction());
+  RTDyldObjectLinkingLayer(ExecutionSession &ES,
+                           GetMemoryManagerFunction GetMemoryManager);
 
   /// Emit the object.
   void emit(MaterializationResponsibility R,
             std::unique_ptr<MemoryBuffer> O) override;
 
+  /// Set the NotifyLoaded callback.
+  RTDyldObjectLinkingLayer &setNotifyLoaded(NotifyLoadedFunction NotifyLoaded) {
+    this->NotifyLoaded = std::move(NotifyLoaded);
+    return *this;
+  }
+
+  /// Set the NotifyEmitted callback.
+  RTDyldObjectLinkingLayer &
+  setNotifyEmitted(NotifyEmittedFunction NotifyEmitted) {
+    this->NotifyEmitted = std::move(NotifyEmitted);
+    return *this;
+  }
+
   /// Set the 'ProcessAllSections' flag.
   ///
   /// If set to true, all sections in each object file will be allocated using
@@ -108,7 +120,8 @@ private:
                   std::map<StringRef, JITEvaluatedSymbol> Resolved,
                   std::set<StringRef> &InternalSymbols);
 
-  void onObjEmit(VModuleKey K, MaterializationResponsibility &R, Error Err);
+  void onObjEmit(VModuleKey K, std::unique_ptr<MemoryBuffer> ObjBuffer,
+                 MaterializationResponsibility &R, Error Err);
 
   mutable std::mutex RTDyldLayerMutex;
   GetMemoryManagerFunction GetMemoryManager;
diff --git a/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp b/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp
index 29c5445b9d1..97f69ebafa2 100644
--- a/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp
@@ -77,11 +77,8 @@ namespace llvm {
 namespace orc {
 
 RTDyldObjectLinkingLayer::RTDyldObjectLinkingLayer(
-    ExecutionSession &ES, GetMemoryManagerFunction GetMemoryManager,
-    NotifyLoadedFunction NotifyLoaded, NotifyEmittedFunction NotifyEmitted)
-    : ObjectLayer(ES), GetMemoryManager(GetMemoryManager),
-      NotifyLoaded(std::move(NotifyLoaded)),
-      NotifyEmitted(std::move(NotifyEmitted)) {}
+    ExecutionSession &ES, GetMemoryManagerFunction GetMemoryManager)
+    : ObjectLayer(ES), GetMemoryManager(GetMemoryManager) {}
 
 void RTDyldObjectLinkingLayer::emit(MaterializationResponsibility R,
                                     std::unique_ptr<MemoryBuffer> O) {
@@ -95,7 +92,13 @@ void RTDyldObjectLinkingLayer::emit(MaterializationResponsibility R,
 
   auto &ES = getExecutionSession();
 
-  auto Obj = object::ObjectFile::createObjectFile(*O);
+  // Create a MemoryBufferRef backed MemoryBuffer (i.e. shallow) copy of the
+  // the underlying buffer to pass into RuntimeDyld. This allows us to hold
+  // ownership of the real underlying buffer and return it to the user once
+  // the object has been emitted.
+  auto ObjBuffer = MemoryBuffer::getMemBuffer(O->getMemBufferRef(), false);
+
+  auto Obj = object::ObjectFile::createObjectFile(*ObjBuffer);
 
   if (!Obj) {
     getExecutionSession().reportError(Obj.takeError());
@@ -133,13 +136,8 @@ void RTDyldObjectLinkingLayer::emit(MaterializationResponsibility R,
 
   JITDylibSearchOrderResolver Resolver(*SharedR);
 
-  /* Thoughts on proper cross-dylib weak symbol handling:
-   *
-   * Change selection of canonical defs to be a manually triggered process, and
-   * add a 'canonical' bit to symbol definitions. When canonical def selection
-   * is triggered, sweep the JITDylibs to mark defs as canonical, discard
-   * duplicate defs.
-   */
+  // FIXME: Switch to move-capture for the 'O' buffer once we have c++14.
+  MemoryBuffer *UnownedObjBuffer = O.release();
   jitLinkForORC(
       **Obj, std::move(O), *MemMgr, Resolver, ProcessAllSections,
       [this, K, SharedR, &Obj, InternalSymbols](
@@ -148,8 +146,9 @@ void RTDyldObjectLinkingLayer::emit(MaterializationResponsibility R,
         return onObjLoad(K, *SharedR, **Obj, std::move(LoadedObjInfo),
                          ResolvedSymbols, *InternalSymbols);
       },
-      [this, K, SharedR](Error Err) {
-        onObjEmit(K, *SharedR, std::move(Err));
+      [this, K, SharedR, UnownedObjBuffer](Error Err) {
+        std::unique_ptr<MemoryBuffer> ObjBuffer(UnownedObjBuffer);
+        onObjEmit(K, std::move(ObjBuffer), *SharedR, std::move(Err));
       });
 }
 
@@ -196,9 +195,9 @@ Error RTDyldObjectLinkingLayer::onObjLoad(
   return Error::success();
 }
 
-void RTDyldObjectLinkingLayer::onObjEmit(VModuleKey K,
-                                          MaterializationResponsibility &R,
-                                          Error Err) {
+void RTDyldObjectLinkingLayer::onObjEmit(
+    VModuleKey K, std::unique_ptr<MemoryBuffer> ObjBuffer,
+    MaterializationResponsibility &R, Error Err) {
   if (Err) {
     getExecutionSession().reportError(std::move(Err));
     R.failMaterialization();
@@ -208,7 +207,7 @@ void RTDyldObjectLinkingLayer::onObjEmit(VModuleKey K,
   R.emit();
 
   if (NotifyEmitted)
-    NotifyEmitted(K);
+    NotifyEmitted(K, std::move(ObjBuffer));
 }
 
 } // End namespace orc.