[orc-rt] Add SPS serialization for AllocAction and AllocActionPair. (#157620)

Author: lhames

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

@@ -15,10 +15,57 @@
 #define ORC_RT_SPSALLOCACTION_H
 
 #include "orc-rt/AllocAction.h"
+#include "orc-rt/SPSWrapperFunctionBuffer.h"
 #include "orc-rt/SimplePackedSerialization.h"
 
 namespace orc_rt {
 
+struct SPSAllocAction;
+
+template <> class SPSSerializationTraits<SPSAllocAction, AllocAction> {
+public:
+  static size_t size(const AllocAction &AA) {
+    return SPSArgList<SPSExecutorAddr, SPSWrapperFunctionBuffer>::size(
+        ExecutorAddr::fromPtr(AA.Fn), AA.ArgData);
+  }
+
+  static bool serialize(SPSOutputBuffer &OB, const AllocAction &AA) {
+    return SPSArgList<SPSExecutorAddr, SPSWrapperFunctionBuffer>::serialize(
+        OB, ExecutorAddr::fromPtr(AA.Fn), AA.ArgData);
+  }
+
+  static bool deserialize(SPSInputBuffer &IB, AllocAction &AA) {
+    ExecutorAddr Fn;
+    WrapperFunctionBuffer ArgData;
+    if (!SPSArgList<SPSExecutorAddr, SPSWrapperFunctionBuffer>::deserialize(
+            IB, Fn, ArgData))
+      return false;
+    AA.Fn = Fn.toPtr<AllocActionFn>();
+    AA.ArgData = std::move(ArgData);
+    return true;
+  }
+};
+
+struct SPSAllocActionPair;
+
+template <> class SPSSerializationTraits<SPSAllocActionPair, AllocActionPair> {
+public:
+  static size_t size(const AllocActionPair &AAP) {
+    return SPSArgList<SPSAllocAction, SPSAllocAction>::size(AAP.Finalize,
+                                                            AAP.Dealloc);
+  }
+
+  static bool serialize(SPSOutputBuffer &OB, const AllocActionPair &AAP) {
+    return SPSArgList<SPSAllocAction, SPSAllocAction>::serialize(
+        OB, AAP.Finalize, AAP.Dealloc);
+  }
+
+  static bool deserialize(SPSInputBuffer &IB, AllocActionPair &AAP) {
+    return SPSArgList<SPSAllocAction, SPSAllocAction>::deserialize(
+        IB, AAP.Finalize, AAP.Dealloc);
+  }
+};
+
 template <typename... SPSArgTs> struct AllocActionSPSDeserializer {
   template <typename... ArgTs>
   bool deserialize(const char *ArgData, size_t ArgSize, ArgTs &...Args) {

orc-rt/unittests/CMakeLists.txt

@@ -24,6 +24,7 @@ add_orc_rt_unittest(CoreTests
   RTTITest.cpp
   ScopeExitTest.cpp
   SimplePackedSerializationTest.cpp
+  SPSAllocActionTest.cpp
   SPSMemoryFlagsTest.cpp
   SPSWrapperFunctionTest.cpp
   SPSWrapperFunctionBufferTest.cpp

orc-rt/unittests/SPSAllocActionTest.cpp

@@ -0,0 +1,49 @@
+//===-- SPSAllocActionTest.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
+//
+//===----------------------------------------------------------------------===//
+//
+// Test SPS serialization for AllocAction APIs.
+//
+//===----------------------------------------------------------------------===//
+
+#include "orc-rt/SPSAllocAction.h"
+
+#include "SimplePackedSerializationTestUtils.h"
+#include "gtest/gtest.h"
+
+using namespace orc_rt;
+
+static bool AAEQ(const AllocAction &LHS, const AllocAction &RHS) {
+  if (LHS.Fn != RHS.Fn)
+    return false;
+  if (LHS.ArgData.size() != RHS.ArgData.size())
+    return false;
+  return memcmp(LHS.ArgData.data(), RHS.ArgData.data(), LHS.ArgData.size()) ==
+         0;
+}
+
+static bool AAPEQ(const AllocActionPair &LHS, const AllocActionPair &RHS) {
+  return AAEQ(LHS.Finalize, RHS.Finalize) && AAEQ(LHS.Dealloc, RHS.Dealloc);
+}
+
+static orc_rt_WrapperFunctionBuffer noopAction(const char *ArgData,
+                                               size_t ArgSize) {
+  return WrapperFunctionBuffer().release();
+}
+
+TEST(SPSAllocActionTest, AllocActionSerialization) {
+  AllocAction AA(noopAction, WrapperFunctionBuffer::copyFrom("hello, world!"));
+  blobSerializationRoundTrip<SPSAllocAction>(AA, AAEQ);
+}
+
+TEST(SPSAllocActionTest, AllocActionPairSerialization) {
+  AllocActionPair AAP;
+  AAP.Finalize = {noopAction, WrapperFunctionBuffer::copyFrom("foo")};
+  AAP.Dealloc = {noopAction, WrapperFunctionBuffer::copyFrom("foo")};
+
+  blobSerializationRoundTrip<SPSAllocActionPair>(AAP, AAPEQ);
+}