Handle formatter flags in WTF-8 OsStr Display

The Display implementation for `OsStr` and `Path` on Windows (the WTF-8
version) only handles formatter flags when the entire string is valid
UTF-8. As most paths are valid UTF-8, the common case is formatted like
`str`; however, flags are ignored when they contain an unpaired
surrogate. Implement its Display with the same logic as that of `str`.

Fixes rust-lang#136617 for Windows.

Author: thaliaarchi

library/core/src/fmt/mod.rs

@@ -1516,8 +1516,11 @@ unsafe fn getcount(args: &[rt::Argument<'_>], cnt: &rt::Count) -> Option<u16> {
 }
 
 /// Padding after the end of something. Returned by `Formatter::padding`.
+#[doc(hidden)]
 #[must_use = "don't forget to write the post padding"]
-pub(crate) struct PostPadding {
+#[unstable(feature = "fmt_internals", reason = "internal to standard library", issue = "none")]
+#[derive(Debug)]
+pub struct PostPadding {
     fill: char,
     padding: u16,
 }
@@ -1528,7 +1531,9 @@ impl PostPadding {
     }
 
     /// Writes this post padding.
-    pub(crate) fn write(self, f: &mut Formatter<'_>) -> Result {
+    #[doc(hidden)]
+    #[unstable(feature = "fmt_internals", reason = "internal to standard library", issue = "none")]
+    pub fn write(self, f: &mut Formatter<'_>) -> Result {
         for _ in 0..self.padding {
             f.buf.write_char(self.fill)?;
         }
@@ -1738,7 +1743,9 @@ impl<'a> Formatter<'a> {
     ///
     /// Callers are responsible for ensuring post-padding is written after the
     /// thing that is being padded.
-    pub(crate) fn padding(
+    #[doc(hidden)]
+    #[unstable(feature = "fmt_internals", reason = "internal to standard library", issue = "none")]
+    pub fn padding(
         &mut self,
         padding: u16,
         default: Alignment,

library/std/src/ffi/os_str/tests.rs

@@ -105,6 +105,22 @@ fn test_os_string_join() {
     assert_eq!("a b c", strings_abc.join(OsStr::new(" ")));
 }
 
+#[test]
+fn display() {
+    let os_string = OsString::from("bcd");
+    assert_eq!(format!("a{:^10}e", os_string.display()), "a   bcd    e");
+}
+
+#[cfg(windows)]
+#[test]
+fn display_invalid_wtf8_windows() {
+    use crate::os::windows::ffi::OsStringExt;
+
+    let os_string = OsString::from_wide(&[b'b' as _, 0xD800, b'd' as _]);
+    assert_eq!(format!("a{:^10}e", os_string.display()), "a   b�d    e");
+    assert_eq!(format!("a{:^10}e", os_string.as_os_str().display()), "a   b�d    e");
+}
+
 #[test]
 fn test_os_string_default() {
     let os_string: OsString = Default::default();

library/std/src/lib.rs

@@ -298,6 +298,7 @@
 #![feature(formatting_options)]
 #![feature(if_let_guard)]
 #![feature(intra_doc_pointers)]
+#![feature(iter_advance_by)]
 #![feature(lang_items)]
 #![feature(let_chains)]
 #![feature(link_cfg)]

library/std/src/sys_common/wtf8.rs

@@ -588,23 +588,48 @@ impl fmt::Debug for Wtf8 {
 /// Formats the string with unpaired surrogates substituted with the replacement
 /// character, U+FFFD.
 impl fmt::Display for Wtf8 {
-    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
-        let wtf8_bytes = &self.bytes;
-        let mut pos = 0;
-        loop {
-            match self.next_surrogate(pos) {
-                Some((surrogate_pos, _)) => {
-                    formatter.write_str(unsafe {
-                        str::from_utf8_unchecked(&wtf8_bytes[pos..surrogate_pos])
-                    })?;
-                    formatter.write_str(UTF8_REPLACEMENT_CHARACTER)?;
-                    pos = surrogate_pos + 3;
-                }
-                None => {
-                    let s = unsafe { str::from_utf8_unchecked(&wtf8_bytes[pos..]) };
-                    if pos == 0 { return s.fmt(formatter) } else { return formatter.write_str(s) }
-                }
-            }
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        // Corresponds to `Formatter::pad`, but for `Wtf8` instead of `str`.
+
+        // Make sure there's a fast path up front.
+        if f.options().get_width().is_none() && f.options().get_precision().is_none() {
+            return self.write_lossy(f);
+        }
+
+        // The `precision` field can be interpreted as a maximum width for the
+        // string being formatted.
+        let (s, code_point_count) = if let Some(max_code_point_count) = f.options().get_precision()
+        {
+            let mut iter = self.code_point_indices();
+            let remaining = match iter.advance_by(max_code_point_count as usize) {
+                Ok(()) => 0,
+                Err(remaining) => remaining.get(),
+            };
+            // SAFETY: The offset of `.code_point_indices()` is guaranteed to be
+            // in-bounds and between code point boundaries.
+            let truncated = unsafe {
+                Wtf8::from_bytes_unchecked(self.bytes.get_unchecked(..iter.front_offset))
+            };
+            (truncated, max_code_point_count as usize - remaining)
+        } else {
+            // Use the optimized code point counting algorithm for the full
+            // string.
+            (self, self.code_points().count())
+        };
+
+        // The `width` field is more of a minimum width parameter at this point.
+        if let Some(width) = f.options().get_width()
+            && code_point_count < width as usize
+        {
+            // If we're under the minimum width, then fill up the minimum width
+            // with the specified string + some alignment.
+            let post_padding = f.padding(width - code_point_count as u16, fmt::Alignment::Left)?;
+            s.write_lossy(f)?;
+            post_padding.write(f)
+        } else {
+            // If we're over the minimum width or there is no minimum width, we
+            // can just emit the string.
+            s.write_lossy(f)
         }
     }
 }
@@ -726,6 +751,19 @@ impl Wtf8 {
         }
     }
 
+    /// Writes the string as lossy UTF-8 like [`Wtf8::to_string_lossy`].
+    /// It ignores formatter flags.
+    fn write_lossy(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let wtf8_bytes = &self.bytes;
+        let mut pos = 0;
+        while let Some((surrogate_pos, _)) = self.next_surrogate(pos) {
+            f.write_str(unsafe { str::from_utf8_unchecked(&wtf8_bytes[pos..surrogate_pos]) })?;
+            f.write_str(UTF8_REPLACEMENT_CHARACTER)?;
+            pos = surrogate_pos + 3;
+        }
+        f.write_str(unsafe { str::from_utf8_unchecked(&wtf8_bytes[pos..]) })
+    }
+
     /// Converts the WTF-8 string to potentially ill-formed UTF-16
     /// and return an iterator of 16-bit code units.
     ///

library/std/src/sys_common/wtf8/tests.rs

@@ -749,3 +749,18 @@ fn unwobbly_wtf8_plus_utf8_is_utf8() {
     string.push_str("some utf-8");
     assert!(string.is_known_utf8);
 }
+
+#[test]
+fn display_wtf8() {
+    let string = Wtf8Buf::from_wide(&[b'b' as _, 0xD800, b'd' as _]);
+    assert!(!string.is_known_utf8);
+    assert_eq!(format!("a{:^10}e", string), "a   b�d    e");
+    assert_eq!(format!("a{:^10}e", string.as_slice()), "a   b�d    e");
+
+    let mut string = Wtf8Buf::from_str("bcd");
+    assert!(string.is_known_utf8);
+    assert_eq!(format!("a{:^10}e", string), "a   bcd    e");
+    assert_eq!(format!("a{:^10}e", string.as_slice()), "a   bcd    e");
+    string.is_known_utf8 = false;
+    assert_eq!(format!("a{:^10}e", string), "a   bcd    e");
+}

library/std/tests/path.rs

@@ -1819,6 +1819,18 @@ fn test_clone_into() {
 fn display_format_flags() {
     assert_eq!(format!("a{:#<5}b", Path::new("").display()), "a#####b");
     assert_eq!(format!("a{:#<5}b", Path::new("a").display()), "aa####b");
+    assert_eq!(format!("a{:^10}e", Path::new("bcd").display()), "a   bcd    e");
+}
+
+#[cfg(windows)]
+#[test]
+fn display_invalid_wtf8_windows() {
+    use std::ffi::OsString;
+    use std::os::windows::ffi::OsStringExt;
+
+    let path_buf = PathBuf::from(OsString::from_wide(&[b'b' as _, 0xD800, b'd' as _]));
+    assert_eq!(format!("a{:^10}e", path_buf.display()), "a   b�d    e");
+    assert_eq!(format!("a{:^10}e", path_buf.as_path().display()), "a   b�d    e");
 }
 
 #[test]