[ORC] Add helper functions for running finalize / dealloc actions.

runFinalizeActions takes an AllocActions vector and attempts to run its finalize
actions. If any finalize action fails then all paired dealloc actions up to the
failing pair are run, and the error(s) returned. If all finalize actions succeed
then a vector containing the dealloc actions is returned.

runDeallocActions takes a vector<WrapperFunctionCall> containing dealloc action
calls and runs them all, returning any error(s).

These helpers are intended to simplify the implementation of
JITLinkMemoryManager::InFlightAlloc::finalize and
JITLinkMemoryManager::deallocate overrides by taking care of execution (and
potential roll-back) of allocation actions.

6 files changed, 98 insertions, 22 deletions

diff --git a/llvm/include/llvm/ExecutionEngine/Orc/Shared/AllocationActions.h b/llvm/include/llvm/ExecutionEngine/Orc/Shared/AllocationActions.h
index 26acb702d33a..6469b87c816f 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/Shared/AllocationActions.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/Shared/AllocationActions.h
@@ -43,6 +43,33 @@ struct AllocActionCallPair {
 /// actions will be run in reverse order at deallocation time.
 using AllocActions = std::vector<AllocActionCallPair>;
 
+/// Returns the number of deallocaton actions in the given AllocActions array.
+///
+/// This can be useful if clients want to pre-allocate room for deallocation
+/// actions with the rest of their memory.
+inline size_t numDeallocActions(const AllocActions &AAs) {
+  return llvm::count_if(
+      AAs, [](const AllocActionCallPair &P) { return !!P.Dealloc; });
+}
+
+/// Run finalize actions.
+///
+/// If any finalize action fails then the corresponding dealloc actions will be
+/// run in reverse order (not including the deallocation action for the failed
+/// finalize action), and the error for the failing action will be returned.
+///
+/// If all finalize actions succeed then a vector of deallocation actions will
+/// be returned. The dealloc actions should be run by calling
+/// runDeallocationActions. If this function succeeds then the AA argument will
+/// be cleared before the function returns.
+Expected<std::vector<WrapperFunctionCall>>
+runFinalizeActions(AllocActions &AAs);
+
+/// Run deallocation actions.
+/// Dealloc actions will be run in reverse order (from last element of DAs to
+/// first).
+Error runDeallocActions(ArrayRef<WrapperFunctionCall> DAs);
+
 using SPSAllocActionCallPair =
     SPSTuple<SPSWrapperFunctionCall, SPSWrapperFunctionCall>;
 
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/Shared/WrapperFunctionUtils.h b/llvm/include/llvm/ExecutionEngine/Orc/Shared/WrapperFunctionUtils.h
index b1991c880a7a..eb3fb084b28b 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/Shared/WrapperFunctionUtils.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/Shared/WrapperFunctionUtils.h
@@ -660,7 +660,7 @@ public:
   explicit operator bool() const { return !!FnAddr; }
 
   /// Run call returning raw WrapperFunctionResult.
-  shared::WrapperFunctionResult run() {
+  shared::WrapperFunctionResult run() const {
     using FnTy =
         shared::CWrapperFunctionResult(const char *ArgData, size_t ArgSize);
     return shared::WrapperFunctionResult(
@@ -670,7 +670,7 @@ public:
   /// Run call and deserialize result using SPS.
   template <typename SPSRetT, typename RetT>
   std::enable_if_t<!std::is_same<SPSRetT, void>::value, Error>
-  runWithSPSRet(RetT &RetVal) {
+  runWithSPSRet(RetT &RetVal) const {
     auto WFR = run();
     if (const char *ErrMsg = WFR.getOutOfBandError())
       return make_error<StringError>(ErrMsg, inconvertibleErrorCode());
@@ -684,14 +684,15 @@ public:
 
   /// Overload for SPS functions returning void.
   template <typename SPSRetT>
-  std::enable_if_t<std::is_same<SPSRetT, void>::value, Error> runWithSPSRet() {
+  std::enable_if_t<std::is_same<SPSRetT, void>::value, Error>
+  runWithSPSRet() const {
     shared::SPSEmpty E;
     return runWithSPSRet<shared::SPSEmpty>(E);
   }
 
   /// Run call and deserialize an SPSError result. SPSError returns and
   /// deserialization failures are merged into the returned error.
-  Error runWithSPSRetErrorMerged() {
+  Error runWithSPSRetErrorMerged() const {
     detail::SPSSerializableError RetErr;
     if (auto Err = runWithSPSRet<SPSError>(RetErr))
       return Err;
diff --git a/llvm/lib/ExecutionEngine/JITLink/JITLinkMemoryManager.cpp b/llvm/lib/ExecutionEngine/JITLink/JITLinkMemoryManager.cpp
index 182a452e9495..9315ac4f6120 100644
--- a/llvm/lib/ExecutionEngine/JITLink/JITLinkMemoryManager.cpp
+++ b/llvm/lib/ExecutionEngine/JITLink/JITLinkMemoryManager.cpp
@@ -247,19 +247,11 @@ public:
     }
 
     // Run finalization actions.
-    // FIXME: Roll back previous successful actions on failure.
-    std::vector<orc::shared::WrapperFunctionCall> DeallocActions;
-    DeallocActions.reserve(G.allocActions().size());
-    for (auto &ActPair : G.allocActions()) {
-      if (ActPair.Finalize)
-        if (auto Err = ActPair.Finalize.runWithSPSRetErrorMerged()) {
-          OnFinalized(std::move(Err));
-          return;
-        }
-      if (ActPair.Dealloc)
-        DeallocActions.push_back(ActPair.Dealloc);
+    auto DeallocActions = runFinalizeActions(G.allocActions());
+    if (!DeallocActions) {
+      OnFinalized(DeallocActions.takeError());
+      return;
     }
-    G.allocActions().clear();
 
     // Release the finalize segments slab.
     if (auto EC = sys::Memory::releaseMappedMemory(FinalizationSegments)) {
@@ -269,7 +261,7 @@ public:
 
     // Continue with finalized allocation.
     OnFinalized(MemMgr.createFinalizedAlloc(std::move(StandardSegments),
-                                            std::move(DeallocActions)));
+                                            std::move(*DeallocActions)));
   }
 
   void abandon(OnAbandonedFunction OnAbandoned) override {
diff --git a/llvm/lib/ExecutionEngine/Orc/Shared/AllocationActions.cpp b/llvm/lib/ExecutionEngine/Orc/Shared/AllocationActions.cpp
new file mode 100644
index 000000000000..91f2899449ef
--- /dev/null
+++ b/llvm/lib/ExecutionEngine/Orc/Shared/AllocationActions.cpp
@@ -0,0 +1,44 @@
+//===----- AllocationActions.gpp -- JITLink allocation support calls  -----===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ExecutionEngine/Orc/Shared/AllocationActions.h"
+
+namespace llvm {
+namespace orc {
+namespace shared {
+
+Expected<std::vector<WrapperFunctionCall>>
+runFinalizeActions(AllocActions &AAs) {
+  std::vector<WrapperFunctionCall> DeallocActions;
+  DeallocActions.reserve(numDeallocActions(AAs));
+
+  for (auto &AA : AAs) {
+    if (AA.Finalize)
+      if (auto Err = AA.Finalize.runWithSPSRetErrorMerged())
+        return joinErrors(std::move(Err), runDeallocActions(DeallocActions));
+
+    if (AA.Dealloc)
+      DeallocActions.push_back(std::move(AA.Dealloc));
+  }
+
+  AAs.clear();
+  return DeallocActions;
+}
+
+Error runDeallocActions(ArrayRef<WrapperFunctionCall> DAs) {
+  Error Err = Error::success();
+  while (!DAs.empty()) {
+    Err = joinErrors(std::move(Err), DAs.back().runWithSPSRetErrorMerged());
+    DAs = DAs.drop_back();
+  }
+  return Err;
+}
+
+} // namespace shared
+} // namespace orc
+} // namespace llvm
diff --git a/llvm/lib/ExecutionEngine/Orc/Shared/CMakeLists.txt b/llvm/lib/ExecutionEngine/Orc/Shared/CMakeLists.txt
index bdc731e96af2..b4e1b8a88c24 100644
--- a/llvm/lib/ExecutionEngine/Orc/Shared/CMakeLists.txt
+++ b/llvm/lib/ExecutionEngine/Orc/Shared/CMakeLists.txt
@@ -1,4 +1,5 @@
 add_llvm_component_library(LLVMOrcShared
+  AllocationActions.cpp
   OrcError.cpp
   OrcRTBridge.cpp
   SimpleRemoteEPCUtils.cpp
diff --git a/llvm/tools/llvm-jitlink/llvm-jitlink.cpp b/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
index dec4b6301062..8d0fd0403d4d 100644
--- a/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
+++ b/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
@@ -397,6 +397,9 @@ static void dumpSectionContents(raw_ostream &OS, LinkGraph &G) {
 class JITLinkSlabAllocator final : public JITLinkMemoryManager {
 private:
   struct FinalizedAllocInfo {
+    FinalizedAllocInfo(sys::MemoryBlock Mem,
+                       std::vector<shared::WrapperFunctionCall> DeallocActions)
+        : Mem(Mem), DeallocActions(std::move(DeallocActions)) {}
     sys::MemoryBlock Mem;
     std::vector<shared::WrapperFunctionCall> DeallocActions;
   };
@@ -430,12 +433,20 @@ public:
           return;
         }
 
-        // FIXME: Run finalize actions.
-        assert(BL.graphAllocActions().empty() &&
-               "Support function calls not supported yet");
+        auto DeallocActions = runFinalizeActions(BL.graphAllocActions());
+        if (!DeallocActions) {
+          OnFinalized(DeallocActions.takeError());
+          return;
+        }
+
+        if (auto Err = Parent.freeBlock(FinalizeSegs)) {
+          OnFinalized(
+              joinErrors(std::move(Err), runDeallocActions(*DeallocActions)));
+          return;
+        }
 
-        OnFinalized(
-            FinalizedAlloc(ExecutorAddr::fromPtr(new FinalizedAllocInfo())));
+        OnFinalized(FinalizedAlloc(ExecutorAddr::fromPtr(
+            new FinalizedAllocInfo(StandardSegs, std::move(*DeallocActions)))));
       }
 
       void abandon(OnAbandonedFunction OnAbandoned) override {