Fix Unicode width, precision and escaping for wide strings

include/fmt/format.h

@@ -1979,15 +1979,22 @@ template <typename Char> struct find_escape_result {
   uint32_t cp;
 };
 
+template <typename Char> struct escape_scanner {
+  static auto find(const Char* begin, const Char* end)
+      -> find_escape_result<Char> {
+    for (; begin != end; ++begin) {
+      uint32_t cp = static_cast<unsigned_char<Char>>(*begin);
+      if (sizeof(Char) == 1 && cp >= 0x80) continue;
+      if (needs_escape(cp)) return {begin, begin + 1, cp};
+    }
+    return {begin, nullptr, 0};
+  }
+};
+
 template <typename Char>
 auto find_escape(const Char* begin, const Char* end)
     -> find_escape_result<Char> {
-  for (; begin != end; ++begin) {
-    uint32_t cp = static_cast<unsigned_char<Char>>(*begin);
-    if (sizeof(Char) == 1 && cp >= 0x80) continue;
-    if (needs_escape(cp)) return {begin, begin + 1, cp};
-  }
-  return {begin, nullptr, 0};
+  return escape_scanner<Char>::find(begin, end);
 }
 
 inline auto find_escape(const char* begin, const char* end)
@@ -2499,25 +2506,7 @@ FMT_CONSTEXPR auto write(OutputIt out, basic_string_view<Char> s,
 template <typename Char, typename OutputIt,
           FMT_ENABLE_IF(!std::is_same<Char, char>::value)>
 FMT_CONSTEXPR auto write(OutputIt out, basic_string_view<Char> s,
-                         const format_specs& specs) -> OutputIt {
-  auto data = s.data();
-  auto size = s.size();
-  if (specs.precision >= 0 && to_unsigned(specs.precision) < size)
-    size = to_unsigned(specs.precision);
-
-  bool is_debug = specs.type() == presentation_type::debug;
-  if (is_debug) {
-    auto buf = counting_buffer<Char>();
-    write_escaped_string(basic_appender<Char>(buf), s);
-    size = buf.count();
-  }
-
-  return write_padded<Char>(
-      out, specs, size, [=](reserve_iterator<OutputIt> it) {
-        return is_debug ? write_escaped_string(it, s)
-                        : copy<Char>(data, data + size, it);
-      });
-}
+                         const format_specs& specs) -> OutputIt;
 
 template <typename Char, typename OutputIt>
 FMT_CONSTEXPR auto write(OutputIt out, basic_string_view<Char> s,

include/fmt/xchar.h

@@ -23,6 +23,139 @@
 FMT_BEGIN_NAMESPACE
 namespace detail {
 
+template <typename Char, typename F,
+          FMT_ENABLE_IF(!std::is_same<Char, char>::value)>
+FMT_CONSTEXPR void for_each_codepoint(basic_string_view<Char> s, F f) {
+  auto p = s.begin(), end = s.end();
+  while (p != end) {
+    auto begin = p;
+    uint32_t cp = static_cast<unsigned_char<Char>>(*p++);
+    // Leave unpaired surrogates intact so they can be escaped individually.
+    if (sizeof(Char) == 2 && cp >= 0xd800 && cp <= 0xdbff && p != end) {
+      uint32_t low = static_cast<unsigned_char<Char>>(*p);
+      if (low >= 0xdc00 && low <= 0xdfff) {
+        cp = 0x10000 + ((cp - 0xd800) << 10) + (low - 0xdc00);
+        ++p;
+      }
+    }
+    if (!f(cp, basic_string_view<Char>(begin, to_unsigned(p - begin)))) return;
+  }
+}
+
+template <typename Char> struct unicode_escape_scanner {
+  static auto find(const Char* begin, const Char* end)
+      -> find_escape_result<Char> {
+    auto result = find_escape_result<Char>{end, nullptr, 0};
+    for_each_codepoint(basic_string_view<Char>(begin, to_unsigned(end - begin)),
+                       [&](uint32_t cp, basic_string_view<Char> sv) {
+                         if (sizeof(Char) == 1 && cp >= 0x80) return true;
+                         if (needs_escape(cp)) {
+                           result = {sv.begin(), sv.end(), cp};
+                           return false;
+                         }
+                         return true;
+                       });
+    return result;
+  }
+};
+
+template <> struct escape_scanner<wchar_t> : unicode_escape_scanner<wchar_t> {};
+template <>
+struct escape_scanner<char16_t> : unicode_escape_scanner<char16_t> {};
+template <>
+struct escape_scanner<char32_t> : unicode_escape_scanner<char32_t> {};
+
+template <typename Char, typename OutputIt,
+          enable_if_t<!std::is_same<Char, char>::value, int>>
+FMT_CONSTEXPR auto write(OutputIt out, basic_string_view<Char> s,
+                         const format_specs& specs) -> OutputIt {
+  bool is_debug = specs.type() == presentation_type::debug;
+  if (specs.precision < 0 && specs.width == 0) {
+    auto&& it = reserve(out, s.size());
+    return is_debug ? write_escaped_string(it, s) : copy<Char>(s, it);
+  }
+
+  size_t display_width_limit =
+      specs.precision < 0 ? SIZE_MAX : to_unsigned(specs.precision);
+  size_t display_width =
+      !is_debug || specs.precision == 0 ? 0 : 1;  // Account for opening '"'.
+  size_t size = !is_debug || specs.precision == 0 ? 0 : 1;
+  for_each_codepoint(s, [&](uint32_t cp, basic_string_view<Char> sv) {
+    if (is_debug && needs_escape(cp) &&
+        (sizeof(Char) != 1 || std::is_same<Char, char>::value || cp < 0x80)) {
+      counting_buffer<Char> buf;
+      write_escaped_cp(basic_appender<Char>(buf),
+                       find_escape_result<Char>{sv.begin(), sv.end(), cp});
+      // We're reinterpreting code units as display width. That's okay
+      // because write_escaped_cp() only writes ASCII characters.
+      size_t cp_width = buf.count();
+      if (display_width + cp_width <= display_width_limit) {
+        display_width += cp_width;
+        size += cp_width;
+        // If this is the end of the string, account for closing '"'.
+        if (display_width < display_width_limit && sv.end() == s.end()) {
+          ++display_width;
+          ++size;
+        }
+        return true;
+      }
+
+      size += display_width_limit - display_width;
+      display_width = display_width_limit;
+      return false;
+    }
+
+    size_t cp_width = display_width_of(cp);
+    if (cp_width + display_width <= display_width_limit) {
+      display_width += cp_width;
+      size += sv.size();
+      // If this is the end of the string, account for closing '"'.
+      if (is_debug && display_width < display_width_limit &&
+          sv.end() == s.end()) {
+        ++display_width;
+        ++size;
+      }
+      return true;
+    }
+
+    return false;
+  });
+
+  if (is_debug && s.size() == 0 && specs.precision != 0 &&
+      display_width < display_width_limit) {
+    ++display_width;
+    ++size;
+  }
+
+  struct bounded_output_iterator {
+    reserve_iterator<OutputIt> underlying_iterator;
+    size_t bound;
+
+    FMT_CONSTEXPR auto operator*() -> bounded_output_iterator& { return *this; }
+    FMT_CONSTEXPR auto operator++() -> bounded_output_iterator& {
+      return *this;
+    }
+    FMT_CONSTEXPR auto operator++(int) -> bounded_output_iterator& {
+      return *this;
+    }
+    FMT_CONSTEXPR auto operator=(Char c) -> bounded_output_iterator& {
+      if (bound > 0) {
+        *underlying_iterator++ = c;
+        --bound;
+      }
+      return *this;
+    }
+  };
+
+  return write_padded<Char>(
+      out, specs, size, display_width, [=](reserve_iterator<OutputIt> it) {
+        return is_debug
+                   ? write_escaped_string(bounded_output_iterator{it, size}, s)
+                         .underlying_iterator
+                   : copy<Char>(s.data(), s.data() + size, it);
+      });
+}
+
 template <typename T>
 using is_exotic_char = bool_constant<!std::is_same<T, char>::value>;
 

test/xchar-test.cc

@@ -298,6 +298,75 @@ TEST(xchar_test, escape_string) {
   EXPECT_EQ(fmt::format(L"{}", vec{L"понедельник"}), L"[\"понедельник\"]");
 }
 
+TEST(xchar_test, utf16_string_width) {
+  EXPECT_EQ(fmt::format(u"{:>4}", u"\u4e2d"), u"  \u4e2d");
+  EXPECT_EQ(fmt::format(u"{:^5}", u"\U0001f600"), u" \U0001f600  ");
+  EXPECT_EQ(fmt::format(u"{:<5}", u"a\U0001f600"), u"a\U0001f600  ");
+}
+
+TEST(xchar_test, utf16_string_precision) {
+  EXPECT_EQ(fmt::format(u"{:.1}", u"\U00010000x"), u"\U00010000");
+  EXPECT_EQ(fmt::format(u"{:.1}", u"\U0001f600x"), u"");
+  EXPECT_EQ(fmt::format(u"{:.2}", u"\U0001f600x"), u"\U0001f600");
+  EXPECT_EQ(fmt::format(u"{:.3}", u"a\U0001f600x"), u"a\U0001f600");
+  EXPECT_EQ(fmt::format(u"{:>4.2}", u"\U0001f600x"), u"  \U0001f600");
+  EXPECT_EQ(fmt::format(u"{:.0}", u"\U0001f600"), u"");
+  const char16_t input[] = {u'a', 0, 0xd800, 0xdc00};
+  auto view = fmt::basic_string_view<char16_t>(input, 4);
+  EXPECT_EQ(fmt::format(u"{:.3}", view), std::u16string(input, 4));
+}
+
+TEST(xchar_test, utf16_debug_string) {
+  EXPECT_EQ(fmt::format(u"{:?}", u"\U0001f600"), u"\"\U0001f600\"");
+  EXPECT_EQ(fmt::format(u"{:?}", u"a\U00010000\n\U0001f600"),
+            u"\"a\U00010000\\n\U0001f600\"");
+  EXPECT_EQ(fmt::format(u"{:?}", u"\U0001d173"), u"\"\\U0001d173\"");
+  EXPECT_EQ(fmt::format(u"{:?}", u"\U0010ffff"), u"\"\\U0010ffff\"");
+}
+
+TEST(xchar_test, utf16_debug_width_and_precision) {
+  EXPECT_EQ(fmt::format(u"{:>6?}", u"\U0001f600"), u"  \"\U0001f600\"");
+  EXPECT_EQ(fmt::format(u"{:.0?}", u"\U0001f600"), u"");
+  EXPECT_EQ(fmt::format(u"{:.1?}", u"\U0001f600"), u"\"");
+  EXPECT_EQ(fmt::format(u"{:.2?}", u"\U0001f600"), u"\"");
+  EXPECT_EQ(fmt::format(u"{:.3?}", u"\U0001f600"), u"\"\U0001f600");
+  EXPECT_EQ(fmt::format(u"{:.4?}", u"\U0001f600"), u"\"\U0001f600\"");
+  EXPECT_EQ(fmt::format(u"{:>6.3?}", u"\U0001f600"), u"   \"\U0001f600");
+  EXPECT_EQ(fmt::format(u"{:.4?}", u"\U0001d173"), u"\"\\U0");
+  EXPECT_EQ(fmt::format(u"{:.1?}", u""), u"\"");
+  EXPECT_EQ(fmt::format(u"{:.2?}", u""), u"\"\"");
+}
+
+TEST(xchar_test, utf16_debug_invalid_sequences) {
+  const char16_t high[] = {0xd800, 0};
+  const char16_t low[] = {0xdc00, 0};
+  const char16_t high_ascii[] = {0xd800, u'x', 0};
+  const char16_t high_pair[] = {0xd800, 0xd800, 0xdc00, 0};
+  EXPECT_EQ(fmt::format(u"{:?}", high), u"\"\\ud800\"");
+  EXPECT_EQ(fmt::format(u"{:?}", low), u"\"\\udc00\"");
+  EXPECT_EQ(fmt::format(u"{:?}", high_ascii), u"\"\\ud800x\"");
+  EXPECT_EQ(fmt::format(u"{:?}", high_pair), u"\"\\ud800\U00010000\"");
+  EXPECT_EQ(fmt::format(u"{:.1}", high_pair), std::u16string(high_pair, 1));
+}
+
+TEST(xchar_test, utf32_string_width_and_precision) {
+  EXPECT_EQ(fmt::format(U"{:>4}", U"\u4e2d"), U"  \u4e2d");
+  EXPECT_EQ(fmt::format(U"{:^5}", U"\U0001f600"), U" \U0001f600  ");
+  EXPECT_EQ(fmt::format(U"{:.1}", U"\U0001f600x"), U"");
+  EXPECT_EQ(fmt::format(U"{:.2}", U"\U0001f600x"), U"\U0001f600");
+  EXPECT_EQ(fmt::format(U"{:>6?}", U"\U0001f600"), U"  \"\U0001f600\"");
+  EXPECT_EQ(fmt::format(U"{:.4?}", U"\U0001d173"), U"\"\\U0");
+}
+
+TEST(xchar_test, wide_string_width_and_precision) {
+  EXPECT_EQ(fmt::format(L"{:>4}", L"\u4e2d"), L"  \u4e2d");
+  EXPECT_EQ(fmt::format(L"{:.1}", L"\U0001f600x"), L"");
+  EXPECT_EQ(fmt::format(L"{:.2}", L"\U0001f600x"), L"\U0001f600");
+  EXPECT_EQ(fmt::format(L"{:?}", L"\U0001f600"), L"\"\U0001f600\"");
+  EXPECT_EQ(fmt::format(L"{:>6?}", L"\U0001f600"), L"  \"\U0001f600\"");
+  EXPECT_EQ(fmt::format(L"{:.4?}", L"\U0001d173"), L"\"\\U0");
+}
+
 TEST(xchar_test, to_wstring) { EXPECT_EQ(L"42", fmt::to_wstring(42)); }
 
 #ifndef FMT_STATIC_THOUSANDS_SEPARATOR