Support wide strings (and string_views) in Out<> / IOutputStream

2 files changed, 62 insertions, 15 deletions

diff --git a/util/stream/output.cpp b/util/stream/output.cpp
index 719c877764..8b3f78527d 100644
--- a/util/stream/output.cpp
+++ b/util/stream/output.cpp
@@ -70,24 +70,13 @@ void Out<wchar32>(IOutputStream& o, wchar32 ch) {
     o.Write(buffer, length);
 }
 
-static void WriteString(IOutputStream& o, const wchar16* w, size_t n) {
+template <typename TCharType>
+static void WriteString(IOutputStream& o, const TCharType* w, size_t n) {
     const size_t buflen = (n * MAX_UTF8_BYTES); // * 4 because the conversion functions can convert unicode character into maximum 4 bytes of UTF8
     TTempBuf buffer(buflen + 1);
-    char* const data = buffer.Data();
     size_t written = 0;
-    WideToUTF8(w, n, data, written);
-    data[written] = 0;
-    o.Write(data, written);
-}
-
-static void WriteString(IOutputStream& o, const wchar32* w, size_t n) {
-    const size_t buflen = (n * MAX_UTF8_BYTES); // * 4 because the conversion functions can convert unicode character into maximum 4 bytes of UTF8
-    TTempBuf buffer(buflen + 1);
-    char* const data = buffer.Data();
-    size_t written = 0;
-    WideToUTF8(w, n, data, written);
-    data[written] = 0;
-    o.Write(data, written);
+    WideToUTF8(w, n, buffer.Data(), written);
+    o.Write(buffer.Data(), written);
 }
 
 template <>
@@ -101,11 +90,31 @@ void Out<std::string>(IOutputStream& o, const std::string& p) {
 }
 
 template <>
+void Out<std::wstring>(IOutputStream& o, const std::wstring& p) {
+    WriteString(o, p.data(), p.length());
+}
+
+template <>
+void Out<std::u16string>(IOutputStream& o, const std::u16string& p) {
+    WriteString(o, p.data(), p.length());
+}
+
+template <>
+void Out<std::u32string>(IOutputStream& o, const std::u32string& p) {
+    WriteString(o, p.data(), p.length());
+}
+
+template <>
 void Out<std::string_view>(IOutputStream& o, const std::string_view& p) {
     o.Write(p.data(), p.length());
 }
 
 template <>
+void Out<std::wstring_view>(IOutputStream& o, const std::wstring_view& p) {
+    WriteString(o, p.data(), p.length());
+}
+
+template <>
 void Out<std::u16string_view>(IOutputStream& o, const std::u16string_view& p) {
     WriteString(o, p.data(), p.length());
 }
diff --git a/util/stream/str_ut.cpp b/util/stream/str_ut.cpp
index fc6b46c31a..534b58d71c 100644
--- a/util/stream/str_ut.cpp
+++ b/util/stream/str_ut.cpp
@@ -149,4 +149,42 @@ Y_UNIT_TEST_SUITE(TStringInputOutputTest) {
         // Check old stream is in a valid state
         output1 << "baz";
     }
+
+    // There is no distinct tests for Out<> via IOutputStream.
+    // Let's tests strings output here.
+    Y_UNIT_TEST(TestWritingWideStrings) {
+        using namespace std::literals::string_literals;
+        TString str;
+        TStringOutput stream(str);
+
+        // test char16_t
+        const char16_t* utf16Data = u"Быть или не быть? Вот в чём вопрос";
+        stream << std::u16string(utf16Data);
+        UNIT_ASSERT_STRINGS_EQUAL(str, "Быть или не быть? Вот в чём вопрос");
+        str.clear();
+
+        stream << std::u16string_view(utf16Data);
+        UNIT_ASSERT_STRINGS_EQUAL(str, "Быть или не быть? Вот в чём вопрос");
+        str.clear();
+
+        // test char32_t
+        const char32_t* utf32Data = U"Быть или не быть? Вот в чём вопрос";
+        stream << std::u32string(utf32Data);
+        UNIT_ASSERT_STRINGS_EQUAL(str, "Быть или не быть? Вот в чём вопрос");
+        str.clear();
+
+        stream << std::u32string_view(utf32Data);
+        UNIT_ASSERT_STRINGS_EQUAL(str, "Быть или не быть? Вот в чём вопрос");
+        str.clear();
+
+        // test wchar_t
+        const wchar_t* wcharData = L"Быть или не быть? Вот в чём вопрос";
+        stream << std::wstring(wcharData);
+        UNIT_ASSERT_STRINGS_EQUAL(str, "Быть или не быть? Вот в чём вопрос");
+        str.clear();
+
+        stream << std::wstring_view(wcharData);
+        UNIT_ASSERT_STRINGS_EQUAL(str, "Быть или не быть? Вот в чём вопрос");
+        str.clear();
+    }
 }