use thrust iterator categories in iterator wrapper

Author: upsj

testing/unittest/iterator_helpers.h

@@ -1,21 +1,30 @@
 #pragma once
 
 #include <thrust/iterator/iterator_traits.h>
+#include <thrust/iterator/iterator_categories.h>
+#include <type_traits>
 
 
 // Wraps an existing iterator into a forward iterator,
 // thus removing some of its functionality
 template <typename Iterator>
 struct forward_iterator_wrapper {
     // LegacyIterator requirements
+    using iterator_system_tag = typename thrust::iterator_system<Iterator>::type;
     using reference = typename thrust::iterator_traits<Iterator>::reference;
     using pointer = typename thrust::iterator_traits<Iterator>::pointer;
     using value_type = typename thrust::iterator_traits<Iterator>::value_type;
     using difference_type = typename thrust::iterator_traits<Iterator>::difference_type;
-    using iterator_category = std::forward_iterator_tag;
+    using iterator_category = typename std::conditional<
+        std::is_convertible<iterator_system_tag, thrust::device_system_tag>::value,
+        thrust::forward_device_iterator_tag,
+        typename std::conditional<
+            std::is_convertible<iterator_system_tag, thrust::host_system_tag>::value,
+            thrust::forward_host_iterator_tag,
+            std::forward_iterator_tag>::type>::type;
     using base_iterator_category = typename thrust::iterator_traits<Iterator>::iterator_category;
     static_assert(
-        std::is_convertible<base_iterator_category, std::forward_iterator_tag>::value, 
+        std::is_convertible<base_iterator_category, std::forward_iterator_tag>::value,
         "Cannot create forward_iterator_wrapper around an iterator that is not itself at least a forward iterator");
 
     __host__ __device__ reference operator*() const {