2 files changed, 11 insertions, 4 deletions
diff --git a/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.cpp b/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.cpp
index 6cf1d827fef..38ad1ec7c59 100644
--- a/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.cpp
+++ b/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.cpp
@@ -90,7 +90,7 @@ void JITLinkerBase::linkPhase2(std::unique_ptr<JITLinkerBase> Self,
                                Expected<AsyncLookupResult> LR) {
   // If the lookup failed, bail out.
   if (!LR)
-    return Ctx->notifyFailed(LR.takeError());
+    return deallocateAndBailOut(LR.takeError());
 
   // Assign addresses to external atoms.
   applyLookupResult(*LR);
@@ -102,7 +102,7 @@ void JITLinkerBase::linkPhase2(std::unique_ptr<JITLinkerBase> Self,
 
   // Copy atom content to working memory and fix up.
   if (auto Err = copyAndFixUpAllAtoms(Layout, *Alloc))
-    return Ctx->notifyFailed(std::move(Err));
+    return deallocateAndBailOut(std::move(Err));
 
   LLVM_DEBUG({
     dbgs() << "Atom graph \"" << G->getName() << "\" after copy-and-fixup:\n";
@@ -110,7 +110,7 @@ void JITLinkerBase::linkPhase2(std::unique_ptr<JITLinkerBase> Self,
   });
 
   if (auto Err = runPasses(Passes.PostFixupPasses, *G))
-    return Ctx->notifyFailed(std::move(Err));
+    return deallocateAndBailOut(std::move(Err));
 
   // FIXME: Use move capture once we have c++14.
   auto *UnownedSelf = Self.release();
@@ -124,7 +124,7 @@ void JITLinkerBase::linkPhase2(std::unique_ptr<JITLinkerBase> Self,
 
 void JITLinkerBase::linkPhase3(std::unique_ptr<JITLinkerBase> Self, Error Err) {
   if (Err)
-    return Ctx->notifyFailed(std::move(Err));
+    return deallocateAndBailOut(std::move(Err));
   Ctx->notifyFinalized(std::move(Alloc));
 }
 
@@ -356,6 +356,12 @@ void JITLinkerBase::applyLookupResult(AsyncLookupResult Result) {
          "All atoms should have been resolved by this point");
 }
 
+void JITLinkerBase::deallocateAndBailOut(Error Err) {
+  assert(Err && "Should not be bailing out on success value");
+  assert(Alloc && "can not call deallocateAndBailOut before allocation");
+  Ctx->notifyFailed(joinErrors(std::move(Err), Alloc->deallocate()));
+}
+
 void JITLinkerBase::dumpGraph(raw_ostream &OS) {
   assert(G && "Graph is not set yet");
   G->dump(dbgs(), [this](Edge::Kind K) { return getEdgeKindName(K); });
diff --git a/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.h b/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.h
index 74db48fdff2..9d8238cf2e5 100644
--- a/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.h
+++ b/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.h
@@ -102,6 +102,7 @@ private:
   Error allocateSegments(const SegmentLayoutMap &Layout);
   DenseSet<StringRef> getExternalSymbolNames() const;
   void applyLookupResult(AsyncLookupResult LR);
+  void deallocateAndBailOut(Error Err);
 
   void dumpGraph(raw_ostream &OS);