Fix transform iterator for non-copy-constructible types (#3542)

Fixes: #3541

Co-authored-by: Michael Schellenberger Costa <[email protected]>

thrust/testing/cuda/transform_iterator.cmake

@@ -0,0 +1,12 @@
+target_compile_options(${test_target} PRIVATE $<$<COMPILE_LANG_AND_ID:CUDA,NVIDIA>: --extended-lambda>)
+
+# this check is actually not correct, because we must check the host compiler, not the CXX compiler.
+# We rely on that those are usually the same ;)
+if ("Clang" STREQUAL "${CMAKE_CXX_COMPILER_ID}" AND CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 13)
+    # When clang >= 13 is used as host compiler, we get the following warning:
+    #   nvcc_internal_extended_lambda_implementation:312:22: error: definition of implicit copy constructor for '__nv_hdl_wrapper_t<false, true, false, __nv_dl_tag<void (*)(), &TestAddressStabilityLambda, 2>, int (const int &)>' is deprecated because it has a user-declared copy assignment operator [-Werror,-Wdeprecated-copy]
+    #   312 | __nv_hdl_wrapper_t & operator=(const __nv_hdl_wrapper_t &in) = delete;
+    #   |                      ^
+    # Let's suppress it until NVBug 4980157 is resolved.
+    target_compile_options(${test_target} PRIVATE $<$<COMPILE_LANG_AND_ID:CUDA,NVIDIA>: -Wno-deprecated-copy>)
+endif ()

thrust/testing/cuda/transform_iterator.cu

@@ -0,0 +1,19 @@
+#include <thrust/device_vector.h>
+#include <thrust/host_vector.h>
+#include <thrust/iterator/transform_iterator.h>
+#include <thrust/logical.h>
+
+#include <unittest/unittest.h>
+
+// see also: https://github.com/NVIDIA/cccl/issues/3541
+void TestTransformWithLambda()
+{
+  auto l = [] __host__ __device__(int v) { return v < 4; };
+  thrust::host_vector<int> A{1, 2, 3, 4, 5, 6, 7};
+  ASSERT_EQUAL(thrust::any_of(A.begin(), A.end(), l), true);
+
+  thrust::device_vector<int> B{1, 2, 3, 4, 5, 6, 7};
+  ASSERT_EQUAL(thrust::any_of(B.begin(), B.end(), l), true);
+}
+
+DECLARE_UNITTEST(TestTransformWithLambda);

thrust/testing/transform_iterator.cu

@@ -2,6 +2,7 @@
 #include <thrust/functional.h>
 #include <thrust/iterator/counting_iterator.h>
 #include <thrust/iterator/transform_iterator.h>
+#include <thrust/logical.h>
 #include <thrust/reduce.h>
 #include <thrust/sequence.h>
 

thrust/thrust/iterator/transform_iterator.h

@@ -47,6 +47,9 @@
 #include <thrust/iterator/iterator_facade.h>
 #include <thrust/iterator/iterator_traits.h>
 
+#include <cuda/std/__memory/construct_at.h>
+#include <cuda/std/type_traits>
+
 THRUST_NAMESPACE_BEGIN
 
 /*! \addtogroup iterators
@@ -238,7 +241,26 @@ class transform_iterator
       , m_f(other.functor())
   {}
 
-  transform_iterator& operator=(const transform_iterator&) = default;
+  _CCCL_HOST_DEVICE transform_iterator& operator=(transform_iterator const& other)
+  {
+    super_t::operator=(other);
+    if constexpr (_CCCL_TRAIT(::cuda::std::is_copy_assignable, AdaptableUnaryFunction))
+    {
+      m_f = other.m_f;
+    }
+    else if constexpr (_CCCL_TRAIT(::cuda::std::is_copy_constructible, AdaptableUnaryFunction))
+    {
+      ::cuda::std::__destroy_at(&m_f);
+      ::cuda::std::__construct_at(&m_f, other.m_f);
+    }
+    else
+    {
+      static_assert(_CCCL_TRAIT(::cuda::std::is_copy_constructible, AdaptableUnaryFunction),
+                    "Cannot use thrust::transform_iterator with a functor that is neither copy constructible nor "
+                    "copy assignable");
+    }
+    return *this;
+  }
 
   /*! This method returns a copy of this \p transform_iterator's \c AdaptableUnaryFunction.
    *  \return A copy of this \p transform_iterator's \c AdaptableUnaryFunction.