Fix proclaim_return_type copy constructor (#386)

NVIDIA · Mar 9, 2023 · 34fd806 · 34fd806
1 parent 395dd3b
commit 34fd806
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 13 deletions.
diff --git a/.upstream-tests/test/cuda/proclaim_return_type.pass.cpp b/.upstream-tests/test/cuda/proclaim_return_type.pass.cpp
@@ -19,17 +19,21 @@ __host__ __device__
 void test_proclaim_return_type(Fn&& fn, T expected, As... as)
 {
   {
-    auto f = cuda::proclaim_return_type<T>(cuda::std::forward<Fn>(fn));
+    auto f1 = cuda::proclaim_return_type<T>(cuda::std::forward<Fn>(fn));
 
-    assert(f(as...) == expected);
-    assert(cuda::std::move(f)(as...) == expected);
+    assert(f1(as...) == expected);
+
+    auto f2{f1};
+    assert(cuda::std::move(f2)(as...) == expected);
   }
 
   {
-    const auto f = cuda::proclaim_return_type<T>(fn);
+    const auto f1 = cuda::proclaim_return_type<T>(fn);
+
+    assert(f1(as...) == expected);
 
-    assert(f(as...) == expected);
-    assert(cuda::std::move(f)(as...) == expected);
+    auto f2{f1};
+    assert(cuda::std::move(f2)(as...) == expected);
   }
 }
 

diff --git a/include/cuda/functional b/include/cuda/functional
@@ -71,10 +71,14 @@ class __return_type_wrapper {
  public:
   __return_type_wrapper() = delete;
 
-  template <class _Fn>
-  _LIBCUDACXX_CONSTEXPR_AFTER_CXX11 _LIBCUDACXX_INLINE_VISIBILITY
-  explicit __return_type_wrapper(_Fn&& __fn) noexcept
-      : __fn_(_CUDA_VSTD::forward<_Fn>(__fn)) {}
+  _LIBCUDACXX_CONSTEXPR_AFTER_CXX11 
+  __return_type_wrapper(__return_type_wrapper const&) = default;
+
+  template <class _Fn,
+            class = _CUDA_VSTD::__enable_if_t<_CUDA_VSTD::is_same<_CUDA_VSTD::__decay_t<_Fn>, _DecayFn>::value>>
+  _LIBCUDACXX_CONSTEXPR_AFTER_CXX11 _LIBCUDACXX_INLINE_VISIBILITY 
+  explicit __return_type_wrapper(_Fn &&__fn) noexcept
+    : __fn_(_CUDA_VSTD::forward<_Fn>(__fn)) {}
 
   template <class... _As>
   _LIBCUDACXX_CONSTEXPR_AFTER_CXX11 _LIBCUDACXX_INLINE_VISIBILITY _Ret
@@ -143,10 +147,9 @@ class __return_type_wrapper {
 
 template <class _Ret, class _Fn>
 inline _LIBCUDACXX_INLINE_VISIBILITY
-__detail::__return_type_wrapper<_Ret, typename _CUDA_VSTD::decay<_Fn>::type>
+__detail::__return_type_wrapper<_Ret, _CUDA_VSTD::__decay_t<_Fn>>
 proclaim_return_type(_Fn&& __fn) noexcept {
-  return __detail::__return_type_wrapper<_Ret,
-                                         typename _CUDA_VSTD::decay<_Fn>::type>(
+  return __detail::__return_type_wrapper<_Ret, _CUDA_VSTD::__decay_t<_Fn>>(
       _CUDA_VSTD::forward<_Fn>(__fn));
 }
 _LIBCUDACXX_END_NAMESPACE_CUDA