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)>; + /// 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 P) { std::lock_guard Lock(LayerMutex); @@ -138,6 +148,7 @@ private: jitlink::JITLinkMemoryManager &MemMgr; bool OverrideObjectFlags = false; bool AutoClaimObjectSymbols = false; + ReturnObjectBufferFunction ReturnObjectBuffer; DenseMap TrackedAllocs; std::vector UntrackedAllocs; std::vector> 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 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 {