simplified code and added extra test

Author: Sriya Pratipati

libc/src/__support/wchar/character_converter.cpp

@@ -27,7 +27,6 @@ constexpr uint32_t MASK_ENCODED_BITS =
     mask_trailing_ones<uint32_t, ENCODED_BITS_PER_UTF8>();
 // Maximum value for utf-32 for a utf-8 sequence of a given length
 constexpr char32_t MAX_VALUE_PER_UTF8_LEN[] = {0x7f, 0x7ff, 0xffff, 0x10ffff};
-constexpr int MAX_UTF8_LENGTH = 4;
 
 CharacterConverter::CharacterConverter(mbstate *mbstate) { state = mbstate; }
 

libc/src/__support/wchar/character_converter.h

@@ -17,6 +17,8 @@
 
 namespace LIBC_NAMESPACE_DECL {
 namespace internal {
+  
+constexpr int MAX_UTF8_LENGTH = 4;
 
 class CharacterConverter {
 private:

libc/src/__support/wchar/mbrtowc.cpp

@@ -36,7 +36,7 @@ ErrorOr<size_t> mbrtowc(wchar_t *__restrict pwc, const char *__restrict s,
       return Error(EILSEQ);
   }
   auto wc = char_conv.pop_utf32();
-  if (wc.has_value()) {
+  if (wc.has_value() && pwc != nullptr) {
     *pwc = wc.value();
     // null terminator -> return 0
     if (wc.value() == L'\0')

libc/src/__support/wchar/mbsrtowcs.cpp

@@ -15,27 +15,32 @@
 #include "src/__support/macros/config.h"
 #include "src/__support/wchar/mbrtowc.h"
 #include "src/__support/wchar/mbstate.h"
+#include "src/__support/wchar/character_converter.h"
 
 namespace LIBC_NAMESPACE_DECL {
 namespace internal {
 
 ErrorOr<size_t> mbsrtowcs(wchar_t *__restrict dst, const char **__restrict src,
                           size_t len, mbstate *__restrict ps) {
   size_t i = 0;
-  constexpr size_t MAX_UTF8_LENGTH = 4;
   // Converting characters until we reach error or null terminator
-  for (; i < len; ++i, ++dst) {
-    auto check = mbrtowc(dst, *src, MAX_UTF8_LENGTH, ps);
+  for (; i < len; ++i) {
+    wchar_t temp;
+    auto check = internal::mbrtowc(dst == nullptr ? &temp : dst, *src,
+                                   MAX_UTF8_LENGTH, ps);
     // Encoding error/invalid mbstate
     if (!check.has_value())
       return check;
     // Successfully encoded, check for null terminator
-    if (*dst == L'\0') {
+    if (temp == L'\0' || (dst != nullptr && *dst == L'\0')) {
       *src = nullptr;
       return i;
     }
     // Set src to point right after the last character converted
     *src = *src + check.value();
+    // Incrementing destination
+    if (dst != nullptr)
+      ++dst;
   }
   return i;
 }

libc/src/wchar/mbsrtowcs.cpp

@@ -25,9 +25,9 @@ LLVM_LIBC_FUNCTION(size_t, mbsrtowcs,
                     size_t len, mbstate_t *__restrict ps)) {
   LIBC_CRASH_ON_NULLPTR(src);
   static internal::mbstate internal_mbstate;
-  wchar_t temp[len];
+  len = dst == nullptr ? SIZE_MAX : len;
   auto ret = internal::mbsrtowcs(
-      dst == nullptr ? temp : dst, src, dst == nullptr ? SIZE_MAX : len,
+      dst, src, len,
       ps == nullptr ? &internal_mbstate
                     : reinterpret_cast<internal::mbstate *>(ps));
   if (!ret.has_value()) {

libc/test/src/wchar/mbsrtowcs_test.cpp

@@ -59,6 +59,23 @@ TEST_F(LlvmLibcMBSRToWCSTest, MultiByteTwoCharacters) {
   ASSERT_EQ(src, nullptr);
 }
 
+TEST_F(LlvmLibcMBSRToWCSTest, MixedNumberOfBytes) {
+  // 'A', sigma symbol 'Σ', recycling symbol '♻', laughing cat emoji '😹'
+  const char *src = "A\xce\xa3\xe2\x99\xbb\xf0\x9f\x98\xb9";
+  wchar_t dest[5];
+  size_t n = LIBC_NAMESPACE::mbsrtowcs(dest, &src, 5, nullptr);
+  ASSERT_ERRNO_SUCCESS();
+  ASSERT_EQ(static_cast<char>(dest[0]), 'A');
+  ASSERT_EQ(static_cast<int>(dest[1]), 931);
+  ASSERT_EQ(static_cast<int>(dest[2]), 9851);
+  ASSERT_EQ(static_cast<int>(dest[3]), 128569);
+  ASSERT_TRUE(dest[4] == L'\0');
+  // Should not count null terminator in number
+  ASSERT_EQ(static_cast<int>(n), 4);
+  // Should set ch to nullptr after reading null terminator
+  ASSERT_EQ(src, nullptr);
+}
+
 TEST_F(LlvmLibcMBSRToWCSTest, ReadLessThanStringLength) {
   // Four laughing cat emojis "😹😹😹😹"
   const char *src =