Auto merge of rust-lang#153025 - joboet:bytestr_precision_display, r=Mark-Simulacrum

core: respect precision in `ByteStr` `Display`

Fixes rust-lang#153022.

`ByteStr`'s `Display` implementation didn't respect the precision parameter. Just like `Formatter::pad`, this is fixed by counting off the characters in the string and truncating after the requested length – with the added complication that the `ByteStr` needs to be divided into chunks first. By including a fast path that avoids counting the characters when no parameters were specified this should also fix the performance regressions caused by rust-lang#152865.

Author: bors

library/core/src/bstr/mod.rs

@@ -6,7 +6,7 @@ mod traits;
 pub use traits::{impl_partial_eq, impl_partial_eq_n, impl_partial_eq_ord};
 
 use crate::borrow::{Borrow, BorrowMut};
-use crate::fmt;
+use crate::fmt::{self, Alignment};
 use crate::ops::{Deref, DerefMut, DerefPure};
 
 /// A wrapper for `&[u8]` representing a human-readable string that's conventionally, but not
@@ -174,43 +174,85 @@ impl fmt::Debug for ByteStr {
 #[unstable(feature = "bstr", issue = "134915")]
 impl fmt::Display for ByteStr {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        let nchars: usize = self
-            .utf8_chunks()
-            .map(|chunk| {
-                chunk.valid().chars().count() + if chunk.invalid().is_empty() { 0 } else { 1 }
-            })
-            .sum();
-
-        let padding = f.width().unwrap_or(0).saturating_sub(nchars);
-        let fill = f.fill();
-
-        let (lpad, rpad) = match f.align() {
-            Some(fmt::Alignment::Right) => (padding, 0),
-            Some(fmt::Alignment::Center) => {
-                let half = padding / 2;
-                (half, half + padding % 2)
+        fn emit(byte_str: &ByteStr, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+            for chunk in byte_str.utf8_chunks() {
+                f.write_str(chunk.valid())?;
+                if !chunk.invalid().is_empty() {
+                    f.write_str("\u{FFFD}")?;
+                }
             }
-            // Either alignment is not specified or it's left aligned
-            // which behaves the same with padding
-            _ => (0, padding),
-        };
 
-        for _ in 0..lpad {
-            write!(f, "{fill}")?;
+            Ok(())
         }
 
-        for chunk in self.utf8_chunks() {
-            f.write_str(chunk.valid())?;
-            if !chunk.invalid().is_empty() {
-                f.write_str("\u{FFFD}")?;
-            }
+        let requested_width = f.width().unwrap_or(0);
+        if requested_width == 0 && f.precision().is_none() {
+            // Avoid counting the characters if no truncation or padding was
+            // requested.
+            return emit(self, f);
         }
 
-        for _ in 0..rpad {
-            write!(f, "{fill}")?;
-        }
+        let (truncated, actual_width) = match f.precision() {
+            // The entire string is truncated away. Weird, but ok.
+            Some(0) => (ByteStr::new(&[]), 0),
+            // Advance through string until we run out of space.
+            Some(precision) => {
+                let mut remaining_width = precision;
+                let mut chunks = self.utf8_chunks();
+                let mut current_width = 0;
+                let mut offset = 0;
+                loop {
+                    let Some(chunk) = chunks.next() else {
+                        // We reached the end of the string without running out
+                        // of space, so print the entire string.
+                        break (self, current_width);
+                    };
+
+                    let mut chars = chunk.valid().char_indices();
+                    let Err(remaining) = chars.advance_by(remaining_width) else {
+                        // We've counted off `precision` characters, so truncate
+                        // the string at the current offset.
+                        break (&self[..offset + chars.offset()], precision);
+                    };
+
+                    offset += chunk.valid().len();
+                    current_width += remaining_width - remaining.get();
+                    remaining_width = remaining.get();
+
+                    // `remaining_width` cannot be zero, there is still space
+                    // remaining. So next, count the � character emitted for
+                    // the invalid chunk (if it exists).
+                    if !chunk.invalid().is_empty() {
+                        offset += chunk.invalid().len();
+                        current_width += 1;
+                        remaining_width -= 1;
+
+                        if remaining_width == 0 {
+                            break (&self[..offset], precision);
+                        }
+                    }
+                }
+            }
+            // The string shouldn't be truncated at all, so just count the number
+            // of characters to calculate the padding.
+            None => {
+                let actual_width = self
+                    .utf8_chunks()
+                    .map(|chunk| {
+                        chunk.valid().chars().count()
+                            + if chunk.invalid().is_empty() { 0 } else { 1 }
+                    })
+                    .sum();
+                (self, actual_width)
+            }
+        };
 
-        Ok(())
+        // The width is originally stored as a 16-bit number, so this cannot fail.
+        let padding = u16::try_from(requested_width.saturating_sub(actual_width)).unwrap();
+
+        let post_padding = f.padding(padding, Alignment::Left)?;
+        emit(truncated, f)?;
+        post_padding.write(f)
     }
 }
 

library/coretests/tests/bstr.rs

@@ -49,4 +49,19 @@ fn test_display() {
     assert_eq!(&format!("{b2:->3}!"), "�(��!");
     assert_eq!(&format!("{b2:-^3}!"), "�(��!");
     assert_eq!(&format!("{b2:-^2}!"), "�(��!");
+
+    assert_eq!(&format!("{b1:.1}!"), &format!("{:.1}!", "abc"));
+    assert_eq!(&format!("{b1:.2}!"), &format!("{:.2}!", "abc"));
+    assert_eq!(&format!("{b1:.3}!"), &format!("{:.3}!", "abc"));
+    assert_eq!(&format!("{b1:-<5.2}!"), &format!("{:-<5.2}!", "abc"));
+    assert_eq!(&format!("{b1:-^5.2}!"), &format!("{:-^5.2}!", "abc"));
+    assert_eq!(&format!("{b1:->5.2}!"), &format!("{:->5.2}!", "abc"));
+
+    assert_eq!(&format!("{b2:.1}!"), "�!");
+    assert_eq!(&format!("{b2:.2}!"), "�(!");
+    assert_eq!(&format!("{b2:.3}!"), "�(�!");
+    assert_eq!(&format!("{b2:.4}!"), "�(��!");
+    assert_eq!(&format!("{b2:-<6.3}!"), "�(�---!");
+    assert_eq!(&format!("{b2:-^6.3}!"), "-�(�--!");
+    assert_eq!(&format!("{b2:->6.3}!"), "---�(�!");
 }