Make usage of mutable functors a bit more sane

Author: miscco

libcudacxx/include/cuda/__iterator/transform_iterator.h

@@ -215,10 +215,19 @@ class transform_iterator : public __transform_iterator_category_base<_Iter, _Fn>
   }
 
   //! @brief Invokes the stored functor with the value pointed to by the stored iterator
+  _CCCL_TEMPLATE(class _Iter2 = _Iter)
+  _CCCL_REQUIRES(_CUDA_VSTD::regular_invocable<const _Fn&, _CUDA_VSTD::iter_reference_t<const _Iter2>>)
   [[nodiscard]] _LIBCUDACXX_HIDE_FROM_ABI constexpr decltype(auto) operator*() const
-    noexcept(noexcept(_CUDA_VSTD::invoke(const_cast<_Fn&>(*__func_), *__current_)))
+    noexcept(noexcept(_CUDA_VSTD::invoke(*__func_, *__current_)))
   {
-    return _CUDA_VSTD::invoke(const_cast<_Fn&>(*__func_), *__current_);
+    return _CUDA_VSTD::invoke(*__func_, *__current_);
+  }
+
+  //! @brief Invokes the stored functor with the value pointed to by the stored iterator
+  [[nodiscard]] _LIBCUDACXX_HIDE_FROM_ABI constexpr decltype(auto)
+  operator*() noexcept(noexcept(_CUDA_VSTD::invoke(*__func_, *__current_)))
+  {
+    return _CUDA_VSTD::invoke(*__func_, *__current_);
   }
 
   //! @brief Increments the stored iterator
@@ -289,11 +298,23 @@ class transform_iterator : public __transform_iterator_category_base<_Iter, _Fn>
   //! @param __n The additional offset
   //! @returns _CUDA_VSTD::invoke(__func_, __current_[__n])
   _CCCL_TEMPLATE(class _Iter2 = _Iter)
-  _CCCL_REQUIRES(_CUDA_VSTD::random_access_iterator<_Iter2>)
+  _CCCL_REQUIRES(_CUDA_VSTD::random_access_iterator<_Iter2> _CCCL_AND
+                   _CUDA_VSTD::regular_invocable<const _Fn&, _CUDA_VSTD::iter_reference_t<const _Iter2>>)
   [[nodiscard]] _LIBCUDACXX_HIDE_FROM_ABI constexpr decltype(auto) operator[](difference_type __n) const
-    noexcept(__transform_iterator_nothrow_subscript<_Fn, _Iter2>)
+    noexcept(__transform_iterator_nothrow_subscript<const _Fn, _Iter2>)
+  {
+    return _CUDA_VSTD::invoke(*__func_, __current_[__n]);
+  }
+
+  //! @brief Subscripts the stored iterator by \p __n and applies the stored functor to the result
+  //! @param __n The additional offset
+  //! @returns _CUDA_VSTD::invoke(__func_, __current_[__n])
+  _CCCL_TEMPLATE(class _Iter2 = _Iter)
+  _CCCL_REQUIRES(_CUDA_VSTD::random_access_iterator<_Iter2>)
+  [[nodiscard]] _LIBCUDACXX_HIDE_FROM_ABI constexpr decltype(auto)
+  operator[](difference_type __n) noexcept(__transform_iterator_nothrow_subscript<_Fn, _Iter2>)
   {
-    return _CUDA_VSTD::invoke(const_cast<_Fn&>(*__func_), __current_[__n]);
+    return _CUDA_VSTD::invoke(*__func_, __current_[__n]);
   }
 
   //! @brief Compares two \c transform_iterator for equality, directly comparing the stored iterators

libcudacxx/test/libcudacxx/cuda/iterators/transform_iterator/deref.pass.cpp

@@ -27,6 +27,7 @@ __host__ __device__ constexpr void test()
   {
     cuda::transform_iterator iter{Iter{buffer}, PlusOne{}};
     assert(*iter == 1);
+    assert(*cuda::std::as_const(iter) == 1);
     assert(*buffer == 0);
 
     static_assert(!noexcept(*iter));
@@ -68,6 +69,17 @@ __host__ __device__ constexpr void test()
     static_assert(!noexcept(*iter));
     static_assert(cuda::std::is_same_v<int&&, decltype(*iter)>);
   }
+
+  {
+    cuda::transform_iterator iter{Iter{buffer}, PlusWithMutableMember{3}};
+    assert(*iter == 5);
+    assert(*iter == 6);
+    assert(*iter == 7);
+    assert(*buffer == 2);
+
+    static_assert(noexcept(*iter) == noexcept(*cuda::std::declval<Iter>()));
+    static_assert(cuda::std::is_same_v<int, decltype(*iter)>);
+  }
 }
 
 __host__ __device__ constexpr bool test()

libcudacxx/test/libcudacxx/cuda/iterators/transform_iterator/subscript.pass.cpp

@@ -31,6 +31,7 @@ __host__ __device__ constexpr void test()
     {
       cuda::transform_iterator iter{Iter{buffer}, PlusOne{}};
       assert(iter[4] == 5);
+      assert(cuda::std::as_const(iter)[4] == 5);
       assert(buffer[4] == 4);
 
       static_assert(!noexcept(iter[4]));
@@ -72,6 +73,17 @@ __host__ __device__ constexpr void test()
       static_assert(!noexcept(iter[4]));
       static_assert(cuda::std::is_same_v<int&&, decltype(iter[4])>);
     }
+
+    {
+      cuda::transform_iterator iter{Iter{buffer}, PlusWithMutableMember{3}};
+      assert(iter[4] == 9);
+      assert(iter[4] == 10);
+      assert(iter[4] == 11);
+      assert(buffer[4] == 6);
+
+      static_assert(noexcept(iter[4]) == noexcept(*cuda::std::declval<Iter>()));
+      static_assert(cuda::std::is_same_v<int, decltype(iter[4])>);
+    }
   }
   else
   {

libcudacxx/test/libcudacxx/cuda/iterators/transform_iterator/types.h

@@ -63,4 +63,16 @@ struct PlusOneNoexcept
   }
 };
 
+struct PlusWithMutableMember
+{
+  int val_ = 0;
+  __host__ __device__ constexpr PlusWithMutableMember(const int val) noexcept
+      : val_(val)
+  {}
+  __host__ __device__ constexpr int operator()(int x) noexcept
+  {
+    return x + val_++;
+  }
+};
+
 #endif // TEST_CUDA_TRANSFORM_ITERATOR_TYPES_H