From 59978fe2c159ade45a0195476780e315012bd60b Mon Sep 17 00:00:00 2001
From: Denis Yaroshevskiy <dyaroshev@meta.com>
Date: Mon, 7 Oct 2024 07:20:00 -0700
Subject: [PATCH] change contains interface to be more restrictive (#2309)

Summary:

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.

Differential Revision: D63979345
---
 folly/algorithm/simd/Contains.h            | 46 +++++++++++++++++----
 folly/algorithm/simd/detail/Traits.h       |  4 +-
 folly/algorithm/simd/test/ContainsTest.cpp | 48 ++++++++++++++++++++++
 3 files changed, 87 insertions(+), 11 deletions(-)

diff --git a/folly/algorithm/simd/Contains.h b/folly/algorithm/simd/Contains.h
index 00756bb75ec..3530e90c34e 100644
--- a/folly/algorithm/simd/Contains.h
+++ b/folly/algorithm/simd/Contains.h
@@ -38,6 +38,32 @@ template <typename R>
 using std_range_value_t = typename std::iterator_traits<decltype(std::begin(
     std::declval<R&>()))>::value_type;
 
+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 is_unsigned_v<To> || sizeof(From) < sizeof(To);
+}
+
+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 +79,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/folly/algorithm/simd/detail/Traits.h b/folly/algorithm/simd/detail/Traits.h
index f3e34cb7eee..dc9f1c8a63e 100644
--- a/folly/algorithm/simd/detail/Traits.h
+++ b/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/folly/algorithm/simd/test/ContainsTest.cpp b/folly/algorithm/simd/test/ContainsTest.cpp
index 55f1966d458..2b8302c1c20 100644
--- a/folly/algorithm/simd/test/ContainsTest.cpp
+++ b/folly/algorithm/simd/test/ContainsTest.cpp
@@ -20,6 +20,9 @@
 
 #include <folly/portability/GTest.h>
 
+#include <list>
+#include <vector>
+
 namespace folly::simd {
 namespace not_invocable_tests {
 
@@ -33,6 +36,51 @@ static_assert(std::is_invocable_v< //
               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>>);
+
 } // namespace not_invocable_tests
 
 template <typename T>