[llvm] [orc-rt] Enable transparent SPS conversion for ptrs via ExecutorAddr. (PR #162069)

[Lang Hames via llvm-commits]

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

Allows SPS wrapper function calls and handles to use pointer arguments. These will be converted to ExecutorAddr for serialization / deserialization.

>From 77067a8f8b63a20fc0a71bb7d0ad0018ae098ed3 Mon Sep 17 00:00:00 2001
From: Lang Hames <lhames at gmail.com>
Date: Mon, 6 Oct 2025 22:05:33 +1100
Subject: [PATCH] [orc-rt] Enable transparent SPS conversion for ptrs via
 ExecutorAddr.

Allows SPS wrapper function calls and handles to use pointer arguments. These
will be converted to ExecutorAddr for serialization / deserialization.
---
 orc-rt/include/orc-rt/SPSWrapperFunction.h  | 72 ++++++++++++---------
 orc-rt/unittests/SPSWrapperFunctionTest.cpp | 21 ++++++
 2 files changed, 62 insertions(+), 31 deletions(-)

diff --git a/orc-rt/include/orc-rt/SPSWrapperFunction.h b/orc-rt/include/orc-rt/SPSWrapperFunction.h
index 275e253330f20..e5ed14f5cd721 100644
--- a/orc-rt/include/orc-rt/SPSWrapperFunction.h
+++ b/orc-rt/include/orc-rt/SPSWrapperFunction.h
@@ -33,47 +33,61 @@ template <typename... SPSArgTs> struct WFSPSHelper {
     return std::move(R);
   }
 
-  template <typename T> static const T &toSerializable(const T &Arg) noexcept {
-    return Arg;
-  }
+  template <typename T> struct Serializable {
+    typedef std::decay_t<T> serializable_type;
+    static const T &to(const T &Arg) noexcept { return Arg; }
+    static T &&from(T &&Arg) noexcept { return std::forward<T>(Arg); }
+  };
 
-  static SPSSerializableError toSerializable(Error Err) {
-    return SPSSerializableError(std::move(Err));
-  }
+  template <typename T> struct Serializable<T *> {
+    typedef ExecutorAddr serializable_type;
+    static ExecutorAddr to(T *Arg) { return ExecutorAddr::fromPtr(Arg); }
+    static T *from(ExecutorAddr A) { return A.toPtr<T *>(); }
+  };
 
-  template <typename T>
-  static SPSSerializableExpected<T> toSerializable(Expected<T> Arg) {
-    return SPSSerializableExpected<T>(std::move(Arg));
-  }
+  template <> struct Serializable<Error> {
+    typedef SPSSerializableError serializable_type;
+    static SPSSerializableError to(Error Err) {
+      return SPSSerializableError(std::move(Err));
+    }
+    static Error from(SPSSerializableError Err) { return Err.toError(); }
+  };
+
+  template <typename T> struct Serializable<Expected<T>> {
+    typedef SPSSerializableExpected<T> serializable_type;
+    static SPSSerializableExpected<T> to(Expected<T> Val) {
+      return SPSSerializableExpected<T>(std::move(Val));
+    }
+    static Expected<T> from(SPSSerializableExpected<T> Val) {
+      return Val.toExpected();
+    }
+  };
 
   template <typename... Ts> struct DeserializableTuple;
 
   template <typename... Ts> struct DeserializableTuple<std::tuple<Ts...>> {
-    typedef std::tuple<
-        std::decay_t<decltype(toSerializable(std::declval<Ts>()))>...>
-        type;
+    typedef std::tuple<typename Serializable<Ts>::serializable_type...> type;
   };
 
   template <typename... Ts>
   using DeserializableTuple_t = typename DeserializableTuple<Ts...>::type;
 
-  template <typename T> static T &&fromSerializable(T &&Arg) noexcept {
-    return std::forward<T>(Arg);
-  }
-
-  static Error fromSerializable(SPSSerializableError Err) {
-    return Err.toError();
-  }
-
-  template <typename T>
-  static Expected<T> fromSerializable(SPSSerializableExpected<T> Val) {
-    return Val.toExpected();
+  template <typename ArgTuple, typename... SerializableArgs, std::size_t... Is>
+  std::optional<ArgTuple>
+  applySerializationConversions(std::tuple<SerializableArgs...> &Inputs,
+                                std::index_sequence<Is...>) {
+    static_assert(sizeof...(SerializableArgs) ==
+                      std::index_sequence<Is...>::size(),
+                  "Tuple sizes don't match");
+    return std::optional<ArgTuple>(
+        std::in_place, Serializable<std::tuple_element_t<Is, ArgTuple>>::from(
+                           std::move(std::get<Is>(Inputs)))...);
   }
 
 public:
   template <typename... ArgTs>
   std::optional<WrapperFunctionBuffer> serialize(ArgTs &&...Args) {
-    return serializeImpl(toSerializable(std::forward<ArgTs>(Args))...);
+    return serializeImpl(Serializable<ArgTs>::to(std::forward<ArgTs>(Args))...);
   }
 
   template <typename ArgTuple>
@@ -85,12 +99,8 @@ template <typename... SPSArgTs> struct WFSPSHelper {
     if (!SPSSerializationTraits<SPSTuple<SPSArgTs...>,
                                 decltype(Args)>::deserialize(IB, Args))
       return std::nullopt;
-    return std::apply(
-        [](auto &&...A) {
-          return std::optional<ArgTuple>(std::in_place,
-                                         std::move(fromSerializable(A))...);
-        },
-        std::move(Args));
+    return applySerializationConversions<ArgTuple>(
+        Args, std::make_index_sequence<std::tuple_size_v<ArgTuple>>());
   }
 };
 
diff --git a/orc-rt/unittests/SPSWrapperFunctionTest.cpp b/orc-rt/unittests/SPSWrapperFunctionTest.cpp
index e010e2a067adf..a8409cc10052f 100644
--- a/orc-rt/unittests/SPSWrapperFunctionTest.cpp
+++ b/orc-rt/unittests/SPSWrapperFunctionTest.cpp
@@ -221,6 +221,27 @@ TEST(SPSWrapperFunctionUtilsTest, TestFunctionReturningExpectedFailureCase) {
   EXPECT_EQ(ErrMsg, "N is not a multiple of 2");
 }
 
+static void
+round_trip_int_pointer_sps_wrapper(orc_rt_SessionRef Session, void *CallCtx,
+                                   orc_rt_WrapperFunctionReturn Return,
+                                   orc_rt_WrapperFunctionBuffer ArgBytes) {
+  SPSWrapperFunction<SPSExecutorAddr(SPSExecutorAddr)>::handle(
+      Session, CallCtx, Return, ArgBytes,
+      [](move_only_function<void(int32_t *)> Return, int32_t *P) {
+        Return(P);
+      });
+}
+
+TEST(SPSWrapperFunctionUtilsTest, TestTransparentSerializationForPointers) {
+  int X = 42;
+  int *P = nullptr;
+  SPSWrapperFunction<SPSExecutorAddr(SPSExecutorAddr)>::call(
+      DirectCaller(nullptr, round_trip_int_pointer_sps_wrapper),
+      [&](Expected<int32_t *> R) { P = cantFail(std::move(R)); }, &X);
+
+  EXPECT_EQ(P, &X);
+}
+
 template <size_t N> struct SPSOpCounter {};
 
 namespace orc_rt {