[llvm] a01f772 - [ORC] Add MethodWrapperHandler utility for WrapperFunction.

[Lang Hames via llvm-commits]

Author: Lang Hames
Date: 2021-09-16T14:36:05+10:00
New Revision: a01f772d19d53976a6d6e0eed684440e4f2cc9da

URL: https://github.com/llvm/llvm-project/commit/a01f772d19d53976a6d6e0eed684440e4f2cc9da
DIFF: https://github.com/llvm/llvm-project/commit/a01f772d19d53976a6d6e0eed684440e4f2cc9da.diff

LOG: [ORC] Add MethodWrapperHandler utility for WrapperFunction.

MethodWrapperHandler removes some of the boilerplate when writing wrapper
functions to wrap method calls. It can be used as a handler for wrapper
functions whose first argument is an ExecutorAddress: the address is cast to a
pointer of the given class type, then the given method function pointer is
called on that object pointer (passing the rest of the arguments).

E.g.

class MyClass {
public:
  void myMethod(uint32_t, bool) { ... }
};

// SPS Method signature for myMethod -- note MyClass object address as first
// argument.
using SPSMyMethodWrapperSignature =
  SPSTuple<SPSExecutorAddress, uint32_t, bool>;

// Wrapper function for myMethod.
WrapperFunctionResult
myMethodCallWrapper(const char *ArgData, size_t ArgSize) {
  return WrapperFunction<SPSMyMethodWrapperSignature>::handle(
     ArgData, ArgSize, makeMethodWrapperHandler(&MyClass::myMethod));
}

Added: 
    

Modified: 
    llvm/include/llvm/ExecutionEngine/Orc/Shared/WrapperFunctionUtils.h
    llvm/unittests/ExecutionEngine/Orc/WrapperFunctionUtilsTest.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/ExecutionEngine/Orc/Shared/WrapperFunctionUtils.h b/llvm/include/llvm/ExecutionEngine/Orc/Shared/WrapperFunctionUtils.h
index eb27c9674aabc..a64130d04faab 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/Shared/WrapperFunctionUtils.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/Shared/WrapperFunctionUtils.h
@@ -13,6 +13,7 @@
 #ifndef LLVM_EXECUTIONENGINE_ORC_SHARED_WRAPPERFUNCTIONUTILS_H
 #define LLVM_EXECUTIONENGINE_ORC_SHARED_WRAPPERFUNCTIONUTILS_H
 
+#include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h"
 #include "llvm/ExecutionEngine/Orc/Shared/SimplePackedSerialization.h"
 #include "llvm/Support/Error.h"
 
@@ -566,6 +567,47 @@ class WrapperFunction<void(SPSTagTs...)>
   using WrapperFunction<SPSEmpty(SPSTagTs...)>::handleAsync;
 };
 
+/// A function object that takes an ExecutorAddress as its first argument,
+/// casts that address to a ClassT*, then calls the given method on that
+/// pointer passing in the remaining function arguments. This utility
+/// removes some of the boilerplate from writing wrappers for method calls.
+///
+///   @code{.cpp}
+///   class MyClass {
+///   public:
+///     void myMethod(uint32_t, bool) { ... }
+///   };
+///
+///   // SPS Method signature -- note MyClass object address as first argument.
+///   using SPSMyMethodWrapperSignature =
+///     SPSTuple<SPSExecutorAddress, uint32_t, bool>;
+///
+///   WrapperFunctionResult
+///   myMethodCallWrapper(const char *ArgData, size_t ArgSize) {
+///     return WrapperFunction<SPSMyMethodWrapperSignature>::handle(
+///        ArgData, ArgSize, makeMethodWrapperHandler(&MyClass::myMethod));
+///   }
+///   @endcode
+///
+template <typename RetT, typename ClassT, typename... ArgTs>
+class MethodWrapperHandler {
+public:
+  using MethodT = RetT (ClassT::*)(ArgTs...);
+  MethodWrapperHandler(MethodT M) : M(M) {}
+  RetT operator()(ExecutorAddress ObjAddr, ArgTs &... Args) {
+    return (ObjAddr.toPtr<ClassT*>()->*M)(std::forward<ArgTs>(Args)...);
+  }
+private:
+  MethodT M;
+};
+
+/// Create a MethodWrapperHandler object from the given method pointer.
+template <typename RetT, typename ClassT, typename... ArgTs>
+MethodWrapperHandler<RetT, ClassT, ArgTs...>
+makeMethodWrapperHandler(RetT (ClassT::*Method)(ArgTs...)) {
+  return MethodWrapperHandler<RetT, ClassT, ArgTs...>(Method);
+}
+
 } // end namespace shared
 } // end namespace orc
 } // end namespace llvm

diff  --git a/llvm/unittests/ExecutionEngine/Orc/WrapperFunctionUtilsTest.cpp b/llvm/unittests/ExecutionEngine/Orc/WrapperFunctionUtilsTest.cpp
index 42051836506fb..0437b8b97362a 100644
--- a/llvm/unittests/ExecutionEngine/Orc/WrapperFunctionUtilsTest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/WrapperFunctionUtilsTest.cpp
@@ -13,6 +13,7 @@
 #include <future>
 
 using namespace llvm;
+using namespace llvm::orc;
 using namespace llvm::orc::shared;
 
 namespace {
@@ -58,6 +59,14 @@ TEST(WrapperFunctionUtilsTest, WrapperFunctionResultFromOutOfBandError) {
 
 static void voidNoop() {}
 
+class AddClass {
+public:
+  AddClass(int32_t X) : X(X) {}
+  int32_t addMethod(int32_t Y) { return X + Y; }
+private:
+  int32_t X;
+};
+
 static WrapperFunctionResult voidNoopWrapper(const char *ArgData,
                                              size_t ArgSize) {
   return WrapperFunction<void()>::handle(ArgData, ArgSize, voidNoop);
@@ -68,6 +77,12 @@ static WrapperFunctionResult addWrapper(const char *ArgData, size_t ArgSize) {
       ArgData, ArgSize, [](int32_t X, int32_t Y) -> int32_t { return X + Y; });
 }
 
+static WrapperFunctionResult addMethodWrapper(const char *ArgData,
+                                              size_t ArgSize) {
+  return WrapperFunction<int32_t(SPSExecutorAddress, int32_t)>::handle(
+      ArgData, ArgSize, makeMethodWrapperHandler(&AddClass::addMethod));
+}
+
 TEST(WrapperFunctionUtilsTest, WrapperFunctionCallAndHandleVoid) {
   EXPECT_FALSE(!!WrapperFunction<void()>::call(voidNoopWrapper));
 }
@@ -79,6 +94,14 @@ TEST(WrapperFunctionUtilsTest, WrapperFunctionCallAndHandleRet) {
   EXPECT_EQ(Result, (int32_t)3);
 }
 
+TEST(WrapperFunctionUtilsTest, WrapperFunctionMethodCallAndHandleRet) {
+  int32_t Result;
+  AddClass AddObj(1);
+  EXPECT_FALSE(!!WrapperFunction<int32_t(SPSExecutorAddress, int32_t)>::call(
+      addMethodWrapper, Result, ExecutorAddress::fromPtr(&AddObj), 2));
+  EXPECT_EQ(Result, (int32_t)3);
+}
+
 static void voidNoopAsync(unique_function<void(SPSEmpty)> SendResult) {
   SendResult(SPSEmpty());
 }