Fix f-string padding

Co-Authored-By: yt2b <[email protected]>

Author: youknowone

common/src/format.rs

@@ -1,4 +1,4 @@
-use crate::float_ops;
+use crate::{float_ops, str::BorrowedStr};
 use itertools::{Itertools, PeekingNext};
 use num_bigint::{BigInt, Sign};
 use num_traits::{cast::ToPrimitive, Signed};
@@ -517,8 +517,12 @@ impl FormatSpec {
             }
         };
         let magnitude_str = self.add_magnitude_separators(raw_magnitude_str?, sign_str);
-        self.format_sign_and_align(&magnitude_str, sign_str, FormatAlign::Right)
-            .map_err(|msg| msg.to_owned())
+        self.format_sign_and_align(
+            unsafe { &BorrowedStr::from_ascii_unchecked(magnitude_str.as_bytes()) },
+            sign_str,
+            FormatAlign::Right,
+        )
+        .map_err(|msg| msg.to_owned())
     }
 
     #[inline]
@@ -596,11 +600,15 @@ impl FormatSpec {
         };
         let sign_prefix = format!("{}{}", sign_str, prefix);
         let magnitude_str = self.add_magnitude_separators(raw_magnitude_str?, &sign_prefix);
-        self.format_sign_and_align(&magnitude_str, &sign_prefix, FormatAlign::Right)
-            .map_err(|msg| msg.to_owned())
+        self.format_sign_and_align(
+            unsafe { &BorrowedStr::from_ascii_unchecked(magnitude_str.as_bytes()) },
+            &sign_prefix,
+            FormatAlign::Right,
+        )
+        .map_err(|msg| msg.to_owned())
     }
 
-    pub fn format_string(&self, s: &str) -> Result<String, &'static str> {
+    pub fn format_string(&self, s: &BorrowedStr) -> Result<String, &'static str> {
         match self.format_type {
             Some(FormatType::String) | None => self
                 .format_sign_and_align(s, "", FormatAlign::Left)
@@ -616,14 +624,13 @@ impl FormatSpec {
 
     fn format_sign_and_align(
         &self,
-        magnitude_string: &str,
+        magnitude_str: &BorrowedStr,
         sign_str: &str,
         default_align: FormatAlign,
     ) -> Result<String, &'static str> {
         let align = self.align.unwrap_or(default_align);
 
-        // Use the byte length as the string length since we're in ascii
-        let num_chars = magnitude_string.len();
+        let num_chars = magnitude_str.char_len();
         let fill_char = self.fill.unwrap_or(' ');
         let fill_chars_needed: i32 = self.width.map_or(0, |w| {
             cmp::max(0, (w as i32) - (num_chars as i32) - (sign_str.len() as i32))
@@ -632,20 +639,20 @@ impl FormatSpec {
             FormatAlign::Left => format!(
                 "{}{}{}",
                 sign_str,
-                magnitude_string,
+                magnitude_str,
                 FormatSpec::compute_fill_string(fill_char, fill_chars_needed)
             ),
             FormatAlign::Right => format!(
                 "{}{}{}",
                 FormatSpec::compute_fill_string(fill_char, fill_chars_needed),
                 sign_str,
-                magnitude_string
+                magnitude_str
             ),
             FormatAlign::AfterSign => format!(
                 "{}{}{}",
                 sign_str,
                 FormatSpec::compute_fill_string(fill_char, fill_chars_needed),
-                magnitude_string
+                magnitude_str
             ),
             FormatAlign::Center => {
                 let left_fill_chars_needed = fill_chars_needed / 2;
@@ -654,7 +661,7 @@ impl FormatSpec {
                     FormatSpec::compute_fill_string(fill_char, left_fill_chars_needed);
                 let right_fill_string =
                     FormatSpec::compute_fill_string(fill_char, right_fill_chars_needed);
-                format!("{left_fill_string}{sign_str}{magnitude_string}{right_fill_string}")
+                format!("{left_fill_string}{sign_str}{magnitude_str}{right_fill_string}")
             }
         })
     }

common/src/str.rs

@@ -1,3 +1,7 @@
+use crate::{
+    atomic::{PyAtomic, Radium},
+    hash::PyHash,
+};
 use ascii::AsciiString;
 use once_cell::unsync::OnceCell;
 use std::{
@@ -12,6 +16,119 @@ pub type wchar_t = libc::wchar_t;
 #[allow(non_camel_case_types)]
 pub type wchar_t = u32;
 
+/// Utf8 + state.ascii (+ PyUnicode_Kind in future)
+#[derive(Debug, Copy, Clone, PartialEq)]
+pub enum PyStrKind {
+    Ascii,
+    Utf8,
+}
+
+impl std::ops::BitOr for PyStrKind {
+    type Output = Self;
+    fn bitor(self, other: Self) -> Self {
+        match (self, other) {
+            (Self::Ascii, Self::Ascii) => Self::Ascii,
+            _ => Self::Utf8,
+        }
+    }
+}
+
+impl PyStrKind {
+    #[inline]
+    pub fn new_data(self) -> PyStrKindData {
+        match self {
+            PyStrKind::Ascii => PyStrKindData::Ascii,
+            PyStrKind::Utf8 => PyStrKindData::Utf8(Radium::new(usize::MAX)),
+        }
+    }
+}
+
+#[derive(Debug)]
+pub enum PyStrKindData {
+    Ascii,
+    // uses usize::MAX as a sentinel for "uncomputed"
+    Utf8(PyAtomic<usize>),
+}
+
+impl PyStrKindData {
+    #[inline]
+    pub fn kind(&self) -> PyStrKind {
+        match self {
+            PyStrKindData::Ascii => PyStrKind::Ascii,
+            PyStrKindData::Utf8(_) => PyStrKind::Utf8,
+        }
+    }
+}
+
+pub struct BorrowedStr<'a> {
+    bytes: &'a [u8],
+    kind: PyStrKindData,
+    #[allow(dead_code)]
+    hash: PyAtomic<PyHash>,
+}
+
+impl<'a> BorrowedStr<'a> {
+    /// # Safety
+    /// `s` have to be an ascii string
+    #[inline]
+    pub unsafe fn from_ascii_unchecked(s: &'a [u8]) -> Self {
+        debug_assert!(s.is_ascii());
+        Self {
+            bytes: s,
+            kind: PyStrKind::Ascii.new_data(),
+            hash: PyAtomic::<PyHash>::new(0),
+        }
+    }
+
+    #[inline]
+    pub fn as_str(&self) -> &str {
+        unsafe {
+            // SAFETY: Both PyStrKind::{Ascii, Utf8} are valid utf8 string
+            std::str::from_utf8_unchecked(self.bytes)
+        }
+    }
+
+    #[inline]
+    pub fn char_len(&self) -> usize {
+        match self.kind {
+            PyStrKindData::Ascii => self.bytes.len(),
+            PyStrKindData::Utf8(ref len) => match len.load(core::sync::atomic::Ordering::Relaxed) {
+                usize::MAX => self._compute_char_len(),
+                len => len,
+            },
+        }
+    }
+
+    #[cold]
+    fn _compute_char_len(&self) -> usize {
+        match self.kind {
+            PyStrKindData::Utf8(ref char_len) => {
+                let len = self.as_str().chars().count();
+                // len cannot be usize::MAX, since vec.capacity() < sys.maxsize
+                char_len.store(len, core::sync::atomic::Ordering::Relaxed);
+                len
+            }
+            _ => unsafe {
+                debug_assert!(false); // invalid for non-utf8 strings
+                std::hint::unreachable_unchecked()
+            },
+        }
+    }
+}
+
+impl std::ops::Deref for BorrowedStr<'_> {
+    type Target = str;
+    fn deref(&self) -> &str {
+        self.as_str()
+    }
+}
+
+impl std::fmt::Display for BorrowedStr<'_> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        self.as_str().fmt(f)
+    }
+}
+
 pub fn try_get_chars(s: &str, range: impl RangeBounds<usize>) -> Option<&str> {
     let mut chars = s.chars();
     let start = match range.start_bound() {

vm/src/builtins/str.rs

@@ -7,7 +7,10 @@ use crate::{
     anystr::{self, adjust_indices, AnyStr, AnyStrContainer, AnyStrWrapper},
     atomic_func,
     class::PyClassImpl,
-    common::format::{FormatSpec, FormatString, FromTemplate},
+    common::{
+        format::{FormatSpec, FormatString, FromTemplate},
+        str::{BorrowedStr, PyStrKind, PyStrKindData},
+    },
     convert::{ToPyException, ToPyObject},
     format::{format, format_map},
     function::{ArgIterable, FuncArgs, OptionalArg, OptionalOption, PyComparisonValue},
@@ -39,48 +42,6 @@ use unic_ucd_category::GeneralCategory;
 use unic_ucd_ident::{is_xid_continue, is_xid_start};
 use unicode_casing::CharExt;
 
-/// Utf8 + state.ascii (+ PyUnicode_Kind in future)
-#[derive(Debug, Copy, Clone, PartialEq)]
-pub(crate) enum PyStrKind {
-    Ascii,
-    Utf8,
-}
-
-impl std::ops::BitOr for PyStrKind {
-    type Output = Self;
-    fn bitor(self, other: Self) -> Self {
-        match (self, other) {
-            (Self::Ascii, Self::Ascii) => Self::Ascii,
-            _ => Self::Utf8,
-        }
-    }
-}
-
-impl PyStrKind {
-    fn new_data(self) -> PyStrKindData {
-        match self {
-            PyStrKind::Ascii => PyStrKindData::Ascii,
-            PyStrKind::Utf8 => PyStrKindData::Utf8(Radium::new(usize::MAX)),
-        }
-    }
-}
-
-#[derive(Debug)]
-enum PyStrKindData {
-    Ascii,
-    // uses usize::MAX as a sentinel for "uncomputed"
-    Utf8(PyAtomic<usize>),
-}
-
-impl PyStrKindData {
-    fn kind(&self) -> PyStrKind {
-        match self {
-            PyStrKindData::Ascii => PyStrKind::Ascii,
-            PyStrKindData::Utf8(_) => PyStrKind::Utf8,
-        }
-    }
-}
-
 impl TryFromBorrowedObject for String {
     fn try_from_borrowed_object(vm: &VirtualMachine, obj: &PyObject) -> PyResult<Self> {
         obj.try_value_with(|pystr: &PyStr| Ok(pystr.as_str().to_owned()), vm)
@@ -385,6 +346,10 @@ impl PyStr {
             PyStrKind::Utf8 => self.as_str().chars().all(test),
         }
     }
+
+    fn borrow(&self) -> &BorrowedStr {
+        unsafe { std::mem::transmute(self) }
+    }
 }
 
 #[pyclass(
@@ -475,28 +440,7 @@ impl PyStr {
     #[pymethod(name = "__len__")]
     #[inline]
     pub fn char_len(&self) -> usize {
-        match self.kind {
-            PyStrKindData::Ascii => self.bytes.len(),
-            PyStrKindData::Utf8(ref len) => match len.load(atomic::Ordering::Relaxed) {
-                usize::MAX => self._compute_char_len(),
-                len => len,
-            },
-        }
-    }
-    #[cold]
-    fn _compute_char_len(&self) -> usize {
-        match self.kind {
-            PyStrKindData::Utf8(ref char_len) => {
-                let len = self.as_str().chars().count();
-                // len cannot be usize::MAX, since vec.capacity() < sys.maxsize
-                char_len.store(len, atomic::Ordering::Relaxed);
-                len
-            }
-            _ => unsafe {
-                debug_assert!(false); // invalid for non-utf8 strings
-                std::hint::unreachable_unchecked()
-            },
-        }
+        self.borrow().char_len()
     }
 
     #[pymethod(name = "isascii")]
@@ -787,7 +731,7 @@ impl PyStr {
     #[pymethod(name = "__format__")]
     fn format_str(&self, spec: PyStrRef, vm: &VirtualMachine) -> PyResult<String> {
         FormatSpec::parse(spec.as_str())
-            .and_then(|format_spec| format_spec.format_string(self.as_str()))
+            .and_then(|format_spec| format_spec.format_string(self.borrow()))
             .map_err(|err| vm.new_value_error(err.to_string()))
     }