From 343730c4cb97030103e4d80cd64ed512e32648cf Mon Sep 17 00:00:00 2001
From: Lang Hames <lhames@gmail.com>
Date: Wed, 10 Sep 2025 13:54:22 +1000
Subject: [PATCH] [orc-rt] Add WrapperFunction::handle support for fns,
 fn-ptrs.

Adds support for using functions and function pointers to the
WrapperFunction::handle utility.
---
 orc-rt/include/orc-rt/WrapperFunction.h     |  8 ++++
 orc-rt/unittests/SPSWrapperFunctionTest.cpp | 47 ++++++++++++++++++---
 2 files changed, 50 insertions(+), 5 deletions(-)

diff --git a/orc-rt/include/orc-rt/WrapperFunction.h b/orc-rt/include/orc-rt/WrapperFunction.h
index 46434834bd5c2..bedc097707e38 100644
--- a/orc-rt/include/orc-rt/WrapperFunction.h
+++ b/orc-rt/include/orc-rt/WrapperFunction.h
@@ -121,6 +121,14 @@ struct WFCallableTraits<RetT(ArgT, ArgTs...)> {
   typedef std::tuple<ArgTs...> TailArgTuple;
 };
 
+template <typename RetT, typename... ArgTs>
+struct WFCallableTraits<RetT (*)(ArgTs...)>
+    : public WFCallableTraits<RetT(ArgTs...)> {};
+
+template <typename RetT, typename... ArgTs>
+struct WFCallableTraits<RetT (&)(ArgTs...)>
+    : public WFCallableTraits<RetT(ArgTs...)> {};
+
 template <typename ClassT, typename RetT, typename... ArgTs>
 struct WFCallableTraits<RetT (ClassT::*)(ArgTs...)>
     : public WFCallableTraits<RetT(ArgTs...)> {};
diff --git a/orc-rt/unittests/SPSWrapperFunctionTest.cpp b/orc-rt/unittests/SPSWrapperFunctionTest.cpp
index 919ec2cebd69b..0b65515120b7f 100644
--- a/orc-rt/unittests/SPSWrapperFunctionTest.cpp
+++ b/orc-rt/unittests/SPSWrapperFunctionTest.cpp
@@ -90,9 +90,9 @@ TEST(SPSWrapperFunctionUtilsTest, TestVoidNoop) {
   EXPECT_TRUE(Ran);
 }
 
-static void add_sps_wrapper(orc_rt_SessionRef Session, void *CallCtx,
-                            orc_rt_WrapperFunctionReturn Return,
-                            orc_rt_WrapperFunctionBuffer ArgBytes) {
+static void add_via_lambda_sps_wrapper(orc_rt_SessionRef Session, void *CallCtx,
+                                       orc_rt_WrapperFunctionReturn Return,
+                                       orc_rt_WrapperFunctionBuffer ArgBytes) {
   SPSWrapperFunction<int32_t(int32_t, int32_t)>::handle(
       Session, CallCtx, Return, ArgBytes,
       [](move_only_function<void(int32_t)> Return, int32_t X, int32_t Y) {
@@ -100,10 +100,47 @@ static void add_sps_wrapper(orc_rt_SessionRef Session, void *CallCtx,
       });
 }
 
-TEST(SPSWrapperFunctionUtilsTest, TestAdd) {
+TEST(SPSWrapperFunctionUtilsTest, TestBinaryOpViaLambda) {
   int32_t Result = 0;
   SPSWrapperFunction<int32_t(int32_t, int32_t)>::call(
-      DirectCaller(nullptr, add_sps_wrapper),
+      DirectCaller(nullptr, add_via_lambda_sps_wrapper),
+      [&](Expected<int32_t> R) { Result = cantFail(std::move(R)); }, 41, 1);
+  EXPECT_EQ(Result, 42);
+}
+
+static void add_via_function(move_only_function<void(int32_t)> Return,
+                             int32_t X, int32_t Y) {
+  Return(X + Y);
+}
+
+static void
+add_via_function_sps_wrapper(orc_rt_SessionRef Session, void *CallCtx,
+                             orc_rt_WrapperFunctionReturn Return,
+                             orc_rt_WrapperFunctionBuffer ArgBytes) {
+  SPSWrapperFunction<int32_t(int32_t, int32_t)>::handle(
+      Session, CallCtx, Return, ArgBytes, add_via_function);
+}
+
+TEST(SPSWrapperFunctionUtilsTest, TestBinaryOpViaFunction) {
+  int32_t Result = 0;
+  SPSWrapperFunction<int32_t(int32_t, int32_t)>::call(
+      DirectCaller(nullptr, add_via_function_sps_wrapper),
+      [&](Expected<int32_t> R) { Result = cantFail(std::move(R)); }, 41, 1);
+  EXPECT_EQ(Result, 42);
+}
+
+static void
+add_via_function_pointer_sps_wrapper(orc_rt_SessionRef Session, void *CallCtx,
+                                     orc_rt_WrapperFunctionReturn Return,
+                                     orc_rt_WrapperFunctionBuffer ArgBytes) {
+  SPSWrapperFunction<int32_t(int32_t, int32_t)>::handle(
+      Session, CallCtx, Return, ArgBytes, &add_via_function);
+}
+
+TEST(SPSWrapperFunctionUtilsTest, TestBinaryOpViaFunctionPointer) {
+  int32_t Result = 0;
+  SPSWrapperFunction<int32_t(int32_t, int32_t)>::call(
+      DirectCaller(nullptr, add_via_function_pointer_sps_wrapper),
       [&](Expected<int32_t> R) { Result = cantFail(std::move(R)); }, 41, 1);
   EXPECT_EQ(Result, 42);
 }