Add proper bounded output support to set-algorithms from oneapi::dpl::ranges namespace with CPU policies

include/oneapi/dpl/pstl/parallel_backend_tbb.h

@@ -478,7 +478,7 @@ class __func_task : public __task
     execute()
     {
         return _M_func(this);
-    };
+    }
 
   public:
     template <typename _Fn>

include/oneapi/dpl/pstl/parallel_backend_utils.h

@@ -24,6 +24,7 @@
 #include <utility>
 #include <vector>
 #include <cassert>
+#include <cstdint> // for std::uint8_t (used, e.g., as underlying type of __parallel_set_op_mask)
 #include "utils.h"
 #include "memory_fwd.h"
 #include "functional_impl.h" // for oneapi::dpl::identity, std::invoke
@@ -217,48 +218,135 @@ struct __serial_move_merge
     }
 };
 
-template <typename _ForwardIterator1, typename _ForwardIterator2, typename _OutputIterator,
+template <bool _Bounded>
+struct _MaskSize;
+
+template <>
+struct _MaskSize<false>
+{
+    template <typename _DifferenceType1, typename _DifferenceType2>
+    std::common_type_t<_DifferenceType1, _DifferenceType2>
+    operator()(_DifferenceType1, _DifferenceType2) const
+    {
+        // For unbounded set operations, the maximum possible mask size is always zero
+        return 0;
+    }
+};
+
+template <>
+struct _MaskSize<true>
+{
+    template <typename _DifferenceType1, typename _DifferenceType2>
+    std::common_type_t<_DifferenceType1, _DifferenceType2>
+    operator()(_DifferenceType1 __n, _DifferenceType2 __m) const
+    {
+        using _DifferenceType = std::common_type_t<_DifferenceType1, _DifferenceType2>;
+
+        // For bounded set operations, the maximum possible mask size is the sum of sizes of both input ranges
+        return _DifferenceType{__n} + _DifferenceType{__m};
+    }
+};
+
+enum class __parallel_set_op_mask : std::uint8_t
+{
+    eData1 = 0x10,          // mask for first input data item usage
+    eData2 = 0x01,          // mask for second input data item usage
+    eBoth = eData1 | eData2 // mask for both input data items usage
+};
+
+inline std::nullptr_t
+__set_iterator_mask(std::nullptr_t, __parallel_set_op_mask)
+{
+    return nullptr;
+}
+
+inline __parallel_set_op_mask*
+__set_iterator_mask(__parallel_set_op_mask* __mask, __parallel_set_op_mask __state)
+{
+    *__mask = __state;
+    return ++__mask;
+}
+
+template <typename _Size>
+std::nullptr_t
+__set_iterator_mask_n(std::nullptr_t, __parallel_set_op_mask, _Size)
+{
+    return nullptr;
+}
+
+template <typename _Size>
+__parallel_set_op_mask*
+__set_iterator_mask_n(__parallel_set_op_mask* __mask, __parallel_set_op_mask __state, _Size __count)
+{
+    std::fill_n(__mask, __count, __state);
+    return __mask + __count;
+}
+
+template <typename _ForwardIterator1, typename _ForwardIterator2, typename _OutputIterator, typename _MaskIterator,
           typename _CopyConstructRange, typename _Compare, typename _Proj1, typename _Proj2>
-_OutputIterator
+std::tuple<_OutputIterator, _MaskIterator>
 __set_union_construct(_ForwardIterator1 __first1, _ForwardIterator1 __last1, _ForwardIterator2 __first2,
-                      _ForwardIterator2 __last2, _OutputIterator __result, _CopyConstructRange __cc_range,
-                      _Compare __comp, _Proj1 __proj1, _Proj2 __proj2)
+                      _ForwardIterator2 __last2, _OutputIterator __result, _MaskIterator __mask,
+                      _CopyConstructRange __cc_range, _Compare __comp, _Proj1 __proj1, _Proj2 __proj2)
 {
-    using _Tp = typename ::std::iterator_traits<_OutputIterator>::value_type;
+    assert(__first1 <= __last1);
+    assert(__first2 <= __last2);
+
+    using _Tp = typename std::iterator_traits<_OutputIterator>::value_type;
 
     for (; __first1 != __last1; ++__result)
     {
         if (__first2 == __last2)
-            return __cc_range(__first1, __last1, __result);
+        {
+            __mask = __set_iterator_mask_n(__mask, __parallel_set_op_mask::eData1, __last1 - __first1);
+            return {__cc_range(__first1, __last1, __result), __mask};
+        }
+
         if (std::invoke(__comp, std::invoke(__proj2, *__first2), std::invoke(__proj1, *__first1)))
         {
-            ::new (::std::addressof(*__result)) _Tp(*__first2);
+            new (std::addressof(*__result)) _Tp(*__first2);
             ++__first2;
+            __mask = __set_iterator_mask(__mask, __parallel_set_op_mask::eData2);
         }
         else
         {
-            ::new (::std::addressof(*__result)) _Tp(*__first1);
+            new (std::addressof(*__result)) _Tp(*__first1);
             if (!std::invoke(__comp, std::invoke(__proj1, *__first1), std::invoke(__proj2, *__first2)))
+            {
                 ++__first2;
+                __mask = __set_iterator_mask(__mask, __parallel_set_op_mask::eBoth);
+            }
+            else
+            {
+                __mask = __set_iterator_mask(__mask, __parallel_set_op_mask::eData1);
+            }
             ++__first1;
         }
     }
-    return __cc_range(__first2, __last2, __result);
+
+    __mask = __set_iterator_mask_n(__mask, __parallel_set_op_mask::eData2, __last2 - __first2);
+    return {__cc_range(__first2, __last2, __result), __mask};
 }
 
-template <typename _ForwardIterator1, typename _ForwardIterator2, typename _OutputIterator, typename _CopyFunc,
-          typename _CopyFromFirstSet, typename _Compare, typename _Proj1, typename _Proj2>
-_OutputIterator
+template <typename _ForwardIterator1, typename _ForwardIterator2, typename _OutputIterator, typename _MaskIterator,
+          typename _CopyFunc, typename _CopyFromFirstSet, typename _Compare, typename _Proj1, typename _Proj2>
+std::tuple<_OutputIterator, _MaskIterator>
 __set_intersection_construct(_ForwardIterator1 __first1, _ForwardIterator1 __last1, _ForwardIterator2 __first2,
-                             _ForwardIterator2 __last2, _OutputIterator __result, _CopyFunc _copy, _CopyFromFirstSet,
-                             _Compare __comp, _Proj1 __proj1, _Proj2 __proj2)
+                             _ForwardIterator2 __last2, _OutputIterator __result, _MaskIterator __mask, _CopyFunc _copy,
+                             _CopyFromFirstSet, _Compare __comp, _Proj1 __proj1, _Proj2 __proj2)
 {
     while (__first1 != __last1 && __first2 != __last2)
     {
         if (std::invoke(__comp, std::invoke(__proj1, *__first1), std::invoke(__proj2, *__first2)))
+        {
             ++__first1;
+            __mask = __set_iterator_mask(__mask, __parallel_set_op_mask::eData1);
+        }
         else if (std::invoke(__comp, std::invoke(__proj2, *__first2), std::invoke(__proj1, *__first1)))
+        {
             ++__first2;
+            __mask = __set_iterator_mask(__mask, __parallel_set_op_mask::eData2);
+        }
         else
         {
             if constexpr (_CopyFromFirstSet::value)
@@ -269,74 +357,102 @@ __set_intersection_construct(_ForwardIterator1 __first1, _ForwardIterator1 __las
             ++__first1;
             ++__first2;
             ++__result;
+            __mask = __set_iterator_mask(__mask, __parallel_set_op_mask::eBoth);
         }
     }
-    return __result;
+
+    // This needed to save in mask that we processed all data till the end
+    __mask = __set_iterator_mask_n(__mask, __parallel_set_op_mask::eData1, __last1 - __first1);
+    __mask = __set_iterator_mask_n(__mask, __parallel_set_op_mask::eData2, __last2 - __first2);
+
+    return {__result, __mask};
 }
 
-template <typename _ForwardIterator1, typename _ForwardIterator2, typename _OutputIterator,
+template <typename _ForwardIterator1, typename _ForwardIterator2, typename _OutputIterator, typename _MaskIterator,
           typename _CopyConstructRange, typename _Compare, typename _Proj1, typename _Proj2>
-_OutputIterator
+std::tuple<_OutputIterator, _MaskIterator>
 __set_difference_construct(_ForwardIterator1 __first1, _ForwardIterator1 __last1, _ForwardIterator2 __first2,
-                           _ForwardIterator2 __last2, _OutputIterator __result, _CopyConstructRange __cc_range,
-                           _Compare __comp, _Proj1 __proj1, _Proj2 __proj2)
+                           _ForwardIterator2 __last2, _OutputIterator __result, _MaskIterator __mask,
+                           _CopyConstructRange __cc_range, _Compare __comp, _Proj1 __proj1, _Proj2 __proj2)
 {
     using _Tp = typename ::std::iterator_traits<_OutputIterator>::value_type;
 
-    for (; __first1 != __last1;)
+    while (__first1 != __last1)
     {
         if (__first2 == __last2)
-            return __cc_range(__first1, __last1, __result);
+        {
+            __mask = __set_iterator_mask_n(__mask, __parallel_set_op_mask::eData1, __last1 - __first1);
+            return {__cc_range(__first1, __last1, __result), __mask};
+        }
 
         if (std::invoke(__comp, std::invoke(__proj1, *__first1), std::invoke(__proj2, *__first2)))
         {
-            ::new (::std::addressof(*__result)) _Tp(*__first1);
+            new (std::addressof(*__result)) _Tp(*__first1);
             ++__result;
             ++__first1;
+            __mask = __set_iterator_mask(__mask, __parallel_set_op_mask::eData1);
         }
         else
         {
             if (!std::invoke(__comp, std::invoke(__proj2, *__first2), std::invoke(__proj1, *__first1)))
+            {
                 ++__first1;
+                __mask = __set_iterator_mask(__mask, __parallel_set_op_mask::eBoth);
+            }
+            else
+            {
+                __mask = __set_iterator_mask(__mask, __parallel_set_op_mask::eData2);
+            }
             ++__first2;
         }
     }
-    return __result;
+
+    return {__result, __mask};
 }
 
-template <typename _ForwardIterator1, typename _ForwardIterator2, typename _OutputIterator,
+template <typename _ForwardIterator1, typename _ForwardIterator2, typename _OutputIterator, typename _MaskIterator,
           typename _CopyConstructRange, typename _Compare, typename _Proj1, typename _Proj2>
-_OutputIterator
+std::tuple<_OutputIterator, _MaskIterator>
 __set_symmetric_difference_construct(_ForwardIterator1 __first1, _ForwardIterator1 __last1, _ForwardIterator2 __first2,
-                                     _ForwardIterator2 __last2, _OutputIterator __result,
+                                     _ForwardIterator2 __last2, _OutputIterator __result, _MaskIterator __mask,
                                      _CopyConstructRange __cc_range, _Compare __comp, _Proj1 __proj1, _Proj2 __proj2)
 {
     using _Tp = typename ::std::iterator_traits<_OutputIterator>::value_type;
 
-    for (; __first1 != __last1;)
+    while (__first1 != __last1)
     {
         if (__first2 == __last2)
-            return __cc_range(__first1, __last1, __result);
+        {
+            __mask = __set_iterator_mask_n(__mask, __parallel_set_op_mask::eData1, __last1 - __first1);
+            return {__cc_range(__first1, __last1, __result), __mask};
+        }
 
         if (std::invoke(__comp, std::invoke(__proj1, *__first1), std::invoke(__proj2, *__first2)))
         {
-            ::new (::std::addressof(*__result)) _Tp(*__first1);
+            new (std::addressof(*__result)) _Tp(*__first1);
             ++__result;
             ++__first1;
+            __mask = __set_iterator_mask(__mask, __parallel_set_op_mask::eData1);
         }
         else
         {
             if (std::invoke(__comp, std::invoke(__proj2, *__first2), std::invoke(__proj1, *__first1)))
             {
-                ::new (::std::addressof(*__result)) _Tp(*__first2);
+                new (std::addressof(*__result)) _Tp(*__first2);
                 ++__result;
+                __mask = __set_iterator_mask(__mask, __parallel_set_op_mask::eData2);
             }
             else
+            {
                 ++__first1;
+                __mask = __set_iterator_mask(__mask, __parallel_set_op_mask::eBoth);
+            }
             ++__first2;
         }
     }
-    return __cc_range(__first2, __last2, __result);
+
+    __mask = __set_iterator_mask_n(__mask, __parallel_set_op_mask::eData2, __last2 - __first2);
+    return {__cc_range(__first2, __last2, __result), __mask};
 }
 
 template <template <typename, typename...> typename _Concrete, typename _ValueType, typename... _Args>
@@ -405,6 +521,46 @@ struct __enumerable_thread_local_storage_base
     const std::tuple<_Args...> __args;
 };
 
+template <typename _RandomAccessIterator1, typename _RandomAccessIterator2, typename _RandomAccessOutputIterator>
+struct __set_operations_result
+{
+    _RandomAccessIterator1 __in1;
+    _RandomAccessIterator2 __in2;
+    _RandomAccessOutputIterator __it_out;
+
+    // Get reached input1 and output iterators
+    template <typename TResult>
+    TResult
+    __get_reached_in1_out() const
+    {
+        return {__in1, __it_out};
+    }
+
+    // Get reached input1, input2 and output iterators
+    template <typename TResult>
+    TResult
+    __get_reached_in1_in2_out() const
+    {
+        return {__in1, __in2, __it_out};
+    }
+
+    // Get reached output iterator
+    _RandomAccessOutputIterator
+    __get_reached_out() const
+    {
+        return __it_out;
+    }
+
+    __set_operations_result<_RandomAccessIterator1, _RandomAccessIterator2, _RandomAccessOutputIterator>
+    operator+(std::tuple<typename std::iterator_traits<_RandomAccessIterator1>::difference_type,
+                         typename std::iterator_traits<_RandomAccessIterator2>::difference_type,
+                         typename std::iterator_traits<_RandomAccessOutputIterator>::difference_type>
+                  __offsets) const
+    {
+        return {__in1 + std::get<0>(__offsets), __in2 + std::get<1>(__offsets), __it_out + std::get<2>(__offsets)};
+    }
+};
+
 } // namespace __utils
 } // namespace dpl
 } // namespace oneapi

include/oneapi/dpl/pstl/utils_ranges.h

@@ -834,6 +834,30 @@ __get_subscription_view(_View&& __view)
 }
 #endif // _ONEDPL_CPP20_RANGES_PRESENT
 
+// Returns begin and end of the range
+template <typename _Range>
+auto
+__get_range_bounds(_Range&& __rng)
+{
+    const auto __size = oneapi::dpl::__ranges::__size(__rng);
+
+    auto __begin = oneapi::dpl::__ranges::__begin(__rng);
+
+    return std::make_tuple(__begin, __begin + __size);
+}
+
+// Returns begin, end and size of the range
+template <typename _Range>
+auto
+__get_range_bounds_n(_Range&& __rng)
+{
+    const auto __size = oneapi::dpl::__ranges::__size(__rng);
+
+    auto __begin = oneapi::dpl::__ranges::__begin(__rng);
+
+    return std::make_tuple(__begin, __begin + __size, __size);
+}
+
 } // namespace __ranges
 } // namespace dpl
 } // namespace oneapi