From 4dd45e29e92e1caaeff09435e47f1fedaf38236d Mon Sep 17 00:00:00 2001
From: Denis Yaroshevskiy <dyaroshev@meta.com>
Date: Mon, 7 Oct 2024 06:27:22 -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            | 35 +++++++++++++-----
 folly/algorithm/simd/detail/Traits.h       |  4 +-
 folly/algorithm/simd/test/ContainsTest.cpp | 43 ++++++++++++++++++++++
 3 files changed, 71 insertions(+), 11 deletions(-)

diff --git a/folly/algorithm/simd/Contains.h b/folly/algorithm/simd/Contains.h
index 00756bb75ec..e9b9b47fb0d 100644
--- a/folly/algorithm/simd/Contains.h
+++ b/folly/algorithm/simd/Contains.h
@@ -38,6 +38,21 @@ 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 = (sizeof(From) <= sizeof(To)) &&
+    ((std::is_signed_v<From> && std::is_signed_v<To>) ||
+     (std::is_unsigned_v<From> && std::is_unsigned_v<To>) ||
+     (std::is_unsigned_v<From> && (sizeof(From) < sizeof(To))));
+
+template <typename R, typename T>
+constexpr bool contains_haystack_needle_test_v = //
+    std::is_invocable_v<AsSimdFriendlyUintFn, R> &&
+    has_integral_simd_friendly_equivalent_scalar_v<std_range_value_t<R>> &&
+    has_integral_simd_friendly_equivalent_scalar_v<T> &&
+    convertible_with_no_loss<
+        simd_friendly_equivalent_scalar_t<T>,
+        simd_friendly_equivalent_scalar_t<std_range_value_t<R>>>;
+
 } // namespace detail
 
 /**
@@ -53,23 +68,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_v<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..d4b7479dcd5 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,46 @@ 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>);
+
 } // namespace not_invocable_tests
 
 template <typename T>