[pstl] A hot fix for exclusive_scan (+ lost enable_if in declaration)

It fixes an ambiguity issue in case of a user has a custom policy and
calls a version of exclusive_scan with binary operation.

Differential Revision: https://reviews.llvm.org/D62719

3 files changed, 52 insertions, 6 deletions

diff --git a/pstl/include/pstl/internal/glue_numeric_defs.h b/pstl/include/pstl/internal/glue_numeric_defs.h
index f997b542362e..86cd38b34b13 100644
--- a/pstl/include/pstl/internal/glue_numeric_defs.h
+++ b/pstl/include/pstl/internal/glue_numeric_defs.h
@@ -60,7 +60,7 @@ exclusive_scan(_ExecutionPolicy&& __exec, _ForwardIterator1 __first, _ForwardIte
                _ForwardIterator2 __result, _Tp __init);
 
 template <class _ExecutionPolicy, class _ForwardIterator1, class _ForwardIterator2, class _Tp, class _BinaryOperation>
-_ForwardIterator2
+__pstl::__internal::__enable_if_execution_policy<_ExecutionPolicy, _ForwardIterator2>
 exclusive_scan(_ExecutionPolicy&& __exec, _ForwardIterator1 __first, _ForwardIterator1 __last,
                _ForwardIterator2 __result, _Tp __init, _BinaryOperation __binary_op);
 
diff --git a/pstl/include/pstl/internal/glue_numeric_impl.h b/pstl/include/pstl/internal/glue_numeric_impl.h
index 18b337133515..241b8a3bebc7 100644
--- a/pstl/include/pstl/internal/glue_numeric_impl.h
+++ b/pstl/include/pstl/internal/glue_numeric_impl.h
@@ -101,17 +101,25 @@ __pstl::__internal::__enable_if_execution_policy<_ExecutionPolicy, _ForwardItera
 exclusive_scan(_ExecutionPolicy&& __exec, _ForwardIterator1 __first, _ForwardIterator1 __last,
                _ForwardIterator2 __result, _Tp __init)
 {
-    return transform_exclusive_scan(std::forward<_ExecutionPolicy>(__exec), __first, __last, __result, __init,
-                                    std::plus<_Tp>(), __pstl::__internal::__no_op());
+    using namespace __pstl;
+    return __internal::__pattern_transform_scan(
+        std::forward<_ExecutionPolicy>(__exec), __first, __last, __result, __pstl::__internal::__no_op(), __init,
+        std::plus<_Tp>(), /*inclusive=*/std::false_type(),
+        __internal::__is_vectorization_preferred<_ExecutionPolicy, _ForwardIterator1, _ForwardIterator2>(__exec),
+        __internal::__is_parallelization_preferred<_ExecutionPolicy, _ForwardIterator1, _ForwardIterator2>(__exec));
 }
 
 template <class _ExecutionPolicy, class _ForwardIterator1, class _ForwardIterator2, class _Tp, class _BinaryOperation>
-_ForwardIterator2
+__pstl::__internal::__enable_if_execution_policy<_ExecutionPolicy, _ForwardIterator2>
 exclusive_scan(_ExecutionPolicy&& __exec, _ForwardIterator1 __first, _ForwardIterator1 __last,
                _ForwardIterator2 __result, _Tp __init, _BinaryOperation __binary_op)
 {
-    return transform_exclusive_scan(std::forward<_ExecutionPolicy>(__exec), __first, __last, __result, __init,
-                                    __binary_op, __pstl::__internal::__no_op());
+    using namespace __pstl;
+    return __internal::__pattern_transform_scan(
+        std::forward<_ExecutionPolicy>(__exec), __first, __last, __result, __pstl::__internal::__no_op(), __init,
+        __binary_op, /*inclusive=*/std::false_type(),
+        __internal::__is_vectorization_preferred<_ExecutionPolicy, _ForwardIterator1, _ForwardIterator2>(__exec),
+        __internal::__is_parallelization_preferred<_ExecutionPolicy, _ForwardIterator1, _ForwardIterator2>(__exec));
 }
 
 // [inclusive.scan]
diff --git a/pstl/test/std/numerics/numeric.ops/scan.fail.cpp b/pstl/test/std/numerics/numeric.ops/scan.fail.cpp
new file mode 100644
index 000000000000..13b2c9b751fe
--- /dev/null
+++ b/pstl/test/std/numerics/numeric.ops/scan.fail.cpp
@@ -0,0 +1,38 @@
+// -*- C++ -*-
+//===-- scan.fail.cpp -----------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+// UNSUPPORTED: c++98, c++03, c++11, c++14
+
+#include <execution>
+#include <numeric>
+
+struct CustomPolicy
+{
+    constexpr std::false_type
+    __allow_vector()
+    {
+        return std::false_type{};
+    }
+    constexpr std::false_type
+    __allow_parallel()
+    {
+        return std::false_type{};
+    }
+} policy;
+
+int32_t
+main()
+{
+    int *first = nullptr, *last = nullptr, *result = nullptr;
+
+    std::exclusive_scan(policy, first, last, result, 0); // expected-error {{no matching function for call to 'exclusive_scan'}}
+    std::exclusive_scan(policy, first, last, result, 0, std::plus<int>()); // expected-error {{no matching function for call to 'exclusive_scan'}}
+
+    return 0;
+}