From 72deff436dcd881c23ec2bef65b8f6eb4f309ea7 Mon Sep 17 00:00:00 2001
From: Lang Hames <lhames@gmail.com>
Date: Tue, 15 Oct 2019 21:41:12 +0000
Subject: [PATCH] [Orc] Add a method for ObjectLinkingLayer to return ownership
 of object buffers.

RTDyldObjectLinkingLayer allowed clients to register a NotifyEmitted function to
reclaim ownership of object buffers once they had been linked. This patch adds
similar functionality to ObjectLinkingLayer: Clients can now optionally call the
ObjectLinkingLayer::setReturnObjectBuffer method to register a function that
will be called when discarding object buffers. If set, this function will be
called to return ownership of the object regardless of whether the link
succeeded or failed.

Use cases for this function include debug dumping (it provides a way to dump
all objects linked into JIT'd code) and object re-use (e.g. storing an
object in a cache).

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@374951 91177308-0d34-0410-b5e6-96231b3b80d8
---
 include/llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h | 11 +++++++++++
 lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp        |  7 +++++++
 2 files changed, 18 insertions(+)

diff --git a/include/llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h b/include/llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h
index 0605ee3bb0b..caf8e707516 100644
--- a/include/llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h
+++ b/include/llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h
@@ -73,6 +73,9 @@ public:
     virtual Error notifyRemovingAllModules() { return Error::success(); }
   };
 
+  using ReturnObjectBufferFunction =
+      std::function<void(std::unique_ptr<MemoryBuffer>)>;
+
   /// Construct an ObjectLinkingLayer with the given NotifyLoaded,
   /// and NotifyEmitted functors.
   ObjectLinkingLayer(ExecutionSession &ES,
@@ -81,6 +84,13 @@ public:
   /// Destruct an ObjectLinkingLayer.
   ~ObjectLinkingLayer();
 
+  /// Set an object buffer return function. By default object buffers are
+  /// deleted once the JIT has linked them. If a return function is set then
+  /// it will be called to transfer ownership of the buffer instead.
+  void setReturnObjectBuffer(ReturnObjectBufferFunction ReturnObjectBuffer) {
+    this->ReturnObjectBuffer = std::move(ReturnObjectBuffer);
+  }
+
   /// Add a pass-config modifier.
   ObjectLinkingLayer &addPlugin(std::unique_ptr<Plugin> P) {
     std::lock_guard<std::mutex> Lock(LayerMutex);
@@ -138,6 +148,7 @@ private:
   jitlink::JITLinkMemoryManager &MemMgr;
   bool OverrideObjectFlags = false;
   bool AutoClaimObjectSymbols = false;
+  ReturnObjectBufferFunction ReturnObjectBuffer;
   DenseMap<VModuleKey, AllocPtr> TrackedAllocs;
   std::vector<AllocPtr> UntrackedAllocs;
   std::vector<std::unique_ptr<Plugin>> Plugins;
diff --git a/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp b/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp
index 952ca6071ff..874decb2ade 100644
--- a/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp
+++ b/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp
@@ -29,6 +29,13 @@ public:
                                    std::unique_ptr<MemoryBuffer> ObjBuffer)
       : Layer(Layer), MR(std::move(MR)), ObjBuffer(std::move(ObjBuffer)) {}
 
+  ~ObjectLinkingLayerJITLinkContext() {
+    // If there is an object buffer return function then use it to
+    // return ownership of the buffer.
+    if (Layer.ReturnObjectBuffer)
+      Layer.ReturnObjectBuffer(std::move(ObjBuffer));
+  }
+
   JITLinkMemoryManager &getMemoryManager() override { return Layer.MemMgr; }
 
   MemoryBufferRef getObjectBuffer() const override {