[orc-rt] Add allocation-action execution support.

orc-rt/include/CMakeLists.txt

@@ -3,6 +3,7 @@ set(ORC_RT_HEADERS
     orc-rt-c/ExternC.h
     orc-rt-c/WrapperFunction.h
     orc-rt-c/orc-rt.h
+    orc-rt/AllocAction.h
     orc-rt/BitmaskEnum.h
     orc-rt/Compiler.h
     orc-rt/Error.h
@@ -15,6 +16,7 @@ set(ORC_RT_HEADERS
     orc-rt/WrapperFunction.h
     orc-rt/ScopeExit.h
     orc-rt/SimplePackedSerialization.h
+    orc-rt/SPSAllocAction.h
     orc-rt/SPSWrapperFunction.h
     orc-rt/bind.h
     orc-rt/bit.h

orc-rt/include/orc-rt/AllocAction.h

@@ -0,0 +1,105 @@
+//===---------- AllocAction.h - Allocation action APIs ----------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// AllocAction and related APIs.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef ORC_RT_ALLOCACTION_H
+#define ORC_RT_ALLOCACTION_H
+
+#include "orc-rt/Error.h"
+#include "orc-rt/WrapperFunction.h"
+
+#include <vector>
+
+namespace orc_rt {
+namespace detail {
+
+template <typename Handler>
+struct AAHandlerTraits
+    : public AAHandlerTraits<
+          decltype(&std::remove_cv_t<std::remove_reference_t<Handler>>::
+                   operator())> {};
+
+template <typename... ArgTs>
+struct AAHandlerTraits<WrapperFunctionBuffer(ArgTs...)> {
+  typedef std::tuple<ArgTs...> ArgTuple;
+};
+
+template <typename Class, typename... ArgTs>
+struct AAHandlerTraits<WrapperFunctionBuffer (Class::*)(ArgTs...)>
+    : public AAHandlerTraits<WrapperFunctionBuffer(ArgTs...)> {};
+
+template <typename Class, typename... ArgTs>
+struct AAHandlerTraits<WrapperFunctionBuffer (Class::*)(ArgTs...) const>
+    : public AAHandlerTraits<WrapperFunctionBuffer(ArgTs...)> {};
+
+} // namespace detail
+
+/// An AllocActionFn is a function that takes an argument blob and returns an
+/// empty WrapperFunctionBuffer on success, or an out-of-band error on failure.
+typedef orc_rt_WrapperFunctionBuffer (*AllocActionFn)(const char *ArgData,
+                                                      size_t ArgSize);
+
+struct AllocActionFunction {
+
+  template <typename Deserializer, typename Handler>
+  static WrapperFunctionBuffer handle(const char *ArgData, size_t ArgSize,
+                                      Deserializer &&D, Handler &&H) {
+    typename detail::AAHandlerTraits<Handler>::ArgTuple Args;
+    if (!D.deserialize(ArgData, ArgSize, Args))
+      return WrapperFunctionBuffer::createOutOfBandError(
+          "Could not deserialize allocation action argument buffer");
+
+    return std::apply(std::forward<Handler>(H), std::move(Args)).release();
+  }
+};
+
+/// An AllocAction is a pair of an AllocActionFn and an argument data buffer.
+struct AllocAction {
+  AllocAction() = default;
+  AllocAction(AllocActionFn AA, WrapperFunctionBuffer ArgData)
+      : AA(AA), ArgData(std::move(ArgData)) {}
+
+  [[nodiscard]] WrapperFunctionBuffer operator()() {
+    assert(AA && "Attempt to call null action");
+    return AA(ArgData.data(), ArgData.size());
+  }
+
+  explicit operator bool() const noexcept { return !!AA; }
+
+  AllocActionFn AA = nullptr;
+  WrapperFunctionBuffer ArgData;
+};
+
+/// An AllocActionPair is a pair of a Finalize action and a Dealloc action.
+struct AllocActionPair {
+  AllocAction Finalize;
+  AllocAction Dealloc;
+};
+
+/// Run the finalize actions in the given sequence.
+///
+/// On success, returns the list of deallocation actions to be run in reverse
+/// order at deallocation time.
+///
+/// On failure, runs deallocation actions associated with any previously
+/// successful finalize actions, then returns an error.
+///
+/// Both finalize and dealloc actions are permitted to be null (i.e. have a
+/// null action function) in which case they are ignored.
+[[nodiscard]] Expected<std::vector<AllocAction>>
+runFinalizeActions(std::vector<AllocActionPair> AAPs);
+
+/// Run the given deallocation actions in revwerse order.
+void runDeallocActions(std::vector<AllocAction> DAAs);
+
+} // namespace orc_rt
+
+#endif // ORC_RT_ALLOCACTION_H

orc-rt/include/orc-rt/SPSAllocAction.h

@@ -0,0 +1,46 @@
+//===---- SPSAllocAction.h - SPS-serialized AllocAction utils ---*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Utilities for implementing allocation actions that take an SPS-serialized
+// argument buffer.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef ORC_RT_SPSALLOCACTION_H
+#define ORC_RT_SPSALLOCACTION_H
+
+#include "orc-rt/AllocAction.h"
+#include "orc-rt/SimplePackedSerialization.h"
+
+namespace orc_rt {
+
+template <typename... SPSArgTs> struct AllocActionSPSDeserializer {
+  template <typename... ArgTs>
+  bool deserialize(const char *ArgData, size_t ArgSize, ArgTs &...Args) {
+    SPSInputBuffer IB(ArgData, ArgSize);
+    return SPSArgList<SPSArgTs...>::deserialize(IB, Args...);
+  }
+};
+
+/// Provides call and handle utilities to simplify writing and invocation of
+/// wrapper functions that use SimplePackedSerialization to serialize and
+/// deserialize their arguments and return values.
+template <typename... SPSArgTs> struct SPSAllocActionFunction {
+
+  template <typename Handler>
+  static WrapperFunctionBuffer handle(const char *ArgData, size_t ArgSize,
+                                      Handler &&H) {
+    return AllocActionFunction::handle(
+        ArgData, ArgSize, AllocActionSPSDeserializer<SPSTuple<SPSArgTs...>>(),
+        std::forward<Handler>(H));
+  }
+};
+
+} // namespace orc_rt
+
+#endif // ORC_RT_SPSWRAPPERFUNCTION_H

orc-rt/lib/executor/AllocAction.cpp

@@ -0,0 +1,55 @@
+//===- AllocAction.cpp ----------------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// AllocAction and related APIs.
+//
+//===----------------------------------------------------------------------===//
+
+#include "orc-rt/AllocAction.h"
+#include "orc-rt/ScopeExit.h"
+
+namespace orc_rt {
+
+Expected<std::vector<AllocAction>>
+runFinalizeActions(std::vector<AllocActionPair> AAPs) {
+  std::vector<AllocAction> DeallocActions;
+  auto RunDeallocActions = make_scope_exit([&]() {
+    while (!DeallocActions.empty()) {
+      // TODO: Log errors from cleanup dealloc actions.
+      {
+        [[maybe_unused]] auto B = DeallocActions.back()();
+      }
+      DeallocActions.pop_back();
+    }
+  });
+
+  for (auto &AAP : AAPs) {
+    if (AAP.Finalize) {
+      auto B = AAP.Finalize();
+      if (const char *ErrMsg = B.getOutOfBandError())
+        return make_error<StringError>(ErrMsg);
+    }
+    if (AAP.Dealloc)
+      DeallocActions.push_back(std::move(AAP.Dealloc));
+  }
+
+  RunDeallocActions.release();
+  return DeallocActions;
+}
+
+void runDeallocActions(std::vector<AllocAction> DAAs) {
+  while (!DAAs.empty()) {
+    // TODO: Log errors from cleanup dealloc actions.
+    {
+      [[maybe_unused]] auto B = DAAs.back()();
+    }
+    DAAs.pop_back();
+  }
+}
+
+} // namespace orc_rt

orc-rt/lib/executor/CMakeLists.txt

@@ -1,5 +1,6 @@
 set(files
   orc-rt-executor.cpp
+  AllocAction.cpp
   RTTI.cpp
   )
 

orc-rt/unittests/AllocActionTest.cpp

@@ -0,0 +1,165 @@
+//===- AllocActionTest.cpp ------------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Tests for orc-rt's AllocAction.h APIs.
+//
+//===----------------------------------------------------------------------===//
+
+#include "orc-rt/AllocAction.h"
+#include "orc-rt/ExecutorAddress.h"
+#include "orc-rt/SPSAllocAction.h"
+
+#include "SimplePackedSerializationTestUtils.h"
+#include "gtest/gtest.h"
+
+using namespace orc_rt;
+
+TEST(AllocActionTest, DefaultConstruct) {
+  AllocAction AA;
+  EXPECT_FALSE(AA);
+}
+
+static orc_rt_WrapperFunctionBuffer noopAction(const char *ArgData,
+                                               size_t ArgSize) {
+  return WrapperFunctionBuffer().release();
+}
+
+TEST(AllocActionTest, ConstructWithAction) {
+  AllocAction AA(noopAction, WrapperFunctionBuffer());
+  EXPECT_TRUE(AA);
+}
+
+// Increments int via pointer.
+static orc_rt_WrapperFunctionBuffer
+increment_sps_allocaction(const char *ArgData, size_t ArgSize) {
+  return SPSAllocActionFunction<SPSExecutorAddr>::handle(
+             ArgData, ArgSize,
+             [](ExecutorAddr IntPtr) {
+               *IntPtr.toPtr<int *>() += 1;
+               return WrapperFunctionBuffer();
+             })
+      .release();
+}
+
+// Increments int via pointer.
+static orc_rt_WrapperFunctionBuffer
+decrement_sps_allocaction(const char *ArgData, size_t ArgSize) {
+  return SPSAllocActionFunction<SPSExecutorAddr>::handle(
+             ArgData, ArgSize,
+             [](ExecutorAddr IntPtr) {
+               *IntPtr.toPtr<int *>() -= 1;
+               return WrapperFunctionBuffer();
+             })
+      .release();
+}
+
+template <typename T>
+static WrapperFunctionBuffer makeExecutorAddrBuffer(T *P) {
+  return *spsSerialize<SPSArgList<SPSExecutorAddr>>(ExecutorAddr::fromPtr(P));
+}
+
+TEST(AllocActionTest, RunBasicAction) {
+  int Val = 0;
+  AllocAction IncVal(increment_sps_allocaction, makeExecutorAddrBuffer(&Val));
+  EXPECT_TRUE(IncVal);
+  auto B = IncVal();
+  EXPECT_TRUE(B.empty());
+  EXPECT_EQ(Val, 1);
+}
+
+TEST(AllocActionTest, RunFinalizationActionsComplete) {
+  int Val = 0;
+
+  std::vector<AllocActionPair> InitialActions;
+
+  auto MakeArgBuffer = [&]() { return makeExecutorAddrBuffer(&Val); };
+  InitialActions.push_back({{increment_sps_allocaction, MakeArgBuffer()},
+                            {decrement_sps_allocaction, MakeArgBuffer()}});
+  InitialActions.push_back({{increment_sps_allocaction, MakeArgBuffer()},
+                            {decrement_sps_allocaction, MakeArgBuffer()}});
+
+  auto DeallocActions = cantFail(runFinalizeActions(std::move(InitialActions)));
+
+  EXPECT_EQ(Val, 2);
+
+  runDeallocActions(std::move(DeallocActions));
+
+  EXPECT_EQ(Val, 0);
+}
+
+static orc_rt_WrapperFunctionBuffer fail_sps_allocaction(const char *ArgData,
+                                                         size_t ArgSize) {
+  return WrapperFunctionBuffer::createOutOfBandError("failed").release();
+}
+
+TEST(AllocActionTest, RunFinalizeActionsFail) {
+  int Val = 0;
+
+  std::vector<AllocActionPair> InitialActions;
+
+  auto MakeArgBuffer = [&]() { return makeExecutorAddrBuffer(&Val); };
+  InitialActions.push_back({{increment_sps_allocaction, MakeArgBuffer()},
+                            {decrement_sps_allocaction, MakeArgBuffer()}});
+  InitialActions.push_back({{fail_sps_allocaction, MakeArgBuffer()},
+                            {decrement_sps_allocaction, MakeArgBuffer()}});
+
+  auto DeallocActions = runFinalizeActions(std::move(InitialActions));
+
+  if (DeallocActions) {
+    ADD_FAILURE() << "Failed to report error from runFinalizeActions";
+    return;
+  }
+
+  EXPECT_EQ(toString(DeallocActions.takeError()), std::string("failed"));
+
+  // Check that we ran the decrement corresponding to the first increment.
+  EXPECT_EQ(Val, 0);
+}
+
+TEST(AllocActionTest, RunFinalizeActionsNullFinalize) {
+  int Val = 0;
+
+  std::vector<AllocActionPair> InitialActions;
+
+  auto MakeArgBuffer = [&]() { return makeExecutorAddrBuffer(&Val); };
+  InitialActions.push_back({{increment_sps_allocaction, MakeArgBuffer()},
+                            {decrement_sps_allocaction, MakeArgBuffer()}});
+  InitialActions.push_back({{nullptr, WrapperFunctionBuffer()},
+                            {decrement_sps_allocaction, MakeArgBuffer()}});
+
+  auto DeallocActions = cantFail(runFinalizeActions(std::move(InitialActions)));
+
+  // Both dealloc actions should be included in the returned list, despite one
+  // of them having a null finalize action.
+  EXPECT_EQ(DeallocActions.size(), 2U);
+
+  runDeallocActions(std::move(DeallocActions));
+
+  EXPECT_EQ(Val, -1);
+}
+
+TEST(AllocActionTest, RunFinalizeActionsNullDealloc) {
+  int Val = 0;
+
+  std::vector<AllocActionPair> InitialActions;
+
+  auto MakeArgBuffer = [&]() { return makeExecutorAddrBuffer(&Val); };
+  InitialActions.push_back({{increment_sps_allocaction, MakeArgBuffer()},
+                            {decrement_sps_allocaction, MakeArgBuffer()}});
+  InitialActions.push_back({{increment_sps_allocaction, MakeArgBuffer()},
+                            {nullptr, WrapperFunctionBuffer()}});
+
+  auto DeallocActions = cantFail(runFinalizeActions(std::move(InitialActions)));
+
+  // Null dealloc actions should be filtered out of the returned list.
+  EXPECT_EQ(DeallocActions.size(), 1U);
+
+  runDeallocActions(std::move(DeallocActions));
+
+  EXPECT_EQ(Val, 1);
+}