Refactor Thrust iterators 1/n (#3839)

* Refactor constant_iterator * Refactor counting_iterator
NVIDIA · Feb 19, 2025 · fd86263 · fd86263
1 parent 297447e
commit fd86263
Show file tree

Hide file tree

Showing 6 changed files with 413 additions and 538 deletions.
diff --git a/thrust/testing/copy.cu b/thrust/testing/copy.cu
@@ -717,15 +717,30 @@ template <>
 struct is_non_const_reference<only_set_when_expected_it> : thrust::true_type
 {};
 } // end namespace detail
+THRUST_NAMESPACE_END
 
+namespace std
+{
 template <>
 struct iterator_traits<only_set_when_expected_it>
 {
   using value_type        = long long;
   using reference         = only_set_when_expected_it;
   using iterator_category = thrust::random_access_device_iterator_tag;
+  using difference_type   = ::cuda::std::ptrdiff_t;
 };
-THRUST_NAMESPACE_END
+} // namespace std
+
+_LIBCUDACXX_BEGIN_NAMESPACE_STD
+template <>
+struct iterator_traits<only_set_when_expected_it>
+{
+  using value_type        = long long;
+  using reference         = only_set_when_expected_it;
+  using iterator_category = thrust::random_access_device_iterator_tag;
+  using difference_type   = ::cuda::std::ptrdiff_t;
+};
+_LIBCUDACXX_END_NAMESPACE_STD
 
 void TestCopyWithBigIndexesHelper(int magnitude)
 {

diff --git a/thrust/testing/scan.cu b/thrust/testing/scan.cu
@@ -565,24 +565,27 @@ struct only_set_when_expected_it
   }
 };
 
-THRUST_NAMESPACE_BEGIN
+namespace std
+{
 template <>
 struct iterator_traits<only_set_when_expected_it>
 {
-  using value_type = long long;
-  using reference  = only_set_when_expected_it;
+  using value_type      = long long;
+  using reference       = only_set_when_expected_it;
+  using difference_type = ::cuda::std::ptrdiff_t;
 };
-THRUST_NAMESPACE_END
+} // namespace std
 
-namespace std
-{
+_LIBCUDACXX_BEGIN_NAMESPACE_STD
 template <>
 struct iterator_traits<only_set_when_expected_it>
 {
-  using value_type = long long;
-  using reference  = only_set_when_expected_it;
+  using value_type        = long long;
+  using reference         = only_set_when_expected_it;
+  using iterator_category = thrust::random_access_device_iterator_tag;
+  using difference_type   = ::cuda::std::ptrdiff_t;
 };
-} // namespace std
+_LIBCUDACXX_END_NAMESPACE_STD
 
 void TestInclusiveScanWithBigIndexesHelper(int magnitude)
 {