From 0ce057a16526e4f549ebb73057100795f11f11c2 Mon Sep 17 00:00:00 2001
From: Denis Yaroshevskiy <dyaroshev@meta.com>
Date: Thu, 10 Oct 2024 04:29:40 -0700
Subject: [PATCH] ] change contains interface to be more restrictive

Summary:
X-link: https://github.com/facebook/folly/pull/2309

previous interface allowed for implicit conversions.
So it so happens that the user would have to guard against it.

I think I better guard here.

Reviewed By: yfeldblum

Differential Revision: D63979345

fbshipit-source-id: dacc8037172467f308e81a147cdbe223b838da60
---
 .../folly/src/folly/algorithm/simd/Contains.h | 54 +++++++++++++++---
 .../src/folly/algorithm/simd/detail/Traits.h  |  4 +-
 .../algorithm/simd/test/ContainsTest.cpp      | 57 +++++++++++++++++++
 3 files changed, 104 insertions(+), 11 deletions(-)

diff --git a/third-party/folly/src/folly/algorithm/simd/Contains.h b/third-party/folly/src/folly/algorithm/simd/Contains.h
index 00756bb75ecd89..76980854147cb0 100644
--- a/third-party/folly/src/folly/algorithm/simd/Contains.h
+++ b/third-party/folly/src/folly/algorithm/simd/Contains.h
@@ -38,6 +38,40 @@ template <typename R>
 using std_range_value_t = typename std::iterator_traits<decltype(std::begin(
     std::declval<R&>()))>::value_type;
 
+// Constexpr check that we can always safely cast from From to To.
+// If we don't require this, we might silently get different semantics from
+// standard algorithms.
+template <typename From, typename To>
+constexpr bool convertible_with_no_loss() {
+  if (sizeof(From) > sizeof(To)) {
+    return false;
+  }
+  if (std::is_signed_v<From>) {
+    return std::is_signed_v<To>;
+  }
+
+  return std::is_unsigned_v<To> || sizeof(From) < sizeof(To);
+}
+
+// All the requirements to call contains(haystack, needle);
+//  * both are simd friendly (contigious range, primitive types)
+//  * integrals only
+//  * needle can be converted to the value_type of haystack and
+//    the result of equality comparison will be the same.
+template <typename R, typename T>
+constexpr bool contains_haystack_needle_test() {
+  if constexpr (!std::is_invocable_v<AsSimdFriendlyUintFn, R>) {
+    return false;
+  } else if constexpr (!has_integral_simd_friendly_equivalent_scalar_v<T>) {
+    return false;
+  } else {
+    using simd_haystack_value =
+        simd_friendly_equivalent_scalar_t<std_range_value_t<R>>;
+    using simd_needle = simd_friendly_equivalent_scalar_t<T>;
+    return convertible_with_no_loss<simd_needle, simd_haystack_value>();
+  }
+}
+
 } // namespace detail
 
 /**
@@ -53,23 +87,25 @@ using std_range_value_t = typename std::iterator_traits<decltype(std::begin(
 struct contains_fn {
   template <
       typename R,
-      typename = std::enable_if_t<
-          std::is_invocable_v<detail::AsSimdFriendlyUintFn, R>>>
-  FOLLY_ERASE bool operator()(R&& r, detail::std_range_value_t<R> x) const {
+      typename T,
+      typename =
+          std::enable_if_t<detail::contains_haystack_needle_test<R, T>()>>
+  FOLLY_ERASE bool operator()(R&& r, T x) const {
     auto castR = detail::asSimdFriendlyUint(folly::span(r));
-    auto castX = detail::asSimdFriendlyUint(x);
+    using value_type = detail::std_range_value_t<decltype(castR)>;
 
-    using T = decltype(castX);
+    auto castX = static_cast<value_type>(x);
 
-    if constexpr (std::is_same_v<T, std::uint8_t>) {
+    if constexpr (std::is_same_v<value_type, std::uint8_t>) {
       return detail::containsU8(castR, castX);
-    } else if constexpr (std::is_same_v<T, std::uint16_t>) {
+    } else if constexpr (std::is_same_v<value_type, std::uint16_t>) {
       return detail::containsU16(castR, castX);
-    } else if constexpr (std::is_same_v<T, std::uint32_t>) {
+    } else if constexpr (std::is_same_v<value_type, std::uint32_t>) {
       return detail::containsU32(castR, castX);
     } else {
       static_assert(
-          std::is_same_v<T, std::uint64_t>, "internal error, unknown type");
+          std::is_same_v<value_type, std::uint64_t>,
+          "internal error, unknown type");
       return detail::containsU64(castR, castX);
     }
   }
diff --git a/third-party/folly/src/folly/algorithm/simd/detail/Traits.h b/third-party/folly/src/folly/algorithm/simd/detail/Traits.h
index f3e34cb7eeeb7a..dc9f1c8a63e202 100644
--- a/third-party/folly/src/folly/algorithm/simd/detail/Traits.h
+++ b/third-party/folly/src/folly/algorithm/simd/detail/Traits.h
@@ -53,13 +53,13 @@ using simd_friendly_equivalent_scalar_t = std::enable_if_t<
     like_t<T, decltype(findSimdFriendlyEquivalent<std::remove_const_t<T>>())>>;
 
 template <typename T>
-constexpr bool has_integral_simd_friendly_equivalent_scalar =
+constexpr bool has_integral_simd_friendly_equivalent_scalar_v =
     std::is_integral_v< // void will return false
         decltype(findSimdFriendlyEquivalent<std::remove_const_t<T>>())>;
 
 template <typename T>
 using unsigned_simd_friendly_equivalent_scalar_t = std::enable_if_t<
-    has_integral_simd_friendly_equivalent_scalar<T>,
+    has_integral_simd_friendly_equivalent_scalar_v<T>,
     like_t<T, uint_bits_t<sizeof(T) * 8>>>;
 
 template <typename R>
diff --git a/third-party/folly/src/folly/algorithm/simd/test/ContainsTest.cpp b/third-party/folly/src/folly/algorithm/simd/test/ContainsTest.cpp
index 1838e9a83eb156..4bf594e4fb7f94 100644
--- a/third-party/folly/src/folly/algorithm/simd/test/ContainsTest.cpp
+++ b/third-party/folly/src/folly/algorithm/simd/test/ContainsTest.cpp
@@ -20,6 +20,9 @@
 
 #include <folly/portability/GTest.h>
 
+#include <list>
+#include <vector>
+
 namespace folly::simd {
 
 static_assert( //
@@ -34,6 +37,60 @@ static_assert( //
         std::vector<int>&,
         int>);
 
+static_assert( //
+    std::is_invocable_v< //
+        folly::simd::contains_fn,
+        std::vector<int>&,
+        int>);
+
+static_assert( //
+    std::is_invocable_v< //
+        folly::simd::contains_fn,
+        std::vector<int>&,
+        std::int16_t>);
+
+static_assert( //
+    std::is_invocable_v< //
+        folly::simd::contains_fn,
+        std::vector<int>&,
+        std::uint16_t>);
+
+static_assert( //
+    !std::is_invocable_v< //
+        folly::simd::contains_fn,
+        std::vector<int>&,
+        std::uint32_t>);
+
+static_assert( //
+    !std::is_invocable_v< //
+        folly::simd::contains_fn,
+        std::vector<int>&,
+        std::int64_t>);
+
+static_assert( //
+    !std::is_invocable_v< //
+        folly::simd::contains_fn,
+        std::vector<std::uint32_t>&,
+        std::int16_t>);
+
+static_assert( //
+    std::is_invocable_v< //
+        folly::simd::contains_fn,
+        std::vector<std::uint32_t>&,
+        std::uint16_t>);
+
+static_assert( //
+    !std::is_invocable_v< //
+        folly::simd::contains_fn,
+        std::list<std::int32_t>&,
+        std::int32_t>);
+
+static_assert( //
+    !std::is_invocable_v< //
+        folly::simd::contains_fn,
+        const std::vector<std::vector<std::int32_t>>&,
+        std::vector<std::int32_t>>);
+
 template <typename T>
 struct ContainsTest : ::testing::Test {};