Properly calculate the difference of offset_iterator with different bases

Author: miscco

libcudacxx/include/cuda/__iterator/offset_iterator.h

@@ -317,7 +317,7 @@ class offset_iterator
   [[nodiscard]] _LIBCUDACXX_HIDE_FROM_ABI friend constexpr difference_type
   operator-(const offset_iterator& __lhs, const offset_iterator& __rhs)
   {
-    return static_cast<difference_type>(__rhs.offset() - __lhs.offset());
+    return (__rhs.__iter_ + __rhs.offset()) - (__lhs.__iter_ + __lhs.offset());
   }
 
   //! @brief Returns the offset between an \c offset_iterators and default sentinel, equivalent to -offset()

libcudacxx/test/libcudacxx/cuda/iterators/offset_iterator/minus.iter.pass.cpp

@@ -19,7 +19,7 @@ __host__ __device__ constexpr bool test()
 {
   int buffer[8] = {1, 2, 3, 4, 5, 6, 7, 8};
 
-  {
+  { // Iterators with same base
     using offset_iterator = cuda::offset_iterator<int*>;
     const int offset1     = 4;
     const int offset2     = 2;
@@ -33,7 +33,7 @@ __host__ __device__ constexpr bool test()
     static_assert(cuda::std::is_same_v<decltype(iter1 - iter2), cuda::std::iter_difference_t<int*>>);
   }
 
-  {
+  { // const iterators with same base
     using offset_iterator = cuda::offset_iterator<int*>;
     const int offset1     = 4;
     const int offset2     = 2;
@@ -47,30 +47,90 @@ __host__ __device__ constexpr bool test()
     static_assert(cuda::std::is_same_v<decltype(iter1 - iter2), cuda::std::iter_difference_t<int*>>);
   }
 
-  {
+  { // Iterators with different base
+    using offset_iterator = cuda::offset_iterator<int*>;
+    const int offset1     = 4;
+    const int offset2     = 2;
+    const int base_offset = 3;
+    offset_iterator iter1(buffer + base_offset, offset1);
+    offset_iterator iter2(buffer, offset2);
+    assert(iter1 - iter2 == offset2 - (offset1 + base_offset));
+    assert(iter2 - iter1 == (offset1 + base_offset) - offset2);
+    assert(iter1.offset() == offset1);
+    assert(iter2.offset() == offset2);
+
+    static_assert(cuda::std::is_same_v<decltype(iter1 - iter2), cuda::std::iter_difference_t<int*>>);
+  }
+
+  { // const iterators with different base
+    using offset_iterator = cuda::offset_iterator<int*>;
+    const int offset1     = 4;
+    const int offset2     = 2;
+    const int base_offset = 3;
+    const offset_iterator iter1(buffer + base_offset, offset1);
+    const offset_iterator iter2(buffer, offset2);
+    assert(iter1 - iter2 == offset2 - (offset1 + base_offset));
+    assert(iter2 - iter1 == (offset1 + base_offset) - offset2);
+    assert(iter1.offset() == offset1);
+    assert(iter2.offset() == offset2);
+
+    static_assert(cuda::std::is_same_v<decltype(iter1 - iter2), cuda::std::iter_difference_t<int*>>);
+  }
+
+  { // Iterators with same base and offset array
     using offset_iterator = cuda::offset_iterator<int*, random_access_iterator<const int*>>;
     const int offset1[]   = {4};
     const int offset2[]   = {2};
     offset_iterator iter1(buffer, random_access_iterator<const int*>{offset1});
     offset_iterator iter2(buffer, random_access_iterator<const int*>{offset2});
-    assert(iter1 - iter2 == -2);
-    assert(iter2 - iter1 == 2);
-    assert(iter1.offset() == 4);
-    assert(iter2.offset() == 2);
+    assert(iter1 - iter2 == (*offset2 - *offset1));
+    assert(iter2 - iter1 == (*offset1 - *offset2));
+    assert(iter1.offset() == *offset1);
+    assert(iter2.offset() == *offset2);
 
     static_assert(cuda::std::is_same_v<decltype(iter1 - iter2), cuda::std::iter_difference_t<int*>>);
   }
 
-  {
+  { // const iterators with same base and offset array
     using offset_iterator = cuda::offset_iterator<int*, random_access_iterator<const int*>>;
     const int offset1[]   = {4};
     const int offset2[]   = {2};
     const offset_iterator iter1(buffer, random_access_iterator<const int*>{offset1});
     const offset_iterator iter2(buffer, random_access_iterator<const int*>{offset2});
-    assert(iter1 - iter2 == -2);
-    assert(iter2 - iter1 == 2);
-    assert(iter1.offset() == 4);
-    assert(iter2.offset() == 2);
+    assert(iter1 - iter2 == (*offset2 - *offset1));
+    assert(iter2 - iter1 == (*offset1 - *offset2));
+    assert(iter1.offset() == *offset1);
+    assert(iter2.offset() == *offset2);
+
+    static_assert(cuda::std::is_same_v<decltype(iter1 - iter2), cuda::std::iter_difference_t<int*>>);
+  }
+
+  { // Iterators with different base and offset array
+    using offset_iterator = cuda::offset_iterator<int*, random_access_iterator<const int*>>;
+    const int offset1[]   = {4};
+    const int offset2[]   = {2};
+    const int base_offset = 3;
+    offset_iterator iter1(buffer + base_offset, random_access_iterator<const int*>{offset1});
+    offset_iterator iter2(buffer, random_access_iterator<const int*>{offset2});
+    assert(iter1 - iter2 == (*offset2 - (*offset1 + base_offset)));
+    assert(iter2 - iter1 == (*offset1 + base_offset) - *offset2);
+    assert(iter1.offset() == *offset1);
+    assert(iter2.offset() == *offset2);
+
+    static_assert(cuda::std::is_same_v<decltype(iter1 - iter2), cuda::std::iter_difference_t<int*>>);
+  }
+
+  { // const iterators with different base and offset array
+    using offset_iterator = cuda::offset_iterator<int*, random_access_iterator<const int*>>;
+    const int offset1[]   = {4};
+    const int offset2[]   = {2};
+    const int base_offset = 3;
+    const offset_iterator iter1(buffer + base_offset, random_access_iterator<const int*>{offset1});
+    const offset_iterator iter2(buffer, random_access_iterator<const int*>{offset2});
+    assert(iter1 - iter2 == (*offset2 - (*offset1 + base_offset)));
+    assert(iter2 - iter1 == (*offset1 + base_offset) - *offset2);
+    assert(iter1.offset() == *offset1);
+    assert(iter2.offset() == *offset2);
 
     static_assert(cuda::std::is_same_v<decltype(iter1 - iter2), cuda::std::iter_difference_t<int*>>);
   }