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

[Lang Hames via llvm-commits]

https://github.com/lhames created https://github.com/llvm/llvm-project/pull/157620

None

>From 51b7b292ae768759ebcb202095e5a7d99efc6be1 Mon Sep 17 00:00:00 2001
From: Lang Hames <lhames at gmail.com>
Date: Tue, 9 Sep 2025 16:59:21 +1000
Subject: [PATCH] [orc-rt] Add SPS serialization for AllocAction and
 AllocActionPair.

---
 orc-rt/include/orc-rt/SPSAllocAction.h  | 47 ++++++++++++++++++++++++
 orc-rt/unittests/CMakeLists.txt         |  1 +
 orc-rt/unittests/SPSAllocActionTest.cpp | 49 +++++++++++++++++++++++++
 3 files changed, 97 insertions(+)
 create mode 100644 orc-rt/unittests/SPSAllocActionTest.cpp

diff --git a/orc-rt/include/orc-rt/SPSAllocAction.h b/orc-rt/include/orc-rt/SPSAllocAction.h
index 6107516029cfd..ce9c11606338f 100644
--- a/orc-rt/include/orc-rt/SPSAllocAction.h
+++ b/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) {
diff --git a/orc-rt/unittests/CMakeLists.txt b/orc-rt/unittests/CMakeLists.txt
index ca501e69034bf..f29fb1e3b9d03 100644
--- a/orc-rt/unittests/CMakeLists.txt
+++ b/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
diff --git a/orc-rt/unittests/SPSAllocActionTest.cpp b/orc-rt/unittests/SPSAllocActionTest.cpp
new file mode 100644
index 0000000000000..b7791e3385382
--- /dev/null
+++ b/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);
+}