[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.

Author: lhames

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>;
 

llvm/include/llvm/ExecutionEngine/Orc/Shared/WrapperFunctionUtils.h

@@ -660,7 +660,7 @@ class WrapperFunctionCall {
   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 @@ class WrapperFunctionCall {
   /// 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 @@ class WrapperFunctionCall {
 
   /// 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;

llvm/lib/ExecutionEngine/JITLink/JITLinkMemoryManager.cpp

@@ -247,19 +247,11 @@ class InProcessMemoryManager::IPInFlightAlloc
     }
 
     // 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 @@ class InProcessMemoryManager::IPInFlightAlloc
 
     // Continue with finalized allocation.
     OnFinalized(MemMgr.createFinalizedAlloc(std::move(StandardSegments),
-                                            std::move(DeallocActions)));
+                                            std::move(*DeallocActions)));
   }
 
   void abandon(OnAbandonedFunction OnAbandoned) override {

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

llvm/lib/ExecutionEngine/Orc/Shared/CMakeLists.txt

@@ -1,4 +1,5 @@
 add_llvm_component_library(LLVMOrcShared
+  AllocationActions.cpp
   OrcError.cpp
   OrcRTBridge.cpp
   SimpleRemoteEPCUtils.cpp

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 @@ class JITLinkSlabAllocator final : public JITLinkMemoryManager {
           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 {