[clang][bytecode] Implement __builtin_wmemchr (#132254)

2 files changed, 44 insertions, 5 deletions

diff --git a/clang/lib/AST/ByteCode/InterpBuiltin.cpp b/clang/lib/AST/ByteCode/InterpBuiltin.cpp
index 3fa8fbc..57037b6 100644
--- a/clang/lib/AST/ByteCode/InterpBuiltin.cpp
+++ b/clang/lib/AST/ByteCode/InterpBuiltin.cpp
@@ -2039,11 +2039,23 @@ static bool interp__builtin_memchr(InterpState &S, CodePtr OpPC,
     }
   }
 
-  uint64_t DesiredVal =
-      Desired.trunc(S.getASTContext().getCharWidth()).getZExtValue();
+  uint64_t DesiredVal;
+  if (ID == Builtin::BIwmemchr || ID == Builtin::BI__builtin_wmemchr ||
+      ID == Builtin::BIwcschr || ID == Builtin::BI__builtin_wcschr) {
+    // wcschr and wmemchr are given a wchar_t to look for. Just use it.
+    DesiredVal = Desired.getZExtValue();
+  } else {
+    DesiredVal = Desired.trunc(S.getASTContext().getCharWidth()).getZExtValue();
+  }
+
   bool StopAtZero =
       (ID == Builtin::BIstrchr || ID == Builtin::BI__builtin_strchr);
 
+  PrimType ElemT =
+      IsRawByte
+          ? PT_Sint8
+          : *S.getContext().classify(Ptr.getFieldDesc()->getElemQualType());
+
   size_t Index = Ptr.getIndex();
   size_t Step = 0;
   for (;;) {
@@ -2053,7 +2065,10 @@ static bool interp__builtin_memchr(InterpState &S, CodePtr OpPC,
     if (!CheckLoad(S, OpPC, ElemPtr))
       return false;
 
-    unsigned char V = static_cast<unsigned char>(ElemPtr.deref<char>());
+    uint64_t V;
+    INT_TYPE_SWITCH_NO_BOOL(
+        ElemT, { V = static_cast<uint64_t>(ElemPtr.deref<T>().toUnsigned()); });
+
     if (V == DesiredVal) {
       S.Stk.push<Pointer>(ElemPtr);
       return true;
@@ -2556,11 +2571,11 @@ bool InterpretBuiltin(InterpState &S, CodePtr OpPC, const Function *F,
   case Builtin::BI__builtin_memchr:
   case Builtin::BIstrchr:
   case Builtin::BI__builtin_strchr:
+  case Builtin::BIwmemchr:
+  case Builtin::BI__builtin_wmemchr:
 #if 0
   case Builtin::BIwcschr:
   case Builtin::BI__builtin_wcschr:
-  case Builtin::BImemchr:
-  case Builtin::BI__builtin_wmemchr:
 #endif
   case Builtin::BI__builtin_char_memchr:
     if (!interp__builtin_memchr(S, OpPC, Frame, F, Call))
diff --git a/clang/test/AST/ByteCode/builtin-functions.cpp b/clang/test/AST/ByteCode/builtin-functions.cpp
index 11ff48b..3dd3480 100644
--- a/clang/test/AST/ByteCode/builtin-functions.cpp
+++ b/clang/test/AST/ByteCode/builtin-functions.cpp
@@ -20,6 +20,7 @@ extern "C" {
   extern size_t wcslen(const wchar_t *p);
   extern void *memchr(const void *s, int c, size_t n);
   extern char *strchr(const char *s, int c);
+  extern wchar_t *wmemchr(const wchar_t *s, wchar_t c, size_t n);
 }
 
 namespace strcmp {
@@ -1489,3 +1490,26 @@ namespace Strchr {
   constexpr bool a = !strchr("hello", 'h'); // both-error {{constant expression}} \
                                             // both-note {{non-constexpr function 'strchr' cannot be used in a constant expression}}
 }
+
+namespace WMemChr {
+  constexpr const wchar_t *kStr = L"abca\xffff\0dL";
+  constexpr wchar_t kFoo[] = {L'f', L'o', L'o'};
+
+  static_assert(__builtin_wmemchr(kStr, L'a', 0) == nullptr);
+  static_assert(__builtin_wmemchr(kStr, L'a', 1) == kStr);
+  static_assert(__builtin_wmemchr(kStr, L'\0', 5) == nullptr);
+  static_assert(__builtin_wmemchr(kStr, L'\0', 6) == kStr + 5);
+  static_assert(__builtin_wmemchr(kStr, L'\xffff', 8) == kStr + 4);
+  static_assert(__builtin_wmemchr(kFoo, L'x', 3) == nullptr);
+  static_assert(__builtin_wmemchr(kFoo, L'x', 4) == nullptr); // both-error {{not an integral constant}} \
+                                                              // both-note {{dereferenced one-past-the-end}}
+  static_assert(__builtin_wmemchr(nullptr, L'x', 3) == nullptr); // both-error {{not an integral constant}} \
+                                                                 // both-note {{dereferenced null}}
+  static_assert(__builtin_wmemchr(nullptr, L'x', 0) == nullptr);
+
+  constexpr bool b = !wmemchr(L"hello", L'h', 3); // both-error {{constant expression}} \
+                                                  // both-note {{non-constexpr function 'wmemchr' cannot be used in a constant expression}}
+
+  constexpr wchar_t kStr2[] = {L'f', L'o', L'\xffff', L'o'};
+  static_assert(__builtin_wmemchr(kStr2, L'\xffff', 4) == kStr2 + 2);
+}