Use std::variant in constructors

pentschev · pentschev · commit 3558c0b6c42c · 2023-12-11T11:41:36.000-08:00
diff --git a/cpp/include/ucxx/constructors.h b/cpp/include/ucxx/constructors.h
@@ -56,24 +56,28 @@ std::shared_ptr<Worker> createWorker(std::shared_ptr<Context> context,
                                      const bool enableFuture);
 
 // Transfers
-std::shared_ptr<RequestAm> createRequestAm(std::shared_ptr<Endpoint> endpoint,
-                                           const data::RequestData requestData,
-                                           const bool enablePythonFuture,
-                                           RequestCallbackUserFunction callbackFunction,
-                                           RequestCallbackUserData callbackData);
-
-std::shared_ptr<RequestStream> createRequestStream(std::shared_ptr<Endpoint> endpoint,
-                                                   const data::RequestData requestData,
-                                                   const bool enablePythonFuture);
-
-std::shared_ptr<RequestTag> createRequestTag(std::shared_ptr<Component> endpointOrWorker,
-                                             const data::RequestData requestData,
-                                             const bool enablePythonFuture,
-                                             RequestCallbackUserFunction callbackFunction,
-                                             RequestCallbackUserData callbackData);
-
-std::shared_ptr<RequestTagMulti> createRequestTagMulti(std::shared_ptr<Endpoint> endpoint,
-                                                       const data::RequestData requestData,
-                                                       const bool enablePythonFuture);
+std::shared_ptr<RequestAm> createRequestAm(
+  std::shared_ptr<Endpoint> endpoint,
+  const std::variant<data::AmSend, data::AmReceive> requestData,
+  const bool enablePythonFuture,
+  RequestCallbackUserFunction callbackFunction,
+  RequestCallbackUserData callbackData);
+
+std::shared_ptr<RequestStream> createRequestStream(
+  std::shared_ptr<Endpoint> endpoint,
+  const std::variant<data::StreamSend, data::StreamReceive> requestData,
+  const bool enablePythonFuture);
+
+std::shared_ptr<RequestTag> createRequestTag(
+  std::shared_ptr<Component> endpointOrWorker,
+  const std::variant<data::TagSend, data::TagReceive> requestData,
+  const bool enablePythonFuture,
+  RequestCallbackUserFunction callbackFunction,
+  RequestCallbackUserData callbackData);
+
+std::shared_ptr<RequestTagMulti> createRequestTagMulti(
+  std::shared_ptr<Endpoint> endpoint,
+  const std::variant<data::TagMultiSend, data::TagMultiReceive> requestData,
+  const bool enablePythonFuture);
 
 }  // namespace ucxx
diff --git a/cpp/include/ucxx/request_am.h b/cpp/include/ucxx/request_am.h
@@ -52,7 +52,7 @@ class RequestAm : public Request {
    * @param[in] callbackData        user-defined data to pass to the `callbackFunction`.
    */
   RequestAm(std::shared_ptr<Component> endpointOrWorker,
-            const data::RequestData requestData,
+            const std::variant<data::AmSend, data::AmReceive> requestData,
             const std::string operationName,
             const bool enablePythonFuture                = false,
             RequestCallbackUserFunction callbackFunction = nullptr,
@@ -83,11 +83,12 @@ class RequestAm : public Request {
    *
    * @returns The `shared_ptr<ucxx::RequestAm>` object
    */
-  friend std::shared_ptr<RequestAm> createRequestAm(std::shared_ptr<Endpoint> endpoint,
-                                                    const data::RequestData requestData,
-                                                    const bool enablePythonFuture,
-                                                    RequestCallbackUserFunction callbackFunction,
-                                                    RequestCallbackUserData callbackData);
+  friend std::shared_ptr<RequestAm> createRequestAm(
+    std::shared_ptr<Endpoint> endpoint,
+    const std::variant<data::AmSend, data::AmReceive> requestData,
+    const bool enablePythonFuture,
+    RequestCallbackUserFunction callbackFunction,
+    RequestCallbackUserData callbackData);
 
   virtual void populateDelayedSubmission();
 
diff --git a/cpp/include/ucxx/request_data.h b/cpp/include/ucxx/request_data.h
@@ -197,6 +197,12 @@ dispatch(Ts...) -> dispatch<Ts...>;
 template <class... Ts>
 dispatch(Ts&...) -> dispatch<Ts...>;
 
+template <class T>
+RequestData getRequestData(T t)
+{
+  return std::visit([](auto arg) -> RequestData { return arg; }, t);
+}
+
 }  // namespace data
 
 }  // namespace ucxx
diff --git a/cpp/include/ucxx/request_stream.h b/cpp/include/ucxx/request_stream.h
@@ -39,7 +39,7 @@ class RequestStream : public Request {
    *                                subsequently notified.
    */
   RequestStream(std::shared_ptr<Endpoint> endpoint,
-                const data::RequestData requestData,
+                const std::variant<data::StreamSend, data::StreamReceive> requestData,
                 const std::string operationName,
                 const bool enablePythonFuture = false);
 
@@ -61,9 +61,10 @@ class RequestStream : public Request {
    *
    * @returns The `shared_ptr<ucxx::RequestStream>` object
    */
-  friend std::shared_ptr<RequestStream> createRequestStream(std::shared_ptr<Endpoint> endpoint,
-                                                            const data::RequestData requestData,
-                                                            const bool enablePythonFuture);
+  friend std::shared_ptr<RequestStream> createRequestStream(
+    std::shared_ptr<Endpoint> endpoint,
+    const std::variant<data::StreamSend, data::StreamReceive> requestData,
+    const bool enablePythonFuture);
 
   virtual void populateDelayedSubmission();
 
diff --git a/cpp/include/ucxx/request_tag.h b/cpp/include/ucxx/request_tag.h
@@ -47,7 +47,7 @@ class RequestTag : public Request {
    * @param[in] callbackData        user-defined data to pass to the `callbackFunction`.
    */
   RequestTag(std::shared_ptr<Component> endpointOrWorker,
-             const data::RequestData requestData,
+             const std::variant<data::TagSend, data::TagReceive> requestData,
              const std::string operationName,
              const bool enablePythonFuture                = false,
              RequestCallbackUserFunction callbackFunction = nullptr,
@@ -78,11 +78,12 @@ class RequestTag : public Request {
    *
    * @returns The `shared_ptr<ucxx::RequestTag>` object
    */
-  friend std::shared_ptr<RequestTag> createRequestTag(std::shared_ptr<Component> endpointOrWorker,
-                                                      const data::RequestData requestData,
-                                                      const bool enablePythonFuture,
-                                                      RequestCallbackUserFunction callbackFunction,
-                                                      RequestCallbackUserData callbackData);
+  friend std::shared_ptr<RequestTag> createRequestTag(
+    std::shared_ptr<Component> endpointOrWorker,
+    const std::variant<data::TagSend, data::TagReceive> requestData,
+    const bool enablePythonFuture,
+    RequestCallbackUserFunction callbackFunction,
+    RequestCallbackUserData callbackData);
 
   virtual void populateDelayedSubmission();
 
diff --git a/cpp/include/ucxx/request_tag_multi.h b/cpp/include/ucxx/request_tag_multi.h
@@ -76,7 +76,7 @@ class RequestTagMulti : public Request {
    *                                subsequently notified.
    */
   RequestTagMulti(std::shared_ptr<Endpoint> endpoint,
-                  const data::RequestData requestData,
+                  const std::variant<data::TagMultiSend, data::TagMultiReceive> requestData,
                   const std::string operationName,
                   const bool enablePythonFuture);
 
@@ -152,9 +152,10 @@ class RequestTagMulti : public Request {
    *
    * @returns Request to be subsequently checked for the completion and its state.
    */
-  friend std::shared_ptr<RequestTagMulti> createRequestTagMulti(std::shared_ptr<Endpoint> endpoint,
-                                                                const data::RequestData requestData,
-                                                                const bool enablePythonFuture);
+  friend std::shared_ptr<RequestTagMulti> createRequestTagMulti(
+    std::shared_ptr<Endpoint> endpoint,
+    const std::variant<data::TagMultiSend, data::TagMultiReceive> requestData,
+    const bool enablePythonFuture);
 
   /**
    * @brief `ucxx::RequestTagMulti` destructor.
diff --git a/cpp/include/ucxx/worker.h b/cpp/include/ucxx/worker.h
@@ -55,11 +55,12 @@ class Worker : public Component {
   std::shared_ptr<DelayedSubmissionCollection> _delayedSubmissionCollection{
     nullptr};  ///< Collection of enqueued delayed submissions
 
-  friend std::shared_ptr<RequestAm> createRequestAm(std::shared_ptr<Endpoint> endpoint,
-                                                    const data::RequestData requestData,
-                                                    const bool enablePythonFuture,
-                                                    RequestCallbackUserFunction callbackFunction,
-                                                    RequestCallbackUserData callbackData);
+  friend std::shared_ptr<RequestAm> createRequestAm(
+    std::shared_ptr<Endpoint> endpoint,
+    const std::variant<data::AmSend, data::AmReceive> requestData,
+    const bool enablePythonFuture,
+    RequestCallbackUserFunction callbackFunction,
+    RequestCallbackUserData callbackData);
 
  protected:
   bool _enableFuture{
diff --git a/cpp/src/request_am.cpp b/cpp/src/request_am.cpp
@@ -16,11 +16,12 @@
 
 namespace ucxx {
 
-std::shared_ptr<RequestAm> createRequestAm(std::shared_ptr<Endpoint> endpoint,
-                                           const data::RequestData requestData,
-                                           const bool enablePythonFuture                = false,
-                                           RequestCallbackUserFunction callbackFunction = nullptr,
-                                           RequestCallbackUserData callbackData         = nullptr)
+std::shared_ptr<RequestAm> createRequestAm(
+  std::shared_ptr<Endpoint> endpoint,
+  const std::variant<data::AmSend, data::AmReceive> requestData,
+  const bool enablePythonFuture                = false,
+  RequestCallbackUserFunction callbackFunction = nullptr,
+  RequestCallbackUserData callbackData         = nullptr)
 {
   std::shared_ptr<RequestAm> req = std::visit(
     data::dispatch{
@@ -49,28 +50,26 @@ std::shared_ptr<RequestAm> createRequestAm(std::shared_ptr<Endpoint> endpoint,
         };
         return worker->getAmRecv(endpoint->getHandle(), createRequest);
       },
-      [](auto) -> decltype(req) { throw std::runtime_error("Unreachable"); },
     },
     requestData);
 
   return req;
 }
 
 RequestAm::RequestAm(std::shared_ptr<Component> endpointOrWorker,
-                     const data::RequestData requestData,
+                     const std::variant<data::AmSend, data::AmReceive> requestData,
                      const std::string operationName,
                      const bool enablePythonFuture,
                      RequestCallbackUserFunction callbackFunction,
                      RequestCallbackUserData callbackData)
-  : Request(endpointOrWorker, requestData, operationName, enablePythonFuture)
+  : Request(endpointOrWorker, data::getRequestData(requestData), operationName, enablePythonFuture)
 {
   std::visit(data::dispatch{
                [this](data::AmSend amSend) {
                  if (_endpoint == nullptr)
                    throw ucxx::Error("An endpoint is required to send active messages");
                },
                [](data::AmReceive amReceive) {},
-               [](auto) { throw std::runtime_error("Unreachable"); },
              },
              requestData);
 
diff --git a/cpp/src/request_stream.cpp b/cpp/src/request_stream.cpp
@@ -12,10 +12,10 @@
 
 namespace ucxx {
 RequestStream::RequestStream(std::shared_ptr<Endpoint> endpoint,
-                             const data::RequestData requestData,
+                             const std::variant<data::StreamSend, data::StreamReceive> requestData,
                              const std::string operationName,
                              const bool enablePythonFuture)
-  : Request(endpoint, requestData, operationName, enablePythonFuture)
+  : Request(endpoint, data::getRequestData(requestData), operationName, enablePythonFuture)
 {
   std::visit(data::dispatch{
                [this](data::StreamSend streamSend) {
@@ -31,9 +31,10 @@ RequestStream::RequestStream(std::shared_ptr<Endpoint> endpoint,
              requestData);
 }
 
-std::shared_ptr<RequestStream> createRequestStream(std::shared_ptr<Endpoint> endpoint,
-                                                   const data::RequestData requestData,
-                                                   const bool enablePythonFuture = false)
+std::shared_ptr<RequestStream> createRequestStream(
+  std::shared_ptr<Endpoint> endpoint,
+  const std::variant<data::StreamSend, data::StreamReceive> requestData,
+  const bool enablePythonFuture = false)
 {
   std::shared_ptr<RequestStream> req =
     std::visit(data::dispatch{
diff --git a/cpp/src/request_tag.cpp b/cpp/src/request_tag.cpp
@@ -14,11 +14,12 @@
 
 namespace ucxx {
 
-std::shared_ptr<RequestTag> createRequestTag(std::shared_ptr<Component> endpointOrWorker,
-                                             const data::RequestData requestData,
-                                             const bool enablePythonFuture                = false,
-                                             RequestCallbackUserFunction callbackFunction = nullptr,
-                                             RequestCallbackUserData callbackData         = nullptr)
+std::shared_ptr<RequestTag> createRequestTag(
+  std::shared_ptr<Component> endpointOrWorker,
+  const std::variant<data::TagSend, data::TagReceive> requestData,
+  const bool enablePythonFuture                = false,
+  RequestCallbackUserFunction callbackFunction = nullptr,
+  RequestCallbackUserData callbackData         = nullptr)
 {
   std::shared_ptr<RequestTag> req =
     std::visit(data::dispatch{
@@ -40,7 +41,6 @@ std::shared_ptr<RequestTag> createRequestTag(std::shared_ptr<Component> endpoint
                                                                      callbackFunction,
                                                                      callbackData));
                  },
-                 [](auto) -> decltype(req) { throw std::runtime_error("Unreachable"); },
                },
                requestData);
 
@@ -54,20 +54,19 @@ std::shared_ptr<RequestTag> createRequestTag(std::shared_ptr<Component> endpoint
 }
 
 RequestTag::RequestTag(std::shared_ptr<Component> endpointOrWorker,
-                       const data::RequestData requestData,
+                       const std::variant<data::TagSend, data::TagReceive> requestData,
                        const std::string operationName,
                        const bool enablePythonFuture,
                        RequestCallbackUserFunction callbackFunction,
                        RequestCallbackUserData callbackData)
-  : Request(endpointOrWorker, requestData, operationName, enablePythonFuture)
+  : Request(endpointOrWorker, data::getRequestData(requestData), operationName, enablePythonFuture)
 {
   std::visit(data::dispatch{
                [this](data::TagSend tagSend) {
                  if (_endpoint == nullptr)
                    throw ucxx::Error("An endpoint is required to send tag messages");
                },
                [](data::TagReceive tagReceive) {},
-               [](auto) { throw std::runtime_error("Unreachable"); },
              },
              requestData);
 
diff --git a/cpp/src/request_tag_multi.cpp b/cpp/src/request_tag_multi.cpp
@@ -25,11 +25,12 @@ BufferRequest::BufferRequest() { ucxx_trace("BufferRequest created: %p", this);
 
 BufferRequest::~BufferRequest() { ucxx_trace("BufferRequest destroyed: %p", this); }
 
-RequestTagMulti::RequestTagMulti(std::shared_ptr<Endpoint> endpoint,
-                                 const data::RequestData requestData,
-                                 const std::string operationName,
-                                 const bool enablePythonFuture)
-  : Request(endpoint, requestData, operationName, enablePythonFuture)
+RequestTagMulti::RequestTagMulti(
+  std::shared_ptr<Endpoint> endpoint,
+  const std::variant<data::TagMultiSend, data::TagMultiReceive> requestData,
+  const std::string operationName,
+  const bool enablePythonFuture)
+  : Request(endpoint, data::getRequestData(requestData), operationName, enablePythonFuture)
 {
   auto worker = endpoint->getWorker();
   if (enablePythonFuture) _future = worker->getFuture();
@@ -53,9 +54,10 @@ RequestTagMulti::~RequestTagMulti()
   }
 }
 
-std::shared_ptr<RequestTagMulti> createRequestTagMulti(std::shared_ptr<Endpoint> endpoint,
-                                                       const data::RequestData requestData,
-                                                       const bool enablePythonFuture)
+std::shared_ptr<RequestTagMulti> createRequestTagMulti(
+  std::shared_ptr<Endpoint> endpoint,
+  const std::variant<data::TagMultiSend, data::TagMultiReceive> requestData,
+  const bool enablePythonFuture)
 {
   std::shared_ptr<RequestTagMulti> req =
     std::visit(data::dispatch{
@@ -71,7 +73,6 @@ std::shared_ptr<RequestTagMulti> createRequestTagMulti(std::shared_ptr<Endpoint>
                    req->recvCallback(UCS_OK);
                    return req;
                  },
-                 [](auto) -> decltype(req) { throw std::runtime_error("Unreachable"); },
                },
                requestData);
 
