From cace112b1ada0a51a8d2a191c656332a2d425c76 Mon Sep 17 00:00:00 2001 From: Noa Date: Thu, 6 Mar 2025 19:20:33 -0600 Subject: [PATCH] Allow surrogates in str --- Cargo.lock | 2 + common/Cargo.toml | 2 + common/src/encodings.rs | 22 +- common/src/lib.rs | 1 + common/src/str.rs | 390 +++++++-- common/src/wtf8/core_char.rs | 113 +++ common/src/wtf8/core_str.rs | 113 +++ common/src/wtf8/mod.rs | 1347 +++++++++++++++++++++++++++++++ stdlib/src/pystruct.rs | 31 +- stdlib/src/unicodedata.rs | 49 +- vm/src/anystr.rs | 126 +-- vm/src/builtins/genericalias.rs | 33 +- vm/src/builtins/object.rs | 2 +- vm/src/builtins/set.rs | 4 +- vm/src/builtins/str.rs | 894 +++++++++++++++----- vm/src/builtins/union.rs | 2 +- vm/src/bytesinner.rs | 38 +- vm/src/codecs.rs | 23 +- vm/src/frame.rs | 7 +- vm/src/import.rs | 10 +- vm/src/stdlib/builtins.rs | 28 +- vm/src/stdlib/codecs.rs | 4 +- vm/src/vm/interpreter.rs | 2 +- vm/src/vm/mod.rs | 2 +- 24 files changed, 2758 insertions(+), 487 deletions(-) create mode 100644 common/src/wtf8/core_char.rs create mode 100644 common/src/wtf8/core_str.rs create mode 100644 common/src/wtf8/mod.rs diff --git a/Cargo.lock b/Cargo.lock index b214780cb..94b0732b8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2338,6 +2338,7 @@ version = "0.4.0" dependencies = [ "ascii", "bitflags 2.8.0", + "bstr", "cfg-if", "itertools 0.14.0", "libc", @@ -2345,6 +2346,7 @@ dependencies = [ "malachite-base", "malachite-bigint", "malachite-q", + "memchr", "num-complex", "num-traits", "once_cell", diff --git a/common/Cargo.toml b/common/Cargo.toml index 589170064..e9aeba745 100644 --- a/common/Cargo.toml +++ b/common/Cargo.toml @@ -16,12 +16,14 @@ rustpython-literal = { workspace = true } ascii = { workspace = true } bitflags = { workspace = true } +bstr = { workspace = true } cfg-if = { workspace = true } itertools = { workspace = true } libc = { workspace = true } malachite-bigint = { workspace = true } malachite-q = { workspace = true } malachite-base = { workspace = true } +memchr = { workspace = true } num-complex = { workspace = true } num-traits = { workspace = true } once_cell = { workspace = true } diff --git a/common/src/encodings.rs b/common/src/encodings.rs index 858d3b8c6..3aaca68c3 100644 --- a/common/src/encodings.rs +++ b/common/src/encodings.rs @@ -1,10 +1,12 @@ use std::ops::Range; +use crate::wtf8::{Wtf8, Wtf8Buf}; + pub type EncodeErrorResult = Result<(EncodeReplace, usize), E>; pub type DecodeErrorResult = Result<(S, Option, usize), E>; -pub trait StrBuffer: AsRef { +pub trait StrBuffer: AsRef { fn is_ascii(&self) -> bool { self.as_ref().is_ascii() } @@ -63,19 +65,19 @@ fn decode_utf8_compatible( errors: &E, decode: DecodeF, handle_error: ErrF, -) -> Result<(String, usize), E::Error> +) -> Result<(Wtf8Buf, usize), E::Error> where DecodeF: Fn(&[u8]) -> Result<&str, DecodeError<'_>>, ErrF: Fn(&[u8], Option) -> HandleResult<'_>, { if data.is_empty() { - return Ok((String::new(), 0)); + return Ok((Wtf8Buf::new(), 0)); } // we need to coerce the lifetime to that of the function body rather than the // anonymous input lifetime, so that we can assign it data borrowed from data_from_err let mut data = data; let mut data_from_err: E::BytesBuf; - let mut out = String::with_capacity(data.len()); + let mut out = Wtf8Buf::with_capacity(data.len()); let mut remaining_index = 0; let mut remaining_data = data; loop { @@ -98,7 +100,7 @@ where err_idx..err_len.map_or_else(|| data.len(), |len| err_idx + len); let (replace, new_data, restart) = errors.handle_decode_error(data, err_range, reason)?; - out.push_str(replace.as_ref()); + out.push_wtf8(replace.as_ref()); if let Some(new_data) = new_data { data_from_err = new_data; data = data_from_err.as_ref(); @@ -130,7 +132,7 @@ pub mod utf8 { data: &[u8], errors: &E, final_decode: bool, - ) -> Result<(String, usize), E::Error> { + ) -> Result<(Wtf8Buf, usize), E::Error> { decode_utf8_compatible( data, errors, @@ -218,7 +220,7 @@ pub mod latin_1 { )?; match replace { EncodeReplace::Str(s) => { - if s.as_ref().chars().any(|c| (c as u32) > 255) { + if s.as_ref().code_points().any(|c| c.to_u32() > 255) { return Err( errors.error_encoding(full_data, char_range, ERR_REASON) ); @@ -240,10 +242,10 @@ pub mod latin_1 { Ok(out) } - pub fn decode(data: &[u8], _errors: &E) -> Result<(String, usize), E::Error> { + pub fn decode(data: &[u8], _errors: &E) -> Result<(Wtf8Buf, usize), E::Error> { let out: String = data.iter().map(|c| *c as char).collect(); let out_len = out.len(); - Ok((out, out_len)) + Ok((out.into(), out_len)) } } @@ -303,7 +305,7 @@ pub mod ascii { Ok(out) } - pub fn decode(data: &[u8], errors: &E) -> Result<(String, usize), E::Error> { + pub fn decode(data: &[u8], errors: &E) -> Result<(Wtf8Buf, usize), E::Error> { decode_utf8_compatible( data, errors, diff --git a/common/src/lib.rs b/common/src/lib.rs index 760ba8c55..e83a9af43 100644 --- a/common/src/lib.rs +++ b/common/src/lib.rs @@ -29,6 +29,7 @@ pub mod static_cell; pub mod str; #[cfg(windows)] pub mod windows; +pub mod wtf8; pub mod vendored { pub use ascii; diff --git a/common/src/str.rs b/common/src/str.rs index 89d2381d3..dbe4e2223 100644 --- a/common/src/str.rs +++ b/common/src/str.rs @@ -1,9 +1,9 @@ -use crate::{ - atomic::{PyAtomic, Radium}, - format::CharLen, - hash::PyHash, -}; -use ascii::AsciiString; +use crate::atomic::{PyAtomic, Radium}; +use crate::format::CharLen; +use crate::wtf8::{CodePoint, Wtf8, Wtf8Buf}; +use ascii::{AsciiChar, AsciiStr, AsciiString}; +use core::fmt; +use core::sync::atomic::Ordering::Relaxed; use std::ops::{Bound, RangeBounds}; #[cfg(not(target_arch = "wasm32"))] @@ -15,132 +15,313 @@ pub type wchar_t = u32; /// Utf8 + state.ascii (+ PyUnicode_Kind in future) #[derive(Debug, Copy, Clone, PartialEq)] -pub enum PyStrKind { +pub enum StrKind { Ascii, Utf8, + Wtf8, } -impl std::ops::BitOr for PyStrKind { +impl std::ops::BitOr for StrKind { type Output = Self; fn bitor(self, other: Self) -> Self { + use StrKind::*; match (self, other) { - (Self::Ascii, Self::Ascii) => Self::Ascii, - _ => Self::Utf8, + (Wtf8, _) | (_, Wtf8) => Wtf8, + (Utf8, _) | (_, Utf8) => Utf8, + (Ascii, Ascii) => Ascii, } } } -impl PyStrKind { - #[inline] - pub fn new_data(self) -> PyStrKindData { - match self { - PyStrKind::Ascii => PyStrKindData::Ascii, - PyStrKind::Utf8 => PyStrKindData::Utf8(Radium::new(usize::MAX)), +impl StrKind { + pub fn is_ascii(&self) -> bool { + matches!(self, Self::Ascii) + } + + pub fn is_utf8(&self) -> bool { + matches!(self, Self::Ascii | Self::Utf8) + } +} + +pub trait DeduceStrKind { + fn str_kind(&self) -> StrKind; +} + +impl DeduceStrKind for str { + fn str_kind(&self) -> StrKind { + if self.is_ascii() { + StrKind::Ascii + } else { + StrKind::Utf8 } } } +impl DeduceStrKind for Wtf8 { + fn str_kind(&self) -> StrKind { + if self.is_ascii() { + StrKind::Ascii + } else if self.is_utf8() { + StrKind::Utf8 + } else { + StrKind::Wtf8 + } + } +} + +impl DeduceStrKind for String { + fn str_kind(&self) -> StrKind { + (**self).str_kind() + } +} + +impl DeduceStrKind for Wtf8Buf { + fn str_kind(&self) -> StrKind { + (**self).str_kind() + } +} + +impl DeduceStrKind for &T { + fn str_kind(&self) -> StrKind { + (**self).str_kind() + } +} + +impl DeduceStrKind for Box { + fn str_kind(&self) -> StrKind { + (**self).str_kind() + } +} + #[derive(Debug)] -pub enum PyStrKindData { - Ascii, - // uses usize::MAX as a sentinel for "uncomputed" - Utf8(PyAtomic), +pub enum PyKindStr<'a> { + Ascii(&'a AsciiStr), + Utf8(&'a str), + Wtf8(&'a Wtf8), } -impl PyStrKindData { - #[inline] - pub fn kind(&self) -> PyStrKind { - match self { - PyStrKindData::Ascii => PyStrKind::Ascii, - PyStrKindData::Utf8(_) => PyStrKind::Utf8, - } +#[derive(Debug, Clone)] +pub struct StrData { + data: Box, + kind: StrKind, + len: StrLen, +} + +struct StrLen(PyAtomic); + +impl From for StrLen { + #[inline(always)] + fn from(value: usize) -> Self { + Self(Radium::new(value)) } } -pub struct BorrowedStr<'a> { - bytes: &'a [u8], - kind: PyStrKindData, - #[allow(dead_code)] - hash: PyAtomic, -} - -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::::new(0), - } - } - - #[inline] - pub fn from_bytes(s: &'a [u8]) -> Self { - let k = if s.is_ascii() { - PyStrKind::Ascii.new_data() +impl fmt::Debug for StrLen { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let len = self.0.load(Relaxed); + if len == usize::MAX { + f.write_str("") } else { - PyStrKind::Utf8.new_data() - }; + len.fmt(f) + } + } +} + +impl StrLen { + #[inline(always)] + fn zero() -> Self { + 0usize.into() + } + #[inline(always)] + fn uncomputed() -> Self { + usize::MAX.into() + } +} + +impl Clone for StrLen { + fn clone(&self) -> Self { + Self(self.0.load(Relaxed).into()) + } +} + +impl Default for StrData { + fn default() -> Self { Self { - bytes: s, - kind: k, - hash: PyAtomic::::new(0), + data: >::default(), + kind: StrKind::Ascii, + len: StrLen::zero(), + } + } +} + +impl From> for StrData { + fn from(value: Box) -> Self { + // doing the check is ~10x faster for ascii, and is actually only 2% slower worst case for + // non-ascii; see https://github.com/RustPython/RustPython/pull/2586#issuecomment-844611532 + let kind = value.str_kind(); + unsafe { Self::new_str_unchecked(value, kind) } + } +} + +impl From> for StrData { + #[inline] + fn from(value: Box) -> Self { + // doing the check is ~10x faster for ascii, and is actually only 2% slower worst case for + // non-ascii; see https://github.com/RustPython/RustPython/pull/2586#issuecomment-844611532 + let kind = value.str_kind(); + unsafe { Self::new_str_unchecked(value.into(), kind) } + } +} + +impl From> for StrData { + #[inline] + fn from(value: Box) -> Self { + Self { + len: value.len().into(), + data: value.into(), + kind: StrKind::Ascii, + } + } +} + +impl From for StrData { + fn from(ch: AsciiChar) -> Self { + AsciiString::from(ch).into_boxed_ascii_str().into() + } +} + +impl From for StrData { + fn from(ch: char) -> Self { + if let Ok(ch) = ascii::AsciiChar::from_ascii(ch) { + ch.into() + } else { + Self { + data: ch.to_string().into(), + kind: StrKind::Utf8, + len: 1.into(), + } + } + } +} + +impl From for StrData { + fn from(ch: CodePoint) -> Self { + if let Some(ch) = ch.to_char() { + ch.into() + } else { + Self { + data: Wtf8Buf::from(ch).into(), + kind: StrKind::Wtf8, + len: 1.into(), + } + } + } +} + +impl StrData { + /// # Safety + /// + /// Given `bytes` must be valid data for given `kind` + pub unsafe fn new_str_unchecked(data: Box, kind: StrKind) -> Self { + let len = match kind { + StrKind::Ascii => data.len().into(), + _ => StrLen::uncomputed(), + }; + Self { data, kind, len } + } + + /// # Safety + /// + /// `char_len` must be accurate. + pub unsafe fn new_with_char_len(data: Box, kind: StrKind, char_len: usize) -> Self { + Self { + data, + kind, + len: char_len.into(), } } #[inline] - pub fn as_str(&self) -> &str { - unsafe { - // SAFETY: Both PyStrKind::{Ascii, Utf8} are valid utf8 string - std::str::from_utf8_unchecked(self.bytes) + pub fn as_wtf8(&self) -> &Wtf8 { + &self.data + } + + #[inline] + pub fn as_str(&self) -> Option<&str> { + self.kind + .is_utf8() + .then(|| unsafe { std::str::from_utf8_unchecked(self.data.as_bytes()) }) + } + + pub fn as_ascii(&self) -> Option<&AsciiStr> { + self.kind + .is_ascii() + .then(|| unsafe { AsciiStr::from_ascii_unchecked(self.data.as_bytes()) }) + } + + pub fn kind(&self) -> StrKind { + self.kind + } + + #[inline] + pub fn as_str_kind(&self) -> PyKindStr<'_> { + match self.kind { + StrKind::Ascii => { + PyKindStr::Ascii(unsafe { AsciiStr::from_ascii_unchecked(self.data.as_bytes()) }) + } + StrKind::Utf8 => { + PyKindStr::Utf8(unsafe { std::str::from_utf8_unchecked(self.data.as_bytes()) }) + } + StrKind::Wtf8 => PyKindStr::Wtf8(&self.data), } } + #[inline] + pub fn len(&self) -> usize { + self.data.len() + } + + pub fn is_empty(&self) -> bool { + self.data.is_empty() + } + #[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, - }, + match self.len.0.load(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() - }, + let len = if let Some(s) = self.as_str() { + // utf8 chars().count() is optimized + s.chars().count() + } else { + self.data.code_points().count() + }; + // len cannot be usize::MAX, since vec.capacity() < sys.maxsize + self.len.0.store(len, Relaxed); + len + } + + pub fn nth_char(&self, index: usize) -> CodePoint { + match self.as_str_kind() { + PyKindStr::Ascii(s) => s[index].into(), + PyKindStr::Utf8(s) => s.chars().nth(index).unwrap().into(), + PyKindStr::Wtf8(w) => w.code_points().nth(index).unwrap(), } } } -impl std::ops::Deref for BorrowedStr<'_> { - type Target = str; - fn deref(&self) -> &str { - self.as_str() - } -} - -impl std::fmt::Display for BorrowedStr<'_> { +impl std::fmt::Display for StrData { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - self.as_str().fmt(f) + self.data.fmt(f) } } -impl CharLen for BorrowedStr<'_> { +impl CharLen for StrData { fn char_len(&self) -> usize { self.char_len() } @@ -181,6 +362,41 @@ pub fn char_range_end(s: &str, nchars: usize) -> Option { Some(i) } +pub fn try_get_codepoints(w: &Wtf8, range: impl RangeBounds) -> Option<&Wtf8> { + let mut chars = w.code_points(); + let start = match range.start_bound() { + Bound::Included(&i) => i, + Bound::Excluded(&i) => i + 1, + Bound::Unbounded => 0, + }; + for _ in 0..start { + chars.next()?; + } + let s = chars.as_wtf8(); + let range_len = match range.end_bound() { + Bound::Included(&i) => i + 1 - start, + Bound::Excluded(&i) => i - start, + Bound::Unbounded => return Some(s), + }; + codepoint_range_end(s, range_len).map(|end| &s[..end]) +} + +pub fn get_codepoints(w: &Wtf8, range: impl RangeBounds) -> &Wtf8 { + try_get_codepoints(w, range).unwrap() +} + +#[inline] +pub fn codepoint_range_end(s: &Wtf8, nchars: usize) -> Option { + let i = match nchars.checked_sub(1) { + Some(last_char_index) => { + let (index, c) = s.code_point_indices().nth(last_char_index)?; + index + c.len_wtf8() + } + None => 0, + }; + Some(i) +} + pub fn zfill(bytes: &[u8], width: usize) -> Vec { if width <= bytes.len() { bytes.to_vec() diff --git a/common/src/wtf8/core_char.rs b/common/src/wtf8/core_char.rs new file mode 100644 index 000000000..1444e8e13 --- /dev/null +++ b/common/src/wtf8/core_char.rs @@ -0,0 +1,113 @@ +//! Unstable functions from [`core::char`] + +use core::slice; + +pub const MAX_LEN_UTF8: usize = 4; +pub const MAX_LEN_UTF16: usize = 2; + +// UTF-8 ranges and tags for encoding characters +const TAG_CONT: u8 = 0b1000_0000; +const TAG_TWO_B: u8 = 0b1100_0000; +const TAG_THREE_B: u8 = 0b1110_0000; +const TAG_FOUR_B: u8 = 0b1111_0000; +const MAX_ONE_B: u32 = 0x80; +const MAX_TWO_B: u32 = 0x800; +const MAX_THREE_B: u32 = 0x10000; + +#[inline] +#[must_use] +pub const fn len_utf8(code: u32) -> usize { + match code { + ..MAX_ONE_B => 1, + ..MAX_TWO_B => 2, + ..MAX_THREE_B => 3, + _ => 4, + } +} + +#[inline] +#[must_use] +const fn len_utf16(code: u32) -> usize { + if (code & 0xFFFF) == code { 1 } else { 2 } +} + +/// Encodes a raw `u32` value as UTF-8 into the provided byte buffer, +/// and then returns the subslice of the buffer that contains the encoded character. +/// +/// Unlike `char::encode_utf8`, this method also handles codepoints in the surrogate range. +/// (Creating a `char` in the surrogate range is UB.) +/// The result is valid [generalized UTF-8] but not valid UTF-8. +/// +/// [generalized UTF-8]: https://simonsapin.github.io/wtf-8/#generalized-utf8 +/// +/// # Panics +/// +/// Panics if the buffer is not large enough. +/// A buffer of length four is large enough to encode any `char`. +#[doc(hidden)] +#[inline] +pub fn encode_utf8_raw(code: u32, dst: &mut [u8]) -> &mut [u8] { + let len = len_utf8(code); + match (len, &mut *dst) { + (1, [a, ..]) => { + *a = code as u8; + } + (2, [a, b, ..]) => { + *a = (code >> 6 & 0x1F) as u8 | TAG_TWO_B; + *b = (code & 0x3F) as u8 | TAG_CONT; + } + (3, [a, b, c, ..]) => { + *a = (code >> 12 & 0x0F) as u8 | TAG_THREE_B; + *b = (code >> 6 & 0x3F) as u8 | TAG_CONT; + *c = (code & 0x3F) as u8 | TAG_CONT; + } + (4, [a, b, c, d, ..]) => { + *a = (code >> 18 & 0x07) as u8 | TAG_FOUR_B; + *b = (code >> 12 & 0x3F) as u8 | TAG_CONT; + *c = (code >> 6 & 0x3F) as u8 | TAG_CONT; + *d = (code & 0x3F) as u8 | TAG_CONT; + } + _ => { + panic!( + "encode_utf8: need {len} bytes to encode U+{code:04X} but buffer has just {dst_len}", + dst_len = dst.len(), + ) + } + }; + // SAFETY: `<&mut [u8]>::as_mut_ptr` is guaranteed to return a valid pointer and `len` has been tested to be within bounds. + unsafe { slice::from_raw_parts_mut(dst.as_mut_ptr(), len) } +} + +/// Encodes a raw `u32` value as UTF-16 into the provided `u16` buffer, +/// and then returns the subslice of the buffer that contains the encoded character. +/// +/// Unlike `char::encode_utf16`, this method also handles codepoints in the surrogate range. +/// (Creating a `char` in the surrogate range is UB.) +/// +/// # Panics +/// +/// Panics if the buffer is not large enough. +/// A buffer of length 2 is large enough to encode any `char`. +#[doc(hidden)] +#[inline] +pub fn encode_utf16_raw(mut code: u32, dst: &mut [u16]) -> &mut [u16] { + let len = len_utf16(code); + match (len, &mut *dst) { + (1, [a, ..]) => { + *a = code as u16; + } + (2, [a, b, ..]) => { + code -= 0x1_0000; + *a = (code >> 10) as u16 | 0xD800; + *b = (code & 0x3FF) as u16 | 0xDC00; + } + _ => { + panic!( + "encode_utf16: need {len} bytes to encode U+{code:04X} but buffer has just {dst_len}", + dst_len = dst.len(), + ) + } + }; + // SAFETY: `<&mut [u16]>::as_mut_ptr` is guaranteed to return a valid pointer and `len` has been tested to be within bounds. + unsafe { slice::from_raw_parts_mut(dst.as_mut_ptr(), len) } +} diff --git a/common/src/wtf8/core_str.rs b/common/src/wtf8/core_str.rs new file mode 100644 index 000000000..56f715cf7 --- /dev/null +++ b/common/src/wtf8/core_str.rs @@ -0,0 +1,113 @@ +//! Operations related to UTF-8 validation. +//! +//! Copied from `core::str::validations` + +/// Returns the initial codepoint accumulator for the first byte. +/// The first byte is special, only want bottom 5 bits for width 2, 4 bits +/// for width 3, and 3 bits for width 4. +#[inline] +const fn utf8_first_byte(byte: u8, width: u32) -> u32 { + (byte & (0x7F >> width)) as u32 +} + +/// Returns the value of `ch` updated with continuation byte `byte`. +#[inline] +const fn utf8_acc_cont_byte(ch: u32, byte: u8) -> u32 { + (ch << 6) | (byte & CONT_MASK) as u32 +} + +/// Checks whether the byte is a UTF-8 continuation byte (i.e., starts with the +/// bits `10`). +#[inline] +pub(super) const fn utf8_is_cont_byte(byte: u8) -> bool { + (byte as i8) < -64 +} + +/// Reads the next code point out of a byte iterator (assuming a +/// UTF-8-like encoding). +/// +/// # Safety +/// +/// `bytes` must produce a valid UTF-8-like (UTF-8 or WTF-8) string +#[inline] +pub unsafe fn next_code_point<'a, I: Iterator>(bytes: &mut I) -> Option { + // Decode UTF-8 + let x = *bytes.next()?; + if x < 128 { + return Some(x as u32); + } + + // Multibyte case follows + // Decode from a byte combination out of: [[[x y] z] w] + // NOTE: Performance is sensitive to the exact formulation here + let init = utf8_first_byte(x, 2); + // SAFETY: `bytes` produces an UTF-8-like string, + // so the iterator must produce a value here. + let y = unsafe { *bytes.next().unwrap_unchecked() }; + let mut ch = utf8_acc_cont_byte(init, y); + if x >= 0xE0 { + // [[x y z] w] case + // 5th bit in 0xE0 .. 0xEF is always clear, so `init` is still valid + // SAFETY: `bytes` produces an UTF-8-like string, + // so the iterator must produce a value here. + let z = unsafe { *bytes.next().unwrap_unchecked() }; + let y_z = utf8_acc_cont_byte((y & CONT_MASK) as u32, z); + ch = init << 12 | y_z; + if x >= 0xF0 { + // [x y z w] case + // use only the lower 3 bits of `init` + // SAFETY: `bytes` produces an UTF-8-like string, + // so the iterator must produce a value here. + let w = unsafe { *bytes.next().unwrap_unchecked() }; + ch = (init & 7) << 18 | utf8_acc_cont_byte(y_z, w); + } + } + + Some(ch) +} + +/// Reads the last code point out of a byte iterator (assuming a +/// UTF-8-like encoding). +/// +/// # Safety +/// +/// `bytes` must produce a valid UTF-8-like (UTF-8 or WTF-8) string +#[inline] +pub unsafe fn next_code_point_reverse<'a, I>(bytes: &mut I) -> Option +where + I: DoubleEndedIterator, +{ + // Decode UTF-8 + let w = match *bytes.next_back()? { + next_byte if next_byte < 128 => return Some(next_byte as u32), + back_byte => back_byte, + }; + + // Multibyte case follows + // Decode from a byte combination out of: [x [y [z w]]] + let mut ch; + // SAFETY: `bytes` produces an UTF-8-like string, + // so the iterator must produce a value here. + let z = unsafe { *bytes.next_back().unwrap_unchecked() }; + ch = utf8_first_byte(z, 2); + if utf8_is_cont_byte(z) { + // SAFETY: `bytes` produces an UTF-8-like string, + // so the iterator must produce a value here. + let y = unsafe { *bytes.next_back().unwrap_unchecked() }; + ch = utf8_first_byte(y, 3); + if utf8_is_cont_byte(y) { + // SAFETY: `bytes` produces an UTF-8-like string, + // so the iterator must produce a value here. + let x = unsafe { *bytes.next_back().unwrap_unchecked() }; + ch = utf8_first_byte(x, 4); + ch = utf8_acc_cont_byte(ch, y); + } + ch = utf8_acc_cont_byte(ch, z); + } + ch = utf8_acc_cont_byte(ch, w); + + Some(ch) +} + +/// Mask of the value bits of a continuation byte. +const CONT_MASK: u8 = 0b0011_1111; diff --git a/common/src/wtf8/mod.rs b/common/src/wtf8/mod.rs new file mode 100644 index 000000000..8ec3ae0af --- /dev/null +++ b/common/src/wtf8/mod.rs @@ -0,0 +1,1347 @@ +//! An implementation of [WTF-8], a utf8-compatible encoding that allows for +//! unpaired surrogate codepoints. This implementation additionally allows for +//! paired surrogates that are nonetheless treated as two separate codepoints. +//! +//! +//! RustPython uses this because CPython internally uses a variant of UCS-1/2/4 +//! as its string storage, which treats each `u8`/`u16`/`u32` value (depending +//! on the highest codepoint value in the string) as simply integers, unlike +//! UTF-8 or UTF-16 where some characters are encoded using multi-byte +//! sequences. CPython additionally doesn't disallow the use of surrogates in +//! `str`s (which in UTF-16 pair together to represent codepoints with a value +//! higher than `u16::MAX`) and in fact takes quite extensive advantage of the +//! fact that they're allowed. The `surrogateescape` codec-error handler uses +//! them to represent byte sequences which are invalid in the given codec (e.g. +//! bytes with their high bit set in ASCII or UTF-8) by mapping them into the +//! surrogate range. `surrogateescape` is the default error handler in Python +//! for interacting with the filesystem, and thus if RustPython is to properly +//! support `surrogateescape`, its `str`s must be able to represent surrogates. +//! +//! We use WTF-8 over something more similar to CPython's string implementation +//! because of its compatibility with UTF-8, meaning that in the case where a +//! string has no surrogates, it can be viewed as a UTF-8 Rust [`str`] without +//! needing any copies or re-encoding. +//! +//! This implementation is mostly copied from the WTF-8 implentation in the +//! Rust standard library, which is used as the backing for [`OsStr`] on +//! Windows targets. As previously mentioned, however, it is modified to not +//! join two surrogates into one codepoint when concatenating strings, in order +//! to match CPython's behavior. +//! +//! [WTF-8]: https://simonsapin.github.io/wtf-8 +//! [`OsStr`]: std::ffi::OsStr + +#![allow(clippy::precedence, clippy::match_overlapping_arm)] + +use core::fmt; +use core::hash::{Hash, Hasher}; +use core::iter::FusedIterator; +use core::mem; +use core::ops; +use core::slice; +use core::str; +use core_char::MAX_LEN_UTF8; +use core_char::{MAX_LEN_UTF16, encode_utf8_raw, encode_utf16_raw, len_utf8}; +use core_str::{next_code_point, next_code_point_reverse}; +use itertools::{Either, Itertools}; +use std::borrow::{Borrow, Cow}; +use std::collections::TryReserveError; +use std::string::String; +use std::vec::Vec; + +use bstr::ByteSlice; + +mod core_char; +mod core_str; + +const UTF8_REPLACEMENT_CHARACTER: &str = "\u{FFFD}"; + +/// A Unicode code point: from U+0000 to U+10FFFF. +/// +/// Compares with the `char` type, +/// which represents a Unicode scalar value: +/// a code point that is not a surrogate (U+D800 to U+DFFF). +#[derive(Eq, PartialEq, Ord, PartialOrd, Clone, Copy)] +pub struct CodePoint { + value: u32, +} + +/// Format the code point as `U+` followed by four to six hexadecimal digits. +/// Example: `U+1F4A9` +impl fmt::Debug for CodePoint { + #[inline] + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(formatter, "U+{:04X}", self.value) + } +} + +impl CodePoint { + /// Unsafely creates a new `CodePoint` without checking the value. + /// + /// # Safety + /// + /// `value` must be less than or equal to 0x10FFFF. + #[inline] + pub unsafe fn from_u32_unchecked(value: u32) -> CodePoint { + CodePoint { value } + } + + /// Creates a new `CodePoint` if the value is a valid code point. + /// + /// Returns `None` if `value` is above 0x10FFFF. + #[inline] + pub fn from_u32(value: u32) -> Option { + match value { + 0..=0x10FFFF => Some(CodePoint { value }), + _ => None, + } + } + + /// Creates a new `CodePoint` from a `char`. + /// + /// Since all Unicode scalar values are code points, this always succeeds. + #[inline] + pub fn from_char(value: char) -> CodePoint { + CodePoint { + value: value as u32, + } + } + + /// Returns the numeric value of the code point. + #[inline] + pub fn to_u32(&self) -> u32 { + self.value + } + + /// Returns the numeric value of the code point if it is a leading surrogate. + #[inline] + pub fn to_lead_surrogate(&self) -> Option { + match self.value { + lead @ 0xD800..=0xDBFF => Some(lead as u16), + _ => None, + } + } + + /// Returns the numeric value of the code point if it is a trailing surrogate. + #[inline] + pub fn to_trail_surrogate(&self) -> Option { + match self.value { + trail @ 0xDC00..=0xDFFF => Some(trail as u16), + _ => None, + } + } + + /// Optionally returns a Unicode scalar value for the code point. + /// + /// Returns `None` if the code point is a surrogate (from U+D800 to U+DFFF). + #[inline] + pub fn to_char(&self) -> Option { + match self.value { + 0xD800..=0xDFFF => None, + _ => Some(unsafe { char::from_u32_unchecked(self.value) }), + } + } + + /// Returns a Unicode scalar value for the code point. + /// + /// Returns `'\u{FFFD}'` (the replacement character “�”) + /// if the code point is a surrogate (from U+D800 to U+DFFF). + #[inline] + pub fn to_char_lossy(&self) -> char { + self.to_char().unwrap_or('\u{FFFD}') + } + + pub fn is_char_and(self, f: impl FnOnce(char) -> bool) -> bool { + self.to_char().is_some_and(f) + } + + pub fn encode_wtf8(self, dst: &mut [u8]) -> &mut Wtf8 { + unsafe { Wtf8::from_mut_bytes_unchecked(encode_utf8_raw(self.value, dst)) } + } + + pub fn len_wtf8(&self) -> usize { + len_utf8(self.value) + } +} + +impl From for CodePoint { + fn from(value: u16) -> Self { + unsafe { Self::from_u32_unchecked(value.into()) } + } +} + +impl From for CodePoint { + fn from(value: char) -> Self { + Self::from_char(value) + } +} + +impl From for CodePoint { + fn from(value: ascii::AsciiChar) -> Self { + Self::from_char(value.into()) + } +} + +impl From for Wtf8Buf { + fn from(ch: CodePoint) -> Self { + ch.encode_wtf8(&mut [0; MAX_LEN_UTF8]).to_owned() + } +} + +impl PartialEq for CodePoint { + fn eq(&self, other: &char) -> bool { + self.to_u32() == *other as u32 + } +} +impl PartialEq for char { + fn eq(&self, other: &CodePoint) -> bool { + *self as u32 == other.to_u32() + } +} + +/// An owned, growable string of well-formed WTF-8 data. +/// +/// Similar to `String`, but can additionally contain surrogate code points +/// if they’re not in a surrogate pair. +#[derive(Eq, PartialEq, Ord, PartialOrd, Clone, Default)] +pub struct Wtf8Buf { + bytes: Vec, +} + +impl ops::Deref for Wtf8Buf { + type Target = Wtf8; + + fn deref(&self) -> &Wtf8 { + self.as_slice() + } +} + +impl ops::DerefMut for Wtf8Buf { + fn deref_mut(&mut self) -> &mut Wtf8 { + self.as_mut_slice() + } +} + +impl Borrow for Wtf8Buf { + fn borrow(&self) -> &Wtf8 { + self + } +} + +/// Formats the string in double quotes, with characters escaped according to +/// [`char::escape_debug`] and unpaired surrogates represented as `\u{xxxx}`, +/// where each `x` is a hexadecimal digit. +/// +/// For example, the code units [U+0061, U+D800, U+000A] are formatted as +/// `"a\u{D800}\n"`. +impl fmt::Debug for Wtf8Buf { + #[inline] + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt::Debug::fmt(&**self, formatter) + } +} + +/// Formats the string with unpaired surrogates substituted with the replacement +/// character, U+FFFD. +impl fmt::Display for Wtf8Buf { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt::Display::fmt(&**self, formatter) + } +} + +impl Wtf8Buf { + /// Creates a new, empty WTF-8 string. + #[inline] + pub fn new() -> Wtf8Buf { + Wtf8Buf::default() + } + + /// Creates a new, empty WTF-8 string with pre-allocated capacity for `capacity` bytes. + #[inline] + pub fn with_capacity(capacity: usize) -> Wtf8Buf { + Wtf8Buf { + bytes: Vec::with_capacity(capacity), + } + } + + /// Creates a WTF-8 string from a WTF-8 byte vec. + /// + /// # Safety + /// + /// `value` must contain valid WTF-8. + #[inline] + pub unsafe fn from_bytes_unchecked(value: Vec) -> Wtf8Buf { + Wtf8Buf { bytes: value } + } + + /// Creates a WTF-8 string from a UTF-8 `String`. + /// + /// This takes ownership of the `String` and does not copy. + /// + /// Since WTF-8 is a superset of UTF-8, this always succeeds. + #[inline] + pub fn from_string(string: String) -> Wtf8Buf { + Wtf8Buf { + bytes: string.into_bytes(), + } + } + + pub fn clear(&mut self) { + self.bytes.clear(); + } + + /// Creates a WTF-8 string from a potentially ill-formed UTF-16 slice of 16-bit code units. + /// + /// This is lossless: calling `.encode_wide()` on the resulting string + /// will always return the original code units. + pub fn from_wide(v: &[u16]) -> Wtf8Buf { + let mut string = Wtf8Buf::with_capacity(v.len()); + for item in char::decode_utf16(v.iter().cloned()) { + match item { + Ok(ch) => string.push_char(ch), + Err(surrogate) => { + let surrogate = surrogate.unpaired_surrogate(); + // Surrogates are known to be in the code point range. + let code_point = CodePoint::from(surrogate); + // Skip the WTF-8 concatenation check, + // surrogate pairs are already decoded by decode_utf16 + string.push(code_point); + } + } + } + string + } + + #[inline] + pub fn as_slice(&self) -> &Wtf8 { + unsafe { Wtf8::from_bytes_unchecked(&self.bytes) } + } + + #[inline] + pub fn as_mut_slice(&mut self) -> &mut Wtf8 { + // Safety: `Wtf8` doesn't expose any way to mutate the bytes that would + // cause them to change from well-formed UTF-8 to ill-formed UTF-8, + // which would break the assumptions of the `is_known_utf8` field. + unsafe { Wtf8::from_mut_bytes_unchecked(&mut self.bytes) } + } + + /// Reserves capacity for at least `additional` more bytes to be inserted + /// in the given `Wtf8Buf`. + /// The collection may reserve more space to avoid frequent reallocations. + /// + /// # Panics + /// + /// Panics if the new capacity exceeds `isize::MAX` bytes. + #[inline] + pub fn reserve(&mut self, additional: usize) { + self.bytes.reserve(additional) + } + + /// Tries to reserve capacity for at least `additional` more bytes to be + /// inserted in the given `Wtf8Buf`. The `Wtf8Buf` may reserve more space to + /// avoid frequent reallocations. After calling `try_reserve`, capacity will + /// be greater than or equal to `self.len() + additional`. Does nothing if + /// capacity is already sufficient. This method preserves the contents even + /// if an error occurs. + /// + /// # Errors + /// + /// If the capacity overflows, or the allocator reports a failure, then an error + /// is returned. + #[inline] + pub fn try_reserve(&mut self, additional: usize) -> Result<(), TryReserveError> { + self.bytes.try_reserve(additional) + } + + #[inline] + pub fn reserve_exact(&mut self, additional: usize) { + self.bytes.reserve_exact(additional) + } + + /// Tries to reserve the minimum capacity for exactly `additional` more + /// bytes to be inserted in the given `Wtf8Buf`. After calling + /// `try_reserve_exact`, capacity will be greater than or equal to + /// `self.len() + additional` if it returns `Ok(())`. + /// Does nothing if the capacity is already sufficient. + /// + /// Note that the allocator may give the `Wtf8Buf` more space than it + /// requests. Therefore, capacity can not be relied upon to be precisely + /// minimal. Prefer [`try_reserve`] if future insertions are expected. + /// + /// [`try_reserve`]: Wtf8Buf::try_reserve + /// + /// # Errors + /// + /// If the capacity overflows, or the allocator reports a failure, then an error + /// is returned. + #[inline] + pub fn try_reserve_exact(&mut self, additional: usize) -> Result<(), TryReserveError> { + self.bytes.try_reserve_exact(additional) + } + + #[inline] + pub fn shrink_to_fit(&mut self) { + self.bytes.shrink_to_fit() + } + + #[inline] + pub fn shrink_to(&mut self, min_capacity: usize) { + self.bytes.shrink_to(min_capacity) + } + + #[inline] + pub fn leak<'a>(self) -> &'a mut Wtf8 { + unsafe { Wtf8::from_mut_bytes_unchecked(self.bytes.leak()) } + } + + /// Returns the number of bytes that this string buffer can hold without reallocating. + #[inline] + pub fn capacity(&self) -> usize { + self.bytes.capacity() + } + + /// Append a UTF-8 slice at the end of the string. + #[inline] + pub fn push_str(&mut self, other: &str) { + self.bytes.extend_from_slice(other.as_bytes()) + } + + /// Append a WTF-8 slice at the end of the string. + #[inline] + pub fn push_wtf8(&mut self, other: &Wtf8) { + self.bytes.extend_from_slice(&other.bytes); + } + + /// Append a Unicode scalar value at the end of the string. + #[inline] + pub fn push_char(&mut self, c: char) { + self.push(CodePoint::from_char(c)) + } + + /// Append a code point at the end of the string. + #[inline] + pub fn push(&mut self, code_point: CodePoint) { + self.push_wtf8(code_point.encode_wtf8(&mut [0; MAX_LEN_UTF8])) + } + + /// Shortens a string to the specified length. + /// + /// # Panics + /// + /// Panics if `new_len` > current length, + /// or if `new_len` is not a code point boundary. + #[inline] + pub fn truncate(&mut self, new_len: usize) { + assert!(is_code_point_boundary(self, new_len)); + self.bytes.truncate(new_len) + } + + /// Consumes the WTF-8 string and tries to convert it to a vec of bytes. + #[inline] + pub fn into_bytes(self) -> Vec { + self.bytes + } + + /// Consumes the WTF-8 string and tries to convert it to UTF-8. + /// + /// This does not copy the data. + /// + /// If the contents are not well-formed UTF-8 + /// (that is, if the string contains surrogates), + /// the original WTF-8 string is returned instead. + pub fn into_string(self) -> Result { + if self.is_utf8() { + Ok(unsafe { String::from_utf8_unchecked(self.bytes) }) + } else { + Err(self) + } + } + + /// Consumes the WTF-8 string and converts it lossily to UTF-8. + /// + /// This does not copy the data (but may overwrite parts of it in place). + /// + /// Surrogates are replaced with `"\u{FFFD}"` (the replacement character “�”) + pub fn into_string_lossy(mut self) -> String { + let mut pos = 0; + while let Some((surrogate_pos, _)) = self.next_surrogate(pos) { + pos = surrogate_pos + 3; + // Surrogates and the replacement character are all 3 bytes, so + // they can substituted in-place. + self.bytes[surrogate_pos..pos].copy_from_slice(UTF8_REPLACEMENT_CHARACTER.as_bytes()); + } + unsafe { String::from_utf8_unchecked(self.bytes) } + } + + /// Converts this `Wtf8Buf` into a boxed `Wtf8`. + #[inline] + pub fn into_box(self) -> Box { + // SAFETY: relies on `Wtf8` being `repr(transparent)`. + unsafe { mem::transmute(self.bytes.into_boxed_slice()) } + } + + /// Converts a `Box` into a `Wtf8Buf`. + pub fn from_box(boxed: Box) -> Wtf8Buf { + let bytes: Box<[u8]> = unsafe { mem::transmute(boxed) }; + Wtf8Buf { + bytes: bytes.into_vec(), + } + } +} + +/// Creates a new WTF-8 string from an iterator of code points. +/// +/// This replaces surrogate code point pairs with supplementary code points, +/// like concatenating ill-formed UTF-16 strings effectively would. +impl FromIterator for Wtf8Buf { + fn from_iter>(iter: T) -> Wtf8Buf { + let mut string = Wtf8Buf::new(); + string.extend(iter); + string + } +} + +/// Append code points from an iterator to the string. +/// +/// This replaces surrogate code point pairs with supplementary code points, +/// like concatenating ill-formed UTF-16 strings effectively would. +impl Extend for Wtf8Buf { + fn extend>(&mut self, iter: T) { + let iterator = iter.into_iter(); + let (low, _high) = iterator.size_hint(); + // Lower bound of one byte per code point (ASCII only) + self.bytes.reserve(low); + iterator.for_each(move |code_point| self.push(code_point)); + } +} + +impl> FromIterator for Wtf8Buf { + fn from_iter>(iter: T) -> Self { + let mut buf = Wtf8Buf::new(); + iter.into_iter().for_each(|w| buf.push_wtf8(w.as_ref())); + buf + } +} + +impl AsRef for Wtf8Buf { + fn as_ref(&self) -> &Wtf8 { + self + } +} + +impl From for Wtf8Buf { + fn from(s: String) -> Self { + Wtf8Buf::from_string(s) + } +} + +impl From<&str> for Wtf8Buf { + fn from(s: &str) -> Self { + Wtf8Buf::from_string(s.to_owned()) + } +} + +/// A borrowed slice of well-formed WTF-8 data. +/// +/// Similar to `&str`, but can additionally contain surrogate code points +/// if they’re not in a surrogate pair. +#[derive(PartialEq, Eq, PartialOrd, Ord)] +pub struct Wtf8 { + bytes: [u8], +} + +impl AsRef for Wtf8 { + fn as_ref(&self) -> &Wtf8 { + self + } +} + +impl ToOwned for Wtf8 { + type Owned = Wtf8Buf; + fn to_owned(&self) -> Self::Owned { + self.to_wtf8_buf() + } +} + +/// Formats the string in double quotes, with characters escaped according to +/// [`char::escape_debug`] and unpaired surrogates represented as `\u{xxxx}`, +/// where each `x` is a hexadecimal digit. +impl fmt::Debug for Wtf8 { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + fn write_str_escaped(f: &mut fmt::Formatter<'_>, s: &str) -> fmt::Result { + use std::fmt::Write; + for c in s.chars().flat_map(|c| c.escape_debug()) { + f.write_char(c)? + } + Ok(()) + } + + formatter.write_str("\"")?; + let mut pos = 0; + while let Some((surrogate_pos, surrogate)) = self.next_surrogate(pos) { + write_str_escaped(formatter, unsafe { + str::from_utf8_unchecked(&self.bytes[pos..surrogate_pos]) + })?; + write!(formatter, "\\u{{{:x}}}", surrogate)?; + pos = surrogate_pos + 3; + } + write_str_escaped(formatter, unsafe { + str::from_utf8_unchecked(&self.bytes[pos..]) + })?; + formatter.write_str("\"") + } +} + +/// 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); + } + } + } + } + } +} + +impl Default for &Wtf8 { + fn default() -> Self { + unsafe { Wtf8::from_bytes_unchecked(&[]) } + } +} + +impl Wtf8 { + /// Creates a WTF-8 slice from a UTF-8 `&str` slice. + /// + /// Since WTF-8 is a superset of UTF-8, this always succeeds. + #[inline] + pub fn new + ?Sized>(value: &S) -> &Wtf8 { + value.as_ref() + } + + /// Creates a WTF-8 slice from a WTF-8 byte slice. + /// + /// # Safety + /// + /// `value` must contain valid WTF-8. + #[inline] + pub unsafe fn from_bytes_unchecked(value: &[u8]) -> &Wtf8 { + // SAFETY: start with &[u8], end with fancy &[u8] + unsafe { &*(value as *const [u8] as *const Wtf8) } + } + + /// Creates a mutable WTF-8 slice from a mutable WTF-8 byte slice. + /// + /// Since the byte slice is not checked for valid WTF-8, this functions is + /// marked unsafe. + #[inline] + unsafe fn from_mut_bytes_unchecked(value: &mut [u8]) -> &mut Wtf8 { + // SAFETY: start with &mut [u8], end with fancy &mut [u8] + unsafe { &mut *(value as *mut [u8] as *mut Wtf8) } + } + + /// Returns the length, in WTF-8 bytes. + #[inline] + pub fn len(&self) -> usize { + self.bytes.len() + } + + #[inline] + pub fn is_empty(&self) -> bool { + self.bytes.is_empty() + } + + /// Returns the code point at `position` if it is in the ASCII range, + /// or `b'\xFF'` otherwise. + /// + /// # Panics + /// + /// Panics if `position` is beyond the end of the string. + #[inline] + pub fn ascii_byte_at(&self, position: usize) -> u8 { + match self.bytes[position] { + ascii_byte @ 0x00..=0x7F => ascii_byte, + _ => 0xFF, + } + } + + /// Returns an iterator for the string’s code points. + #[inline] + pub fn code_points(&self) -> Wtf8CodePoints<'_> { + Wtf8CodePoints { + bytes: self.bytes.iter(), + } + } + + /// Returns an iterator for the string’s code points and their indices. + #[inline] + pub fn code_point_indices(&self) -> Wtf8CodePointIndices<'_> { + Wtf8CodePointIndices { + front_offset: 0, + iter: self.code_points(), + } + } + + /// Access raw bytes of WTF-8 data + #[inline] + pub fn as_bytes(&self) -> &[u8] { + &self.bytes + } + + /// Tries to convert the string to UTF-8 and return a `&str` slice. + /// + /// Returns `None` if the string contains surrogates. + /// + /// This does not copy the data. + #[inline] + pub fn as_str(&self) -> Result<&str, str::Utf8Error> { + str::from_utf8(&self.bytes) + } + + /// Creates an owned `Wtf8Buf` from a borrowed `Wtf8`. + pub fn to_wtf8_buf(&self) -> Wtf8Buf { + Wtf8Buf { + bytes: self.bytes.to_vec(), + } + } + + /// Lossily converts the string to UTF-8. + /// Returns a UTF-8 `&str` slice if the contents are well-formed in UTF-8. + /// + /// Surrogates are replaced with `"\u{FFFD}"` (the replacement character “�”). + /// + /// This only copies the data if necessary (if it contains any surrogate). + pub fn to_string_lossy(&self) -> Cow<'_, str> { + let Some((surrogate_pos, _)) = self.next_surrogate(0) else { + return Cow::Borrowed(unsafe { str::from_utf8_unchecked(&self.bytes) }); + }; + let wtf8_bytes = &self.bytes; + let mut utf8_bytes = Vec::with_capacity(self.len()); + utf8_bytes.extend_from_slice(&wtf8_bytes[..surrogate_pos]); + utf8_bytes.extend_from_slice(UTF8_REPLACEMENT_CHARACTER.as_bytes()); + let mut pos = surrogate_pos + 3; + loop { + match self.next_surrogate(pos) { + Some((surrogate_pos, _)) => { + utf8_bytes.extend_from_slice(&wtf8_bytes[pos..surrogate_pos]); + utf8_bytes.extend_from_slice(UTF8_REPLACEMENT_CHARACTER.as_bytes()); + pos = surrogate_pos + 3; + } + None => { + utf8_bytes.extend_from_slice(&wtf8_bytes[pos..]); + return Cow::Owned(unsafe { String::from_utf8_unchecked(utf8_bytes) }); + } + } + } + } + + /// Converts the WTF-8 string to potentially ill-formed UTF-16 + /// and return an iterator of 16-bit code units. + /// + /// This is lossless: + /// calling `Wtf8Buf::from_ill_formed_utf16` on the resulting code units + /// would always return the original WTF-8 string. + #[inline] + pub fn encode_wide(&self) -> EncodeWide<'_> { + EncodeWide { + code_points: self.code_points(), + extra: 0, + } + } + + pub fn chunks(&self) -> Wtf8Chunks<'_> { + Wtf8Chunks { wtf8: self } + } + + pub fn map_utf8<'a, I>(&'a self, f: impl Fn(&'a str) -> I) -> impl Iterator + where + I: Iterator, + { + self.chunks().flat_map(move |chunk| match chunk { + Wtf8Chunk::Utf8(s) => Either::Left(f(s).map_into()), + Wtf8Chunk::Surrogate(c) => Either::Right(std::iter::once(c)), + }) + } + + #[inline] + fn next_surrogate(&self, mut pos: usize) -> Option<(usize, u16)> { + let mut iter = self.bytes[pos..].iter(); + loop { + let b = *iter.next()?; + if b < 0x80 { + pos += 1; + } else if b < 0xE0 { + iter.next(); + pos += 2; + } else if b == 0xED { + match (iter.next(), iter.next()) { + (Some(&b2), Some(&b3)) if b2 >= 0xA0 => { + return Some((pos, decode_surrogate(b2, b3))); + } + _ => pos += 3, + } + } else if b < 0xF0 { + iter.next(); + iter.next(); + pos += 3; + } else { + iter.next(); + iter.next(); + iter.next(); + pos += 4; + } + } + } + + pub fn clone_into(&self, buf: &mut Wtf8Buf) { + self.bytes.clone_into(&mut buf.bytes); + } + + /// Boxes this `Wtf8`. + #[inline] + pub fn into_box(&self) -> Box { + let boxed: Box<[u8]> = self.bytes.into(); + unsafe { mem::transmute(boxed) } + } + + /// Creates a boxed, empty `Wtf8`. + pub fn empty_box() -> Box { + let boxed: Box<[u8]> = Default::default(); + unsafe { mem::transmute(boxed) } + } + + #[inline] + pub fn make_ascii_lowercase(&mut self) { + self.bytes.make_ascii_lowercase() + } + + #[inline] + pub fn make_ascii_uppercase(&mut self) { + self.bytes.make_ascii_uppercase() + } + + #[inline] + pub fn to_ascii_lowercase(&self) -> Wtf8Buf { + Wtf8Buf { + bytes: self.bytes.to_ascii_lowercase(), + } + } + + #[inline] + pub fn to_ascii_uppercase(&self) -> Wtf8Buf { + Wtf8Buf { + bytes: self.bytes.to_ascii_uppercase(), + } + } + + #[inline] + pub fn is_ascii(&self) -> bool { + self.bytes.is_ascii() + } + + #[inline] + pub fn is_utf8(&self) -> bool { + self.next_surrogate(0).is_none() + } + + #[inline] + pub fn eq_ignore_ascii_case(&self, other: &Self) -> bool { + self.bytes.eq_ignore_ascii_case(&other.bytes) + } + + pub fn split(&self, pat: &Wtf8) -> impl Iterator { + self.as_bytes() + .split_str(pat) + .map(|w| unsafe { Wtf8::from_bytes_unchecked(w) }) + } + + pub fn splitn(&self, n: usize, pat: &Wtf8) -> impl Iterator { + self.as_bytes() + .splitn_str(n, pat) + .map(|w| unsafe { Wtf8::from_bytes_unchecked(w) }) + } + + pub fn rsplit(&self, pat: &Wtf8) -> impl Iterator { + self.as_bytes() + .rsplit_str(pat) + .map(|w| unsafe { Wtf8::from_bytes_unchecked(w) }) + } + + pub fn rsplitn(&self, n: usize, pat: &Wtf8) -> impl Iterator { + self.as_bytes() + .rsplitn_str(n, pat) + .map(|w| unsafe { Wtf8::from_bytes_unchecked(w) }) + } + + pub fn trim_start_matches(&self, f: impl Fn(CodePoint) -> bool) -> &Self { + let mut iter = self.code_points(); + loop { + let old = iter.clone(); + match iter.next().map(&f) { + Some(true) => continue, + Some(false) => { + iter = old; + break; + } + None => return iter.as_wtf8(), + } + } + iter.as_wtf8() + } + + pub fn trim_end_matches(&self, f: impl Fn(CodePoint) -> bool) -> &Self { + let mut iter = self.code_points(); + loop { + let old = iter.clone(); + match iter.next_back().map(&f) { + Some(true) => continue, + Some(false) => { + iter = old; + break; + } + None => return iter.as_wtf8(), + } + } + iter.as_wtf8() + } + + pub fn trim_matches(&self, f: impl Fn(CodePoint) -> bool) -> &Self { + self.trim_start_matches(&f).trim_end_matches(&f) + } + + pub fn find(&self, pat: &Wtf8) -> Option { + memchr::memmem::find(self.as_bytes(), pat.as_bytes()) + } + + pub fn rfind(&self, pat: &Wtf8) -> Option { + memchr::memmem::rfind(self.as_bytes(), pat.as_bytes()) + } + + pub fn get(&self, range: impl ops::RangeBounds) -> Option<&Self> { + let start = match range.start_bound() { + ops::Bound::Included(&i) => i, + ops::Bound::Excluded(&i) => i.saturating_add(1), + ops::Bound::Unbounded => 0, + }; + let end = match range.end_bound() { + ops::Bound::Included(&i) => i.saturating_add(1), + ops::Bound::Excluded(&i) => i, + ops::Bound::Unbounded => self.len(), + }; + // is_code_point_boundary checks that the index is in [0, .len()] + if start <= end && is_code_point_boundary(self, start) && is_code_point_boundary(self, end) + { + Some(unsafe { slice_unchecked(self, start, end) }) + } else { + None + } + } +} + +impl AsRef for str { + fn as_ref(&self) -> &Wtf8 { + unsafe { Wtf8::from_bytes_unchecked(self.as_bytes()) } + } +} + +impl AsRef<[u8]> for Wtf8 { + fn as_ref(&self) -> &[u8] { + self.as_bytes() + } +} + +/// Returns a slice of the given string for the byte range \[`begin`..`end`). +/// +/// # Panics +/// +/// Panics when `begin` and `end` do not point to code point boundaries, +/// or point beyond the end of the string. +impl ops::Index> for Wtf8 { + type Output = Wtf8; + + #[inline] + fn index(&self, range: ops::Range) -> &Wtf8 { + // is_code_point_boundary checks that the index is in [0, .len()] + if range.start <= range.end + && is_code_point_boundary(self, range.start) + && is_code_point_boundary(self, range.end) + { + unsafe { slice_unchecked(self, range.start, range.end) } + } else { + slice_error_fail(self, range.start, range.end) + } + } +} + +/// Returns a slice of the given string from byte `begin` to its end. +/// +/// # Panics +/// +/// Panics when `begin` is not at a code point boundary, +/// or is beyond the end of the string. +impl ops::Index> for Wtf8 { + type Output = Wtf8; + + #[inline] + fn index(&self, range: ops::RangeFrom) -> &Wtf8 { + // is_code_point_boundary checks that the index is in [0, .len()] + if is_code_point_boundary(self, range.start) { + unsafe { slice_unchecked(self, range.start, self.len()) } + } else { + slice_error_fail(self, range.start, self.len()) + } + } +} + +/// Returns a slice of the given string from its beginning to byte `end`. +/// +/// # Panics +/// +/// Panics when `end` is not at a code point boundary, +/// or is beyond the end of the string. +impl ops::Index> for Wtf8 { + type Output = Wtf8; + + #[inline] + fn index(&self, range: ops::RangeTo) -> &Wtf8 { + // is_code_point_boundary checks that the index is in [0, .len()] + if is_code_point_boundary(self, range.end) { + unsafe { slice_unchecked(self, 0, range.end) } + } else { + slice_error_fail(self, 0, range.end) + } + } +} + +impl ops::Index for Wtf8 { + type Output = Wtf8; + + #[inline] + fn index(&self, _range: ops::RangeFull) -> &Wtf8 { + self + } +} + +#[inline] +fn decode_surrogate(second_byte: u8, third_byte: u8) -> u16 { + // The first byte is assumed to be 0xED + 0xD800 | (second_byte as u16 & 0x3F) << 6 | third_byte as u16 & 0x3F +} + +/// Copied from str::is_char_boundary +#[inline] +pub fn is_code_point_boundary(slice: &Wtf8, index: usize) -> bool { + if index == 0 { + return true; + } + match slice.bytes.get(index) { + None => index == slice.len(), + Some(&b) => (b as i8) >= -0x40, + } +} + +/// Verify that `index` is at the edge of either a valid UTF-8 codepoint +/// (i.e. a codepoint that's not a surrogate) or of the whole string. +/// +/// These are the cases currently permitted by `OsStr::slice_encoded_bytes`. +/// Splitting between surrogates is valid as far as WTF-8 is concerned, but +/// we do not permit it in the public API because WTF-8 is considered an +/// implementation detail. +#[track_caller] +#[inline] +pub fn check_utf8_boundary(slice: &Wtf8, index: usize) { + if index == 0 { + return; + } + match slice.bytes.get(index) { + Some(0xED) => (), // Might be a surrogate + Some(&b) if (b as i8) >= -0x40 => return, + Some(_) => panic!("byte index {index} is not a codepoint boundary"), + None if index == slice.len() => return, + None => panic!("byte index {index} is out of bounds"), + } + if slice.bytes[index + 1] >= 0xA0 { + // There's a surrogate after index. Now check before index. + if index >= 3 && slice.bytes[index - 3] == 0xED && slice.bytes[index - 2] >= 0xA0 { + panic!("byte index {index} lies between surrogate codepoints"); + } + } +} + +/// Copied from core::str::raw::slice_unchecked +/// +/// # Safety +/// +/// `begin` and `end` must be within bounds and on codepoint boundaries. +#[inline] +pub unsafe fn slice_unchecked(s: &Wtf8, begin: usize, end: usize) -> &Wtf8 { + // SAFETY: memory layout of a &[u8] and &Wtf8 are the same + unsafe { + let len = end - begin; + let start = s.as_bytes().as_ptr().add(begin); + Wtf8::from_bytes_unchecked(slice::from_raw_parts(start, len)) + } +} + +/// Copied from core::str::raw::slice_error_fail +#[inline(never)] +pub fn slice_error_fail(s: &Wtf8, begin: usize, end: usize) -> ! { + assert!(begin <= end); + panic!("index {begin} and/or {end} in `{s:?}` do not lie on character boundary"); +} + +/// Iterator for the code points of a WTF-8 string. +/// +/// Created with the method `.code_points()`. +#[derive(Clone)] +pub struct Wtf8CodePoints<'a> { + bytes: slice::Iter<'a, u8>, +} + +impl Iterator for Wtf8CodePoints<'_> { + type Item = CodePoint; + + #[inline] + fn next(&mut self) -> Option { + // SAFETY: `self.bytes` has been created from a WTF-8 string + unsafe { next_code_point(&mut self.bytes).map(|c| CodePoint { value: c }) } + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + let len = self.bytes.len(); + (len.saturating_add(3) / 4, Some(len)) + } + + fn last(mut self) -> Option { + self.next_back() + } +} + +impl DoubleEndedIterator for Wtf8CodePoints<'_> { + #[inline] + fn next_back(&mut self) -> Option { + // SAFETY: `str` invariant says `self.iter` is a valid WTF-8 string and + // the resulting `ch` is a valid Unicode Code Point. + unsafe { + next_code_point_reverse(&mut self.bytes).map(|ch| CodePoint::from_u32_unchecked(ch)) + } + } +} + +impl<'a> Wtf8CodePoints<'a> { + pub fn as_wtf8(&self) -> &'a Wtf8 { + unsafe { Wtf8::from_bytes_unchecked(self.bytes.as_slice()) } + } +} + +#[derive(Clone)] +pub struct Wtf8CodePointIndices<'a> { + pub(super) front_offset: usize, + pub(super) iter: Wtf8CodePoints<'a>, +} + +impl Iterator for Wtf8CodePointIndices<'_> { + type Item = (usize, CodePoint); + + #[inline] + fn next(&mut self) -> Option<(usize, CodePoint)> { + let pre_len = self.iter.bytes.len(); + match self.iter.next() { + None => None, + Some(ch) => { + let index = self.front_offset; + let len = self.iter.bytes.len(); + self.front_offset += pre_len - len; + Some((index, ch)) + } + } + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + self.iter.size_hint() + } + + #[inline] + fn last(mut self) -> Option<(usize, CodePoint)> { + // No need to go through the entire string. + self.next_back() + } +} + +impl DoubleEndedIterator for Wtf8CodePointIndices<'_> { + #[inline] + fn next_back(&mut self) -> Option<(usize, CodePoint)> { + self.iter.next_back().map(|ch| { + let index = self.front_offset + self.iter.bytes.len(); + (index, ch) + }) + } +} + +impl FusedIterator for Wtf8CodePointIndices<'_> {} + +/// Generates a wide character sequence for potentially ill-formed UTF-16. +#[derive(Clone)] +pub struct EncodeWide<'a> { + code_points: Wtf8CodePoints<'a>, + extra: u16, +} + +// Copied from libunicode/u_str.rs +impl Iterator for EncodeWide<'_> { + type Item = u16; + + #[inline] + fn next(&mut self) -> Option { + if self.extra != 0 { + let tmp = self.extra; + self.extra = 0; + return Some(tmp); + } + + let mut buf = [0; MAX_LEN_UTF16]; + self.code_points.next().map(|code_point| { + let n = encode_utf16_raw(code_point.value, &mut buf).len(); + if n == 2 { + self.extra = buf[1]; + } + buf[0] + }) + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + let (low, high) = self.code_points.size_hint(); + let ext = (self.extra != 0) as usize; + // every code point gets either one u16 or two u16, + // so this iterator is between 1 or 2 times as + // long as the underlying iterator. + ( + low + ext, + high.and_then(|n| n.checked_mul(2)) + .and_then(|n| n.checked_add(ext)), + ) + } +} + +impl FusedIterator for EncodeWide<'_> {} + +pub struct Wtf8Chunks<'a> { + wtf8: &'a Wtf8, +} + +impl<'a> Iterator for Wtf8Chunks<'a> { + type Item = Wtf8Chunk<'a>; + + fn next(&mut self) -> Option { + match self.wtf8.next_surrogate(0) { + Some((0, surrogate)) => { + self.wtf8 = &self.wtf8[3..]; + Some(Wtf8Chunk::Surrogate(surrogate.into())) + } + Some((n, _)) => { + let s = unsafe { str::from_utf8_unchecked(&self.wtf8.as_bytes()[..n]) }; + self.wtf8 = &self.wtf8[n..]; + Some(Wtf8Chunk::Utf8(s)) + } + None => { + let s = + unsafe { str::from_utf8_unchecked(std::mem::take(&mut self.wtf8).as_bytes()) }; + (!s.is_empty()).then_some(Wtf8Chunk::Utf8(s)) + } + } + } +} + +pub enum Wtf8Chunk<'a> { + Utf8(&'a str), + Surrogate(CodePoint), +} + +impl Hash for CodePoint { + #[inline] + fn hash(&self, state: &mut H) { + self.value.hash(state) + } +} + +// == BOX IMPLS == + +/// # Safety +/// +/// `value` must be valid WTF-8. +pub unsafe fn from_boxed_wtf8_unchecked(value: Box<[u8]>) -> Box { + unsafe { Box::from_raw(Box::into_raw(value) as *mut Wtf8) } +} + +impl Clone for Box { + fn clone(&self) -> Self { + (&**self).into() + } +} + +impl Default for Box { + fn default() -> Self { + unsafe { from_boxed_wtf8_unchecked(Box::default()) } + } +} + +impl From<&Wtf8> for Box { + fn from(w: &Wtf8) -> Self { + w.into_box() + } +} + +impl From<&str> for Box { + fn from(s: &str) -> Self { + Box::::from(s).into() + } +} + +impl From> for Box { + fn from(s: Box) -> Self { + unsafe { from_boxed_wtf8_unchecked(s.into_boxed_bytes()) } + } +} + +impl From> for Box { + fn from(s: Box) -> Self { + >::from(s).into() + } +} + +impl From> for Box<[u8]> { + fn from(w: Box) -> Self { + unsafe { Box::from_raw(Box::into_raw(w) as *mut [u8]) } + } +} + +impl From for Box { + fn from(w: Wtf8Buf) -> Self { + w.into_box() + } +} + +impl From for Box { + fn from(s: String) -> Self { + s.into_boxed_str().into() + } +} diff --git a/stdlib/src/pystruct.rs b/stdlib/src/pystruct.rs index f8d41414f..702e5b668 100644 --- a/stdlib/src/pystruct.rs +++ b/stdlib/src/pystruct.rs @@ -27,23 +27,20 @@ pub(crate) mod _struct { fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult { // CPython turns str to bytes but we do reversed way here // The only performance difference is this transition cost - let fmt = match_class! { - match obj { - s @ PyStr => if s.is_ascii() { - Some(s) - } else { - None - }, - b @ PyBytes => if b.is_ascii() { - Some(unsafe { - PyStr::new_ascii_unchecked(b.as_bytes().to_vec()) - }.into_ref(&vm.ctx)) - } else { - None - }, - other => return Err(vm.new_type_error(format!("Struct() argument 1 must be a str or bytes object, not {}", other.class().name()))), - } - }.ok_or_else(|| vm.new_unicode_decode_error("Struct format must be a ascii string".to_owned()))?; + let fmt = match_class!(match obj { + s @ PyStr => s.is_ascii().then_some(s), + b @ PyBytes => ascii::AsciiStr::from_ascii(&b) + .ok() + .map(|s| vm.ctx.new_str(s)), + other => + return Err(vm.new_type_error(format!( + "Struct() argument 1 must be a str or bytes object, not {}", + other.class().name() + ))), + }) + .ok_or_else(|| { + vm.new_unicode_decode_error("Struct format must be a ascii string".to_owned()) + })?; Ok(IntoStructFormatBytes(fmt)) } } diff --git a/stdlib/src/unicodedata.rs b/stdlib/src/unicodedata.rs index 49f3ef625..9af921d36 100644 --- a/stdlib/src/unicodedata.rs +++ b/stdlib/src/unicodedata.rs @@ -65,6 +65,7 @@ mod unicodedata { function::OptionalArg, }; use itertools::Itertools; + use rustpython_common::wtf8::{CodePoint, Wtf8Buf}; use ucd::{Codepoint, EastAsianWidth}; use unic_char_property::EnumeratedCharProperty; use unic_normal::StrNormalForm; @@ -84,14 +85,23 @@ mod unicodedata { Self { unic_version } } - fn check_age(&self, c: char) -> bool { - Age::of(c).is_some_and(|age| age.actual() <= self.unic_version) + fn check_age(&self, c: CodePoint) -> bool { + c.to_char() + .is_none_or(|c| Age::of(c).is_some_and(|age| age.actual() <= self.unic_version)) } - fn extract_char(&self, character: PyStrRef, vm: &VirtualMachine) -> PyResult> { - let c = character.as_str().chars().exactly_one().map_err(|_| { - vm.new_type_error("argument must be an unicode character, not str".to_owned()) - })?; + fn extract_char( + &self, + character: PyStrRef, + vm: &VirtualMachine, + ) -> PyResult> { + let c = character + .as_wtf8() + .code_points() + .exactly_one() + .map_err(|_| { + vm.new_type_error("argument must be an unicode character, not str".to_owned()) + })?; Ok(self.check_age(c).then_some(c)) } @@ -103,7 +113,10 @@ mod unicodedata { fn category(&self, character: PyStrRef, vm: &VirtualMachine) -> PyResult { Ok(self .extract_char(character, vm)? - .map_or(GeneralCategory::Unassigned, GeneralCategory::of) + .map_or(GeneralCategory::Unassigned, |c| { + c.to_char() + .map_or(GeneralCategory::Surrogate, GeneralCategory::of) + }) .abbr_name() .to_owned()) } @@ -111,7 +124,7 @@ mod unicodedata { #[pymethod] fn lookup(&self, name: PyStrRef, vm: &VirtualMachine) -> PyResult { if let Some(character) = unicode_names2::character(name.as_str()) { - if self.check_age(character) { + if self.check_age(character.into()) { return Ok(character.to_string()); } } @@ -129,7 +142,7 @@ mod unicodedata { if let Some(c) = c { if self.check_age(c) { - if let Some(name) = unicode_names2::name(c) { + if let Some(name) = c.to_char().and_then(unicode_names2::name) { return Ok(vm.ctx.new_str(name.to_string()).into()); } } @@ -144,7 +157,10 @@ mod unicodedata { vm: &VirtualMachine, ) -> PyResult<&'static str> { let bidi = match self.extract_char(character, vm)? { - Some(c) => BidiClass::of(c).abbr_name(), + Some(c) => c + .to_char() + .map_or(BidiClass::LeftToRight, BidiClass::of) + .abbr_name(), None => "", }; Ok(bidi) @@ -159,19 +175,20 @@ mod unicodedata { ) -> PyResult<&'static str> { Ok(self .extract_char(character, vm)? + .and_then(|c| c.to_char()) .map_or(EastAsianWidth::Neutral, |c| c.east_asian_width()) .abbr_name()) } #[pymethod] - fn normalize(&self, form: super::NormalizeForm, unistr: PyStrRef) -> PyResult { + fn normalize(&self, form: super::NormalizeForm, unistr: PyStrRef) -> PyResult { use super::NormalizeForm::*; - let text = unistr.as_str(); + let text = unistr.as_wtf8(); let normalized_text = match form { - Nfc => text.nfc().collect::(), - Nfkc => text.nfkc().collect::(), - Nfd => text.nfd().collect::(), - Nfkd => text.nfkd().collect::(), + Nfc => text.map_utf8(|s| s.nfc()).collect(), + Nfkc => text.map_utf8(|s| s.nfkc()).collect(), + Nfd => text.map_utf8(|s| s.nfd()).collect(), + Nfkd => text.map_utf8(|s| s.nfkd()).collect(), }; Ok(normalized_text) } diff --git a/vm/src/anystr.rs b/vm/src/anystr.rs index d01136b0f..0a884fc73 100644 --- a/vm/src/anystr.rs +++ b/vm/src/anystr.rs @@ -7,28 +7,13 @@ use crate::{ use num_traits::{cast::ToPrimitive, sign::Signed}; #[derive(FromArgs)] -pub struct SplitArgs { +pub struct SplitArgs { #[pyarg(any, default)] sep: Option, #[pyarg(any, default = "-1")] maxsplit: isize, } -impl SplitArgs { - pub fn get_value(self, vm: &VirtualMachine) -> PyResult<(Option, isize)> { - let sep = if let Some(s) = self.sep { - let sep = s.as_ref(); - if sep.is_empty() { - return Err(vm.new_value_error("empty separator".to_owned())); - } - Some(s) - } else { - None - }; - Ok((sep, self.maxsplit)) - } -} - #[derive(FromArgs)] pub struct SplitLinesArgs { #[pyarg(any, default = "false")] @@ -132,9 +117,9 @@ impl StringRange for std::ops::Range { } } -pub trait AnyStrWrapper { - type Str: ?Sized + AnyStr; - fn as_ref(&self) -> &Self::Str; +pub trait AnyStrWrapper { + fn as_ref(&self) -> Option<&S>; + fn is_empty(&self) -> bool; } pub trait AnyStrContainer @@ -146,15 +131,18 @@ where fn push_str(&mut self, s: &S); } -pub trait AnyStr { - type Char: Copy; - type Container: AnyStrContainer + Extend; +pub trait AnyChar: Copy { + fn is_lowercase(self) -> bool; + fn is_uppercase(self) -> bool; + fn bytes_len(self) -> usize; +} - fn element_bytes_len(c: Self::Char) -> usize; +pub trait AnyStr { + type Char: AnyChar; + type Container: AnyStrContainer + Extend; fn to_container(&self) -> Self::Container; fn as_bytes(&self) -> &[u8]; - fn chars(&self) -> impl Iterator; fn elements(&self) -> impl Iterator; fn get_bytes(&self, range: std::ops::Range) -> &Self; // FIXME: get_chars is expensive for str @@ -172,29 +160,35 @@ pub trait AnyStr { new } - fn py_split( + fn py_split( &self, args: SplitArgs, vm: &VirtualMachine, + full_obj: impl FnOnce() -> PyObjectRef, split: SP, splitn: SN, splitw: SW, - ) -> PyResult> + ) -> PyResult> where - T: TryFromObject + AnyStrWrapper, - SP: Fn(&Self, &Self, &VirtualMachine) -> Vec, - SN: Fn(&Self, &Self, usize, &VirtualMachine) -> Vec, - SW: Fn(&Self, isize, &VirtualMachine) -> Vec, + T: TryFromObject + AnyStrWrapper, + SP: Fn(&Self, &Self, &VirtualMachine) -> Vec, + SN: Fn(&Self, &Self, usize, &VirtualMachine) -> Vec, + SW: Fn(&Self, isize, &VirtualMachine) -> Vec, { - let (sep, maxsplit) = args.get_value(vm)?; - let splits = if let Some(pattern) = sep { - if maxsplit < 0 { - split(self, pattern.as_ref(), vm) + if args.sep.as_ref().is_some_and(|sep| sep.is_empty()) { + return Err(vm.new_value_error("empty separator".to_owned())); + } + let splits = if let Some(pattern) = args.sep { + let Some(pattern) = pattern.as_ref() else { + return Ok(vec![full_obj()]); + }; + if args.maxsplit < 0 { + split(self, pattern, vm) } else { - splitn(self, pattern.as_ref(), (maxsplit + 1) as usize, vm) + splitn(self, pattern, (args.maxsplit + 1) as usize, vm) } } else { - splitw(self, maxsplit, vm) + splitw(self, args.maxsplit, vm) }; Ok(splits) } @@ -242,13 +236,19 @@ pub trait AnyStr { func_default: FD, ) -> &'a Self where - S: AnyStrWrapper, + S: AnyStrWrapper, FC: Fn(&'a Self, &Self) -> &'a Self, FD: Fn(&'a Self) -> &'a Self, { let chars = chars.flatten(); match chars { - Some(chars) => func_chars(self, chars.as_ref()), + Some(chars) => { + if let Some(chars) = chars.as_ref() { + func_chars(self, chars) + } else { + self + } + } None => func_default(self), } } @@ -281,7 +281,7 @@ pub trait AnyStr { fn py_pad(&self, left: usize, right: usize, fillchar: Self::Char) -> Self::Container { let mut u = Self::Container::with_capacity( - (left + right) * Self::element_bytes_len(fillchar) + self.bytes_len(), + (left + right) * fillchar.bytes_len() + self.bytes_len(), ); u.extend(std::iter::repeat(fillchar).take(left)); u.push_str(self); @@ -305,19 +305,17 @@ pub trait AnyStr { fn py_join( &self, - mut iter: impl std::iter::Iterator< - Item = PyResult + TryFromObject>, - >, + mut iter: impl std::iter::Iterator + TryFromObject>>, ) -> PyResult { let mut joined = if let Some(elem) = iter.next() { - elem?.as_ref().to_container() + elem?.as_ref().unwrap().to_container() } else { return Ok(Self::Container::new()); }; for elem in iter { let elem = elem?; joined.push_str(self); - joined.push_str(elem.as_ref()); + joined.push_str(elem.as_ref().unwrap()); } Ok(joined) } @@ -403,25 +401,33 @@ pub trait AnyStr { rustpython_common::str::zfill(self.as_bytes(), width) } - fn py_iscase(&self, is_case: F, is_opposite: G) -> bool - where - F: Fn(char) -> bool, - G: Fn(char) -> bool, - { - // Unified form of CPython functions: - // _Py_bytes_islower - // Py_bytes_isupper - // unicode_islower_impl - // unicode_isupper_impl - let mut cased = false; - for c in self.chars() { - if is_opposite(c) { + // Unified form of CPython functions: + // _Py_bytes_islower + // Py_bytes_isupper + // unicode_islower_impl + // unicode_isupper_impl + fn py_islower(&self) -> bool { + let mut lower = false; + for c in self.elements() { + if c.is_uppercase() { return false; - } else if !cased && is_case(c) { - cased = true + } else if !lower && c.is_lowercase() { + lower = true } } - cased + lower + } + + fn py_isupper(&self) -> bool { + let mut upper = false; + for c in self.elements() { + if c.is_lowercase() { + return false; + } else if !upper && c.is_uppercase() { + upper = true + } + } + upper } } diff --git a/vm/src/builtins/genericalias.rs b/vm/src/builtins/genericalias.rs index c03e3145b..0e0a34227 100644 --- a/vm/src/builtins/genericalias.rs +++ b/vm/src/builtins/genericalias.rs @@ -124,7 +124,7 @@ impl PyGenericAlias { Ok(format!( "{}[{}]", repr_item(self.origin.clone().into(), vm)?, - if self.args.len() == 0 { + if self.args.is_empty() { "()".to_owned() } else { self.args @@ -261,23 +261,20 @@ fn subs_tvars( .and_then(|sub_params| { PyTupleRef::try_from_object(vm, sub_params) .ok() - .and_then(|sub_params| { - if sub_params.len() > 0 { - let sub_args = sub_params - .iter() - .map(|arg| { - if let Some(idx) = tuple_index(params, arg) { - argitems[idx].clone() - } else { - arg.clone() - } - }) - .collect::>(); - let sub_args: PyObjectRef = PyTuple::new_ref(sub_args, &vm.ctx).into(); - Some(obj.get_item(&*sub_args, vm)) - } else { - None - } + .filter(|sub_params| !sub_params.is_empty()) + .map(|sub_params| { + let sub_args = sub_params + .iter() + .map(|arg| { + if let Some(idx) = tuple_index(params, arg) { + argitems[idx].clone() + } else { + arg.clone() + } + }) + .collect::>(); + let sub_args: PyObjectRef = PyTuple::new_ref(sub_args, &vm.ctx).into(); + obj.get_item(&*sub_args, vm) }) }) .unwrap_or(Ok(obj)) diff --git a/vm/src/builtins/object.rs b/vm/src/builtins/object.rs index cce1422d5..be1432754 100644 --- a/vm/src/builtins/object.rs +++ b/vm/src/builtins/object.rs @@ -159,7 +159,7 @@ fn object_getstate_default(obj: &PyObject, required: bool, vm: &VirtualMachine) slots.set_item(name.as_str(), value, vm).unwrap(); } - if slots.len() > 0 { + if !slots.is_empty() { return (state, slots).to_pyresult(vm); } } diff --git a/vm/src/builtins/set.rs b/vm/src/builtins/set.rs index a135af1bd..40823aa37 100644 --- a/vm/src/builtins/set.rs +++ b/vm/src/builtins/set.rs @@ -784,9 +784,7 @@ impl Initializer for PySet { type Args = OptionalArg; fn init(zelf: PyRef, iterable: Self::Args, vm: &VirtualMachine) -> PyResult<()> { - if zelf.len() > 0 { - zelf.clear(); - } + zelf.clear(); if let OptionalArg::Present(it) = iterable { zelf.update(PosArgs::new(vec![it]), vm)?; } diff --git a/vm/src/builtins/str.rs b/vm/src/builtins/str.rs index 76cdca81e..6ff079c64 100644 --- a/vm/src/builtins/str.rs +++ b/vm/src/builtins/str.rs @@ -10,7 +10,7 @@ use crate::{ atomic_func, cformat::cformat_string, class::PyClassImpl, - common::str::{BorrowedStr, PyStrKind, PyStrKindData}, + common::str::{PyKindStr, StrData, StrKind}, convert::{IntoPyException, ToPyException, ToPyObject, ToPyResult}, format::{format, format_map}, function::{ArgIterable, ArgSize, FuncArgs, OptionalArg, OptionalOption, PyComparisonValue}, @@ -24,7 +24,7 @@ use crate::{ PyComparisonOp, Representable, SelfIter, Unconstructible, }, }; -use ascii::{AsciiStr, AsciiString}; +use ascii::{AsciiChar, AsciiStr, AsciiString}; use bstr::ByteSlice; use itertools::Itertools; use num_traits::ToPrimitive; @@ -35,8 +35,10 @@ use rustpython_common::{ format::{FormatSpec, FormatString, FromTemplate}, hash, lock::PyMutex, + str::DeduceStrKind, + wtf8::{CodePoint, Wtf8, Wtf8Buf, Wtf8Chunk}, }; -use std::{char, fmt, ops::Range, string::ToString}; +use std::{char, fmt, ops::Range}; use unic_ucd_bidi::BidiClass; use unic_ucd_category::GeneralCategory; use unic_ucd_ident::{is_xid_continue, is_xid_start}; @@ -57,39 +59,59 @@ impl<'a> TryFromBorrowedObject<'a> for &'a str { #[pyclass(module = false, name = "str")] pub struct PyStr { - bytes: Box<[u8]>, - kind: PyStrKindData, + data: StrData, hash: PyAtomic, } impl fmt::Debug for PyStr { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("PyStr") - .field("value", &self.as_str()) - .field("kind", &self.kind) + .field("value", &self.as_wtf8()) + .field("kind", &self.data.kind()) .field("hash", &self.hash) .finish() } } impl AsRef for PyStr { + #[track_caller] // <- can remove this once it doesn't panic fn as_ref(&self) -> &str { self.as_str() } } impl AsRef for Py { + #[track_caller] // <- can remove this once it doesn't panic fn as_ref(&self) -> &str { self.as_str() } } impl AsRef for PyStrRef { + #[track_caller] // <- can remove this once it doesn't panic fn as_ref(&self) -> &str { self.as_str() } } +impl AsRef for PyStr { + fn as_ref(&self) -> &Wtf8 { + self.as_wtf8() + } +} + +impl AsRef for Py { + fn as_ref(&self) -> &Wtf8 { + self.as_wtf8() + } +} + +impl AsRef for PyStrRef { + fn as_ref(&self) -> &Wtf8 { + self.as_wtf8() + } +} + impl<'a> From<&'a AsciiStr> for PyStr { fn from(s: &'a AsciiStr) -> Self { s.to_owned().into() @@ -98,7 +120,19 @@ impl<'a> From<&'a AsciiStr> for PyStr { impl From for PyStr { fn from(s: AsciiString) -> Self { - unsafe { Self::new_ascii_unchecked(s.into()) } + s.into_boxed_ascii_str().into() + } +} + +impl From> for PyStr { + fn from(s: Box) -> Self { + StrData::from(s).into() + } +} + +impl From for PyStr { + fn from(ch: AsciiChar) -> Self { + AsciiString::from(ch).into() } } @@ -108,12 +142,45 @@ impl<'a> From<&'a str> for PyStr { } } +impl<'a> From<&'a Wtf8> for PyStr { + fn from(s: &'a Wtf8) -> Self { + s.to_owned().into() + } +} + impl From for PyStr { fn from(s: String) -> Self { s.into_boxed_str().into() } } +impl From for PyStr { + fn from(w: Wtf8Buf) -> Self { + w.into_box().into() + } +} + +impl From for PyStr { + fn from(ch: char) -> Self { + StrData::from(ch).into() + } +} + +impl From for PyStr { + fn from(ch: CodePoint) -> Self { + StrData::from(ch).into() + } +} + +impl From for PyStr { + fn from(data: StrData) -> Self { + PyStr { + data, + hash: Radium::new(hash::SENTINEL), + } + } +} + impl<'a> From> for PyStr { fn from(s: std::borrow::Cow<'a, str>) -> Self { s.into_owned().into() @@ -123,19 +190,21 @@ impl<'a> From> for PyStr { impl From> for PyStr { #[inline] fn from(value: Box) -> Self { - // doing the check is ~10x faster for ascii, and is actually only 2% slower worst case for - // non-ascii; see https://github.com/RustPython/RustPython/pull/2586#issuecomment-844611532 - let is_ascii = value.is_ascii(); - let bytes = value.into_boxed_bytes(); - let kind = if is_ascii { - PyStrKind::Ascii - } else { - PyStrKind::Utf8 - } - .new_data(); + StrData::from(value).into() + } +} + +impl From> for PyStr { + #[inline] + fn from(value: Box) -> Self { + StrData::from(value).into() + } +} + +impl Default for PyStr { + fn default() -> Self { Self { - bytes, - kind, + data: StrData::default(), hash: Radium::new(hash::SENTINEL), } } @@ -146,7 +215,7 @@ pub type PyStrRef = PyRef; impl fmt::Display for PyStr { #[inline] fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - fmt::Display::fmt(self.as_str(), f) + self.as_wtf8().fmt(f) } } @@ -237,18 +306,18 @@ impl IterNext for PyStrIterator { let mut internal = zelf.internal.lock(); if let IterStatus::Active(s) = &internal.0.status { - let value = s.as_str(); + let value = s.as_wtf8(); if internal.1 == usize::MAX { - if let Some((offset, ch)) = value.char_indices().nth(internal.0.position) { + if let Some((offset, ch)) = value.code_point_indices().nth(internal.0.position) { internal.0.position += 1; - internal.1 = offset + ch.len_utf8(); + internal.1 = offset + ch.len_wtf8(); return Ok(PyIterReturn::Return(ch.to_pyobject(vm))); } } else if let Some(value) = value.get(internal.1..) { - if let Some(ch) = value.chars().next() { + if let Some(ch) = value.code_points().next() { internal.0.position += 1; - internal.1 += ch.len_utf8(); + internal.1 += ch.len_wtf8(); return Ok(PyIterReturn::Return(ch.to_pyobject(vm))); } } @@ -292,7 +361,7 @@ impl Constructor for PyStr { if string.class().is(&cls) { Ok(string.into()) } else { - PyStr::from(string.as_str()) + PyStr::from(string.as_wtf8()) .into_ref_with_type(vm, cls) .map(Into::into) } @@ -301,20 +370,19 @@ impl Constructor for PyStr { impl PyStr { /// # Safety: Given `bytes` must be valid data for given `kind` - pub(crate) unsafe fn new_str_unchecked(bytes: Vec, kind: PyStrKind) -> Self { - let s = Self { - bytes: bytes.into_boxed_slice(), - kind: kind.new_data(), - hash: Radium::new(hash::SENTINEL), - }; - debug_assert!(matches!(s.kind, PyStrKindData::Ascii) || !s.as_str().is_ascii()); - s + unsafe fn new_str_unchecked(data: Box, kind: StrKind) -> Self { + unsafe { StrData::new_str_unchecked(data, kind) }.into() + } + + unsafe fn new_with_char_len>>(s: T, char_len: usize) -> Self { + let kind = s.str_kind(); + unsafe { StrData::new_with_char_len(s.into(), kind, char_len) }.into() } /// # Safety /// Given `bytes` must be ascii pub unsafe fn new_ascii_unchecked(bytes: Vec) -> Self { - unsafe { Self::new_str_unchecked(bytes, PyStrKind::Ascii) } + unsafe { AsciiString::from_ascii_unchecked(bytes) }.into() } pub fn new_ref(zelf: impl Into, ctx: &Context) -> PyRef { @@ -322,40 +390,56 @@ impl PyStr { PyRef::new_ref(zelf, ctx.types.str_type.to_owned(), None) } - fn new_substr(&self, s: String) -> Self { - let kind = if self.kind.kind() == PyStrKind::Ascii || s.is_ascii() { - PyStrKind::Ascii + fn new_substr(&self, s: Wtf8Buf) -> Self { + let kind = if self.kind().is_ascii() || s.is_ascii() { + StrKind::Ascii + } else if self.kind().is_utf8() || s.is_utf8() { + StrKind::Utf8 } else { - PyStrKind::Utf8 + StrKind::Wtf8 }; unsafe { // SAFETY: kind is properly decided for substring - Self::new_str_unchecked(s.into_bytes(), kind) + Self::new_str_unchecked(s.into(), kind) } } #[inline] + pub fn as_wtf8(&self) -> &Wtf8 { + self.data.as_wtf8() + } + + pub fn as_bytes(&self) -> &[u8] { + self.data.as_wtf8().as_bytes() + } + + // FIXME: make this return an Option + #[inline] + #[track_caller] // <- can remove this once it doesn't panic pub fn as_str(&self) -> &str { - unsafe { - // SAFETY: Both PyStrKind::{Ascii, Utf8} are valid utf8 string - std::str::from_utf8_unchecked(&self.bytes) - } + self.data.as_str().expect("str has surrogates") + } + + pub fn kind(&self) -> StrKind { + self.data.kind() + } + + #[inline] + pub fn as_str_kind(&self) -> PyKindStr<'_> { + self.data.as_str_kind() } fn char_all(&self, test: F) -> bool where F: Fn(char) -> bool, { - match self.kind.kind() { - PyStrKind::Ascii => self.bytes.iter().all(|&x| test(char::from(x))), - PyStrKind::Utf8 => self.as_str().chars().all(test), + match self.as_str_kind() { + PyKindStr::Ascii(s) => s.chars().all(|ch| test(ch.into())), + PyKindStr::Utf8(s) => s.chars().all(test), + PyKindStr::Wtf8(w) => w.code_points().all(|ch| ch.is_char_and(&test)), } } - fn borrow(&self) -> &BorrowedStr<'_> { - unsafe { std::mem::transmute(self) } - } - fn repeat(zelf: PyRef, value: isize, vm: &VirtualMachine) -> PyResult> { if value == 0 && zelf.class().is(vm.ctx.types.str_type) { // Special case: when some `str` is multiplied by `0`, @@ -369,10 +453,10 @@ impl PyStr { // This only works for `str` itself, not its subclasses. return Ok(zelf); } - zelf.as_str() + zelf.as_wtf8() .as_bytes() .mul(vm, value) - .map(|x| Self::from(unsafe { String::from_utf8_unchecked(x) }).into_ref(&vm.ctx)) + .map(|x| Self::from(unsafe { Wtf8Buf::from_bytes_unchecked(x) }).into_ref(&vm.ctx)) } } @@ -394,11 +478,11 @@ impl PyStr { #[pymethod(magic)] fn add(zelf: PyRef, other: PyObjectRef, vm: &VirtualMachine) -> PyResult { if let Some(other) = other.payload::() { - let bytes = zelf.as_str().py_add(other.as_ref()); + let bytes = zelf.as_wtf8().py_add(other.as_wtf8()); Ok(unsafe { // SAFETY: `kind` is safely decided - let kind = zelf.kind.kind() | other.kind.kind(); - Self::new_str_unchecked(bytes.into_bytes(), kind) + let kind = zelf.kind() | other.kind(); + Self::new_str_unchecked(bytes.into(), kind) } .to_pyobject(vm)) } else if let Some(radd) = vm.get_method(other.clone(), identifier!(vm, __radd__)) { @@ -414,7 +498,7 @@ impl PyStr { fn _contains(&self, needle: &PyObject, vm: &VirtualMachine) -> PyResult { if let Some(needle) = needle.payload::() { - Ok(self.as_str().contains(needle.as_str())) + Ok(memchr::memmem::find(self.as_bytes(), needle.as_bytes()).is_some()) } else { Err(vm.new_type_error(format!( "'in ' requires string as left operand, not {}", @@ -429,11 +513,11 @@ impl PyStr { } fn _getitem(&self, needle: &PyObject, vm: &VirtualMachine) -> PyResult { - match SequenceIndex::try_from_borrowed_object(vm, needle, "str")? { - SequenceIndex::Int(i) => self.getitem_by_index(vm, i).map(|x| x.to_string()), - SequenceIndex::Slice(slice) => self.getitem_by_slice(vm, slice), - } - .map(|x| self.new_substr(x).into_ref(&vm.ctx).into()) + let item = match SequenceIndex::try_from_borrowed_object(vm, needle, "str")? { + SequenceIndex::Int(i) => self.getitem_by_index(vm, i)?.to_pyobject(vm), + SequenceIndex::Slice(slice) => self.getitem_by_slice(vm, slice)?.to_pyobject(vm), + }; + Ok(item) } #[pymethod(magic)] @@ -450,7 +534,7 @@ impl PyStr { } #[cold] fn _compute_hash(&self, vm: &VirtualMachine) -> hash::PyHash { - let hash_val = vm.state.hash_secret.hash_str(self.as_str()); + let hash_val = vm.state.hash_secret.hash_bytes(self.as_bytes()); debug_assert_ne!(hash_val, hash::SENTINEL); // like with char_len, we don't need a cmpxchg loop, since it'll always be the same value self.hash.store(hash_val, atomic::Ordering::Relaxed); @@ -459,26 +543,23 @@ impl PyStr { #[inline] pub fn byte_len(&self) -> usize { - self.bytes.len() + self.data.len() } #[inline] pub fn is_empty(&self) -> bool { - self.bytes.is_empty() + self.data.is_empty() } #[pymethod(name = "__len__")] #[inline] pub fn char_len(&self) -> usize { - self.borrow().char_len() + self.data.char_len() } #[pymethod(name = "isascii")] #[inline(always)] pub fn is_ascii(&self) -> bool { - match self.kind { - PyStrKindData::Ascii => true, - PyStrKindData::Utf8(_) => false, - } + matches!(self.kind(), StrKind::Ascii) } #[pymethod(magic)] @@ -495,6 +576,9 @@ impl PyStr { #[inline] pub(crate) fn repr(&self, vm: &VirtualMachine) -> PyResult { use crate::literal::escape::UnicodeEscape; + if !self.kind().is_utf8() { + return Ok(format!("{:?}", self.as_wtf8())); + } let escape = UnicodeEscape::new_repr(self.as_str()); escape .str_repr() @@ -503,10 +587,18 @@ impl PyStr { } #[pymethod] - fn lower(&self) -> String { - match self.kind.kind() { - PyStrKind::Ascii => self.as_str().to_ascii_lowercase(), - PyStrKind::Utf8 => self.as_str().to_lowercase(), + fn lower(&self) -> PyStr { + match self.as_str_kind() { + PyKindStr::Ascii(s) => s.to_ascii_lowercase().into(), + PyKindStr::Utf8(s) => s.to_lowercase().into(), + PyKindStr::Wtf8(w) => w + .chunks() + .map(|c| match c { + Wtf8Chunk::Utf8(s) => s.to_lowercase().into(), + Wtf8Chunk::Surrogate(c) => Wtf8Buf::from(c), + }) + .collect::() + .into(), } } @@ -517,10 +609,18 @@ impl PyStr { } #[pymethod] - fn upper(&self) -> String { - match self.kind.kind() { - PyStrKind::Ascii => self.as_str().to_ascii_uppercase(), - PyStrKind::Utf8 => self.as_str().to_uppercase(), + fn upper(&self) -> PyStr { + match self.as_str_kind() { + PyKindStr::Ascii(s) => s.to_ascii_uppercase().into(), + PyKindStr::Utf8(s) => s.to_uppercase().into(), + PyKindStr::Wtf8(w) => w + .chunks() + .map(|c| match c { + Wtf8Chunk::Utf8(s) => s.to_uppercase().into(), + Wtf8Chunk::Surrogate(c) => Wtf8Buf::from(c), + }) + .collect::() + .into(), } } @@ -539,36 +639,42 @@ impl PyStr { } #[pymethod] - fn split(&self, args: SplitArgs, vm: &VirtualMachine) -> PyResult> { - let elements = match self.kind.kind() { - PyStrKind::Ascii => self.as_str().py_split( + fn split(zelf: &Py, args: SplitArgs, vm: &VirtualMachine) -> PyResult> { + let elements = match zelf.as_str_kind() { + PyKindStr::Ascii(s) => s.py_split( args, vm, + || zelf.as_object().to_owned(), |v, s, vm| { v.as_bytes() .split_str(s) - .map(|s| { - unsafe { PyStr::new_ascii_unchecked(s.to_owned()) }.to_pyobject(vm) - }) + .map(|s| unsafe { AsciiStr::from_ascii_unchecked(s) }.to_pyobject(vm)) .collect() }, |v, s, n, vm| { v.as_bytes() .splitn_str(n, s) - .map(|s| { - unsafe { PyStr::new_ascii_unchecked(s.to_owned()) }.to_pyobject(vm) - }) + .map(|s| unsafe { AsciiStr::from_ascii_unchecked(s) }.to_pyobject(vm)) .collect() }, |v, n, vm| { v.as_bytes().py_split_whitespace(n, |s| { - unsafe { PyStr::new_ascii_unchecked(s.to_owned()) }.to_pyobject(vm) + unsafe { AsciiStr::from_ascii_unchecked(s) }.to_pyobject(vm) }) }, ), - PyStrKind::Utf8 => self.as_str().py_split( + PyKindStr::Utf8(s) => s.py_split( args, vm, + || zelf.as_object().to_owned(), + |v, s, vm| v.split(s).map(|s| vm.ctx.new_str(s).into()).collect(), + |v, s, n, vm| v.splitn(n, s).map(|s| vm.ctx.new_str(s).into()).collect(), + |v, n, vm| v.py_split_whitespace(n, |s| vm.ctx.new_str(s).into()), + ), + PyKindStr::Wtf8(w) => w.py_split( + args, + vm, + || zelf.as_object().to_owned(), |v, s, vm| v.split(s).map(|s| vm.ctx.new_str(s).into()).collect(), |v, s, n, vm| v.splitn(n, s).map(|s| vm.ctx.new_str(s).into()).collect(), |v, n, vm| v.py_split_whitespace(n, |s| vm.ctx.new_str(s).into()), @@ -578,10 +684,11 @@ impl PyStr { } #[pymethod] - fn rsplit(&self, args: SplitArgs, vm: &VirtualMachine) -> PyResult> { - let mut elements = self.as_str().py_split( + fn rsplit(zelf: &Py, args: SplitArgs, vm: &VirtualMachine) -> PyResult> { + let mut elements = zelf.as_wtf8().py_split( args, vm, + || zelf.as_object().to_owned(), |v, s, vm| v.rsplit(s).map(|s| vm.ctx.new_str(s).into()).collect(), |v, s, n, vm| v.rsplitn(n, s).map(|s| vm.ctx.new_str(s).into()).collect(), |v, n, vm| v.py_rsplit_whitespace(n, |s| vm.ctx.new_str(s).into()), @@ -593,14 +700,35 @@ impl PyStr { } #[pymethod] - fn strip(&self, chars: OptionalOption) -> String { - self.as_str() - .py_strip( - chars, - |s, chars| s.trim_matches(|c| chars.contains(c)), - |s| s.trim(), - ) - .to_owned() + fn strip(&self, chars: OptionalOption) -> PyStr { + match self.as_str_kind() { + PyKindStr::Ascii(s) => s + .py_strip( + chars, + |s, chars| { + let s = s + .as_str() + .trim_matches(|c| memchr::memchr(c as _, chars.as_bytes()).is_some()); + unsafe { AsciiStr::from_ascii_unchecked(s.as_bytes()) } + }, + |s| s.trim(), + ) + .into(), + PyKindStr::Utf8(s) => s + .py_strip( + chars, + |s, chars| s.trim_matches(|c| chars.contains(c)), + |s| s.trim(), + ) + .into(), + PyKindStr::Wtf8(w) => w + .py_strip( + chars, + |s, chars| s.trim_matches(|c| chars.code_points().contains(&c)), + |s| s.trim_matches(|c| c.is_char_and(char::is_whitespace)), + ) + .into(), + } } #[pymethod] @@ -701,12 +829,12 @@ impl PyStr { #[pymethod] fn isalnum(&self) -> bool { - !self.bytes.is_empty() && self.char_all(char::is_alphanumeric) + !self.data.is_empty() && self.char_all(char::is_alphanumeric) } #[pymethod] fn isnumeric(&self) -> bool { - !self.bytes.is_empty() && self.char_all(char::is_numeric) + !self.data.is_empty() && self.char_all(char::is_numeric) } #[pymethod] @@ -724,7 +852,7 @@ impl PyStr { #[pymethod] fn isdecimal(&self) -> bool { - !self.bytes.is_empty() + !self.data.is_empty() && self.char_all(|c| GeneralCategory::of(c) == GeneralCategory::DecimalNumber) } @@ -767,7 +895,9 @@ impl PyStr { } let s = FormatSpec::parse(spec) - .and_then(|format_spec| format_spec.format_string(zelf.borrow())) + .and_then(|format_spec| { + format_spec.format_string(&CharLenStr(zelf.as_str(), zelf.char_len())) + }) .map_err(|err| err.into_pyexception(vm))?; Ok(vm.ctx.new_str(s)) } @@ -776,7 +906,7 @@ impl PyStr { /// uppercase character and the remaining characters are lowercase. #[pymethod] fn title(&self) -> String { - let mut title = String::with_capacity(self.bytes.len()); + let mut title = String::with_capacity(self.data.len()); let mut previous_is_cased = false; for c in self.as_str().chars() { if c.is_lowercase() { @@ -803,7 +933,7 @@ impl PyStr { #[pymethod] fn swapcase(&self) -> String { - let mut swapped_str = String::with_capacity(self.bytes.len()); + let mut swapped_str = String::with_capacity(self.data.len()); for c in self.as_str().chars() { // to_uppercase returns an iterator, to_ascii_uppercase returns the char if c.is_lowercase() { @@ -819,7 +949,7 @@ impl PyStr { #[pymethod] fn isalpha(&self) -> bool { - !self.bytes.is_empty() && self.char_all(char::is_alphabetic) + !self.data.is_empty() && self.char_all(char::is_alphabetic) } #[pymethod] @@ -863,7 +993,7 @@ impl PyStr { #[pymethod] fn isspace(&self) -> bool { use unic_ucd_bidi::bidi_class::abbr_names::*; - !self.bytes.is_empty() + !self.data.is_empty() && self.char_all(|c| { GeneralCategory::of(c) == GeneralCategory::SpaceSeparator || matches!(BidiClass::of(c), WS | B | S) @@ -873,41 +1003,39 @@ impl PyStr { // Return true if all cased characters in the string are lowercase and there is at least one cased character, false otherwise. #[pymethod] fn islower(&self) -> bool { - match self.kind.kind() { - PyStrKind::Ascii => self.bytes.py_iscase(char::is_lowercase, char::is_uppercase), - PyStrKind::Utf8 => self - .as_str() - .py_iscase(char::is_lowercase, char::is_uppercase), + match self.as_str_kind() { + PyKindStr::Ascii(s) => s.py_islower(), + PyKindStr::Utf8(s) => s.py_islower(), + PyKindStr::Wtf8(w) => w.py_islower(), } } // Return true if all cased characters in the string are uppercase and there is at least one cased character, false otherwise. #[pymethod] fn isupper(&self) -> bool { - match self.kind.kind() { - PyStrKind::Ascii => self.bytes.py_iscase(char::is_uppercase, char::is_lowercase), - PyStrKind::Utf8 => self - .as_str() - .py_iscase(char::is_uppercase, char::is_lowercase), + match self.as_str_kind() { + PyKindStr::Ascii(s) => s.py_isupper(), + PyKindStr::Utf8(s) => s.py_isupper(), + PyKindStr::Wtf8(w) => w.py_isupper(), } } #[pymethod] fn splitlines(&self, args: anystr::SplitLinesArgs, vm: &VirtualMachine) -> Vec { - let into_wrapper = |s: &str| self.new_substr(s.to_owned()).to_pyobject(vm); + let into_wrapper = |s: &Wtf8| self.new_substr(s.to_owned()).to_pyobject(vm); let mut elements = Vec::new(); let mut last_i = 0; - let self_str = self.as_str(); - let mut enumerated = self_str.char_indices().peekable(); + let self_str = self.as_wtf8(); + let mut enumerated = self_str.code_point_indices().peekable(); while let Some((i, ch)) = enumerated.next() { - let end_len = match ch { + let end_len = match ch.to_char_lossy() { '\n' => 1, '\r' => { let is_rn = enumerated.next_if(|(_, ch)| *ch == '\n').is_some(); if is_rn { 2 } else { 1 } } '\x0b' | '\x0c' | '\x1c' | '\x1d' | '\x1e' | '\u{0085}' | '\u{2028}' - | '\u{2029}' => ch.len_utf8(), + | '\u{2029}' => ch.len_wtf8(), _ => continue, }; let range = if args.keepends { @@ -937,27 +1065,27 @@ impl PyStr { if first.as_object().class().is(vm.ctx.types.str_type) { return Ok(first); } else { - first.as_str().to_owned() + first.as_wtf8().to_owned() } } - Err(iter) => zelf.as_str().py_join(iter)?, + Err(iter) => zelf.as_wtf8().py_join(iter)?, }; Ok(vm.ctx.new_str(joined)) } // FIXME: two traversals of str is expensive #[inline] - fn _to_char_idx(r: &str, byte_idx: usize) -> usize { - r[..byte_idx].chars().count() + fn _to_char_idx(r: &Wtf8, byte_idx: usize) -> usize { + r[..byte_idx].code_points().count() } #[inline] fn _find(&self, args: FindArgs, find: F) -> Option where - F: Fn(&str, &str) -> Option, + F: Fn(&Wtf8, &Wtf8) -> Option, { let (sub, range) = args.get_value(self.len()); - self.as_str().py_find(sub.as_str(), range, find) + self.as_wtf8().py_find(sub.as_wtf8(), range, find) } #[pymethod] @@ -986,9 +1114,9 @@ impl PyStr { #[pymethod] fn partition(&self, sep: PyStrRef, vm: &VirtualMachine) -> PyResult { - let (front, has_mid, back) = self.as_str().py_partition( - sep.as_str(), - || self.as_str().splitn(2, sep.as_str()), + let (front, has_mid, back) = self.as_wtf8().py_partition( + sep.as_wtf8(), + || self.as_wtf8().splitn(2, sep.as_wtf8()), vm, )?; let partition = ( @@ -1005,9 +1133,9 @@ impl PyStr { #[pymethod] fn rpartition(&self, sep: PyStrRef, vm: &VirtualMachine) -> PyResult { - let (back, has_mid, front) = self.as_str().py_partition( - sep.as_str(), - || self.as_str().rsplitn(2, sep.as_str()), + let (back, has_mid, front) = self.as_wtf8().py_partition( + sep.as_wtf8(), + || self.as_wtf8().rsplitn(2, sep.as_wtf8()), vm, )?; Ok(( @@ -1026,7 +1154,7 @@ impl PyStr { /// empty, `false` otherwise. #[pymethod] fn istitle(&self) -> bool { - if self.bytes.is_empty() { + if self.data.is_empty() { return false; } @@ -1247,18 +1375,34 @@ impl PyStr { } } +struct CharLenStr<'a>(&'a str, usize); +impl std::ops::Deref for CharLenStr<'_> { + type Target = str; + fn deref(&self) -> &Self::Target { + self.0 + } +} +impl crate::common::format::CharLen for CharLenStr<'_> { + fn char_len(&self) -> usize { + self.1 + } +} + #[pyclass] impl PyRef { #[pymethod(magic)] fn str(self, vm: &VirtualMachine) -> PyRefExact { - self.into_exact_or(&vm.ctx, |zelf| unsafe { - // Creating a copy with same kind is safe - PyStr::new_str_unchecked(zelf.bytes.to_vec(), zelf.kind.kind()).into_exact_ref(&vm.ctx) + self.into_exact_or(&vm.ctx, |zelf| { + PyStr::from(zelf.data.clone()).into_exact_ref(&vm.ctx) }) } } impl PyStrRef { + pub fn is_empty(&self) -> bool { + (**self).is_empty() + } + pub fn concat_in_place(&mut self, other: &str, vm: &VirtualMachine) { // TODO: call [A]Rc::get_mut on the str to try to mutate the data in place if other.is_empty() { @@ -1296,7 +1440,7 @@ impl Comparable for PyStr { return Ok(res.into()); } let other = class_or_notimplemented!(Self, other); - Ok(op.eval_ord(zelf.as_str().cmp(other.as_str())).into()) + Ok(op.eval_ord(zelf.as_wtf8().cmp(other.as_wtf8())).into()) } } @@ -1352,8 +1496,7 @@ impl AsSequence for PyStr { }), item: atomic_func!(|seq, i, vm| { let zelf = PyStr::sequence_downcast(seq); - zelf.getitem_by_index(vm, i) - .map(|x| zelf.new_substr(x.to_string()).into_ref(&vm.ctx).into()) + zelf.getitem_by_index(vm, i).to_pyresult(vm) }), contains: atomic_func!( |seq, needle, vm| PyStr::sequence_downcast(seq)._contains(needle, vm) @@ -1396,9 +1539,21 @@ impl ToPyObject for String { } } +impl ToPyObject for Wtf8Buf { + fn to_pyobject(self, vm: &VirtualMachine) -> PyObjectRef { + vm.ctx.new_str(self).into() + } +} + impl ToPyObject for char { fn to_pyobject(self, vm: &VirtualMachine) -> PyObjectRef { - vm.ctx.new_str(self.to_string()).into() + vm.ctx.new_str(self).into() + } +} + +impl ToPyObject for CodePoint { + fn to_pyobject(self, vm: &VirtualMachine) -> PyObjectRef { + vm.ctx.new_str(self).into() } } @@ -1426,6 +1581,12 @@ impl ToPyObject for AsciiString { } } +impl ToPyObject for AsciiChar { + fn to_pyobject(self, vm: &VirtualMachine) -> PyObjectRef { + vm.ctx.new_str(self).into() + } +} + type SplitArgs = anystr::SplitArgs; #[derive(FromArgs)] @@ -1452,96 +1613,137 @@ pub fn init(ctx: &Context) { } impl SliceableSequenceOp for PyStr { - type Item = char; - type Sliced = String; + type Item = CodePoint; + type Sliced = PyStr; fn do_get(&self, index: usize) -> Self::Item { - if self.is_ascii() { - self.bytes[index] as char - } else { - self.as_str().chars().nth(index).unwrap() - } + self.data.nth_char(index) } fn do_slice(&self, range: Range) -> Self::Sliced { - let value = self.as_str(); - if self.is_ascii() { - value[range].to_owned() - } else { - rustpython_common::str::get_chars(value, range).to_owned() + match self.as_str_kind() { + PyKindStr::Ascii(s) => s[range].into(), + PyKindStr::Utf8(s) => { + let char_len = range.len(); + let out = rustpython_common::str::get_chars(s, range); + // SAFETY: char_len is accurate + unsafe { PyStr::new_with_char_len(out, char_len) } + } + PyKindStr::Wtf8(w) => { + let char_len = range.len(); + let out = rustpython_common::str::get_codepoints(w, range); + // SAFETY: char_len is accurate + unsafe { PyStr::new_with_char_len(out, char_len) } + } } } fn do_slice_reverse(&self, range: Range) -> Self::Sliced { - if self.is_ascii() { - // this is an ascii string - let mut v = self.bytes[range].to_vec(); - v.reverse(); - unsafe { - // SAFETY: an ascii string is always utf8 - String::from_utf8_unchecked(v) + match self.as_str_kind() { + PyKindStr::Ascii(s) => { + let mut out = s[range].to_owned(); + out.as_mut_slice().reverse(); + out.into() + } + PyKindStr::Utf8(s) => { + let char_len = range.len(); + let mut out = String::with_capacity(2 * char_len); + out.extend( + s.chars() + .rev() + .skip(self.char_len() - range.end) + .take(range.len()), + ); + // SAFETY: char_len is accurate + unsafe { PyStr::new_with_char_len(out, range.len()) } + } + PyKindStr::Wtf8(w) => { + let char_len = range.len(); + let mut out = Wtf8Buf::with_capacity(2 * char_len); + out.extend( + w.code_points() + .rev() + .skip(self.char_len() - range.end) + .take(range.len()), + ); + // SAFETY: char_len is accurate + unsafe { PyStr::new_with_char_len(out, char_len) } } - } else { - let mut s = String::with_capacity(self.bytes.len()); - s.extend( - self.as_str() - .chars() - .rev() - .skip(self.char_len() - range.end) - .take(range.end - range.start), - ); - s } } fn do_stepped_slice(&self, range: Range, step: usize) -> Self::Sliced { - if self.is_ascii() { - let v = self.bytes[range].iter().copied().step_by(step).collect(); - unsafe { - // SAFETY: Any subset of ascii string is a valid utf8 string - String::from_utf8_unchecked(v) + match self.as_str_kind() { + PyKindStr::Ascii(s) => s[range] + .as_slice() + .iter() + .copied() + .step_by(step) + .collect::() + .into(), + PyKindStr::Utf8(s) => { + let char_len = (range.len() / step) + 1; + let mut out = String::with_capacity(2 * char_len); + out.extend(s.chars().skip(range.start).take(range.len()).step_by(step)); + // SAFETY: char_len is accurate + unsafe { PyStr::new_with_char_len(out, char_len) } + } + PyKindStr::Wtf8(w) => { + let char_len = (range.len() / step) + 1; + let mut out = Wtf8Buf::with_capacity(2 * char_len); + out.extend( + w.code_points() + .skip(range.start) + .take(range.len()) + .step_by(step), + ); + // SAFETY: char_len is accurate + unsafe { PyStr::new_with_char_len(out, char_len) } } - } else { - let mut s = String::with_capacity(2 * ((range.len() / step) + 1)); - s.extend( - self.as_str() - .chars() - .skip(range.start) - .take(range.end - range.start) - .step_by(step), - ); - s } } fn do_stepped_slice_reverse(&self, range: Range, step: usize) -> Self::Sliced { - if self.is_ascii() { - // this is an ascii string - let v: Vec = self.bytes[range] - .iter() + match self.as_str_kind() { + PyKindStr::Ascii(s) => s[range] + .chars() .rev() - .copied() .step_by(step) - .collect(); - // TODO: from_utf8_unchecked? - String::from_utf8(v).unwrap() - } else { - // not ascii, so the codepoints have to be at least 2 bytes each - let mut s = String::with_capacity(2 * ((range.len() / step) + 1)); - s.extend( - self.as_str() - .chars() - .rev() - .skip(self.char_len() - range.end) - .take(range.end - range.start) - .step_by(step), - ); - s + .collect::() + .into(), + PyKindStr::Utf8(s) => { + let char_len = (range.len() / step) + 1; + // not ascii, so the codepoints have to be at least 2 bytes each + let mut out = String::with_capacity(2 * char_len); + out.extend( + s.chars() + .rev() + .skip(self.char_len() - range.end) + .take(range.len()) + .step_by(step), + ); + // SAFETY: char_len is accurate + unsafe { PyStr::new_with_char_len(out, char_len) } + } + PyKindStr::Wtf8(w) => { + let char_len = (range.len() / step) + 1; + // not ascii, so the codepoints have to be at least 2 bytes each + let mut out = Wtf8Buf::with_capacity(2 * char_len); + out.extend( + w.code_points() + .rev() + .skip(self.char_len() - range.end) + .take(range.len()) + .step_by(step), + ); + // SAFETY: char_len is accurate + unsafe { PyStr::new_with_char_len(out, char_len) } + } } } fn empty() -> Self::Sliced { - String::new() + PyStr::default() } fn len(&self) -> usize { @@ -1561,10 +1763,30 @@ impl AsRef for PyExact { } } -impl AnyStrWrapper for PyStrRef { - type Str = str; - fn as_ref(&self) -> &str { - self.as_str() +impl AnyStrWrapper for PyStrRef { + fn as_ref(&self) -> Option<&Wtf8> { + Some(self.as_wtf8()) + } + fn is_empty(&self) -> bool { + self.data.is_empty() + } +} + +impl AnyStrWrapper for PyStrRef { + fn as_ref(&self) -> Option<&str> { + self.data.as_str() + } + fn is_empty(&self) -> bool { + self.data.is_empty() + } +} + +impl AnyStrWrapper for PyStrRef { + fn as_ref(&self) -> Option<&AsciiStr> { + self.data.as_ascii() + } + fn is_empty(&self) -> bool { + self.data.is_empty() } } @@ -1582,14 +1804,22 @@ impl AnyStrContainer for String { } } +impl anystr::AnyChar for char { + fn is_lowercase(self) -> bool { + self.is_lowercase() + } + fn is_uppercase(self) -> bool { + self.is_uppercase() + } + fn bytes_len(self) -> usize { + self.len_utf8() + } +} + impl AnyStr for str { type Char = char; type Container = String; - fn element_bytes_len(c: char) -> usize { - c.len_utf8() - } - fn to_container(&self) -> Self::Container { self.to_owned() } @@ -1598,10 +1828,6 @@ impl AnyStr for str { self.as_bytes() } - fn chars(&self) -> impl Iterator { - str::chars(self) - } - fn elements(&self) -> impl Iterator { str::chars(self) } @@ -1675,6 +1901,232 @@ impl AnyStr for str { } } +impl AnyStrContainer for Wtf8Buf { + fn new() -> Self { + Wtf8Buf::new() + } + + fn with_capacity(capacity: usize) -> Self { + Wtf8Buf::with_capacity(capacity) + } + + fn push_str(&mut self, other: &Wtf8) { + self.push_wtf8(other) + } +} + +impl anystr::AnyChar for CodePoint { + fn is_lowercase(self) -> bool { + self.is_char_and(char::is_lowercase) + } + fn is_uppercase(self) -> bool { + self.is_char_and(char::is_uppercase) + } + fn bytes_len(self) -> usize { + self.len_wtf8() + } +} + +impl AnyStr for Wtf8 { + type Char = CodePoint; + type Container = Wtf8Buf; + + fn to_container(&self) -> Self::Container { + self.to_owned() + } + + fn as_bytes(&self) -> &[u8] { + self.as_bytes() + } + + fn elements(&self) -> impl Iterator { + self.code_points() + } + + fn get_bytes(&self, range: std::ops::Range) -> &Self { + &self[range] + } + + fn get_chars(&self, range: std::ops::Range) -> &Self { + rustpython_common::str::get_codepoints(self, range) + } + + fn bytes_len(&self) -> usize { + self.len() + } + + fn is_empty(&self) -> bool { + self.is_empty() + } + + fn py_split_whitespace(&self, maxsplit: isize, convert: F) -> Vec + where + F: Fn(&Self) -> PyObjectRef, + { + // CPython split_whitespace + let mut splits = Vec::new(); + let mut last_offset = 0; + let mut count = maxsplit; + for (offset, _) in self + .code_point_indices() + .filter(|(_, c)| c.is_char_and(|c| c.is_ascii_whitespace() || c == '\x0b')) + { + if last_offset == offset { + last_offset += 1; + continue; + } + if count == 0 { + break; + } + splits.push(convert(&self[last_offset..offset])); + last_offset = offset + 1; + count -= 1; + } + if last_offset != self.len() { + splits.push(convert(&self[last_offset..])); + } + splits + } + + fn py_rsplit_whitespace(&self, maxsplit: isize, convert: F) -> Vec + where + F: Fn(&Self) -> PyObjectRef, + { + // CPython rsplit_whitespace + let mut splits = Vec::new(); + let mut last_offset = self.len(); + let mut count = maxsplit; + for (offset, _) in self + .code_point_indices() + .rev() + .filter(|(_, c)| c.is_char_and(|c| c.is_ascii_whitespace() || c == '\x0b')) + { + if last_offset == offset + 1 { + last_offset -= 1; + continue; + } + if count == 0 { + break; + } + splits.push(convert(&self[offset + 1..last_offset])); + last_offset = offset; + count -= 1; + } + if last_offset != 0 { + splits.push(convert(&self[..last_offset])); + } + splits + } +} + +impl AnyStrContainer for AsciiString { + fn new() -> Self { + AsciiString::new() + } + + fn with_capacity(capacity: usize) -> Self { + AsciiString::with_capacity(capacity) + } + + fn push_str(&mut self, other: &AsciiStr) { + AsciiString::push_str(self, other) + } +} + +impl anystr::AnyChar for ascii::AsciiChar { + fn is_lowercase(self) -> bool { + self.is_lowercase() + } + fn is_uppercase(self) -> bool { + self.is_uppercase() + } + fn bytes_len(self) -> usize { + 1 + } +} + +const ASCII_WHITESPACES: [u8; 6] = [0x20, 0x09, 0x0a, 0x0c, 0x0d, 0x0b]; + +impl AnyStr for AsciiStr { + type Char = AsciiChar; + type Container = AsciiString; + + fn to_container(&self) -> Self::Container { + self.to_ascii_string() + } + + fn as_bytes(&self) -> &[u8] { + self.as_bytes() + } + + fn elements(&self) -> impl Iterator { + self.chars() + } + + fn get_bytes(&self, range: std::ops::Range) -> &Self { + &self[range] + } + + fn get_chars(&self, range: std::ops::Range) -> &Self { + &self[range] + } + + fn bytes_len(&self) -> usize { + self.len() + } + + fn is_empty(&self) -> bool { + self.is_empty() + } + + fn py_split_whitespace(&self, maxsplit: isize, convert: F) -> Vec + where + F: Fn(&Self) -> PyObjectRef, + { + let mut splits = Vec::new(); + let mut count = maxsplit; + let mut haystack = self; + while let Some(offset) = haystack.as_bytes().find_byteset(ASCII_WHITESPACES) { + if offset != 0 { + if count == 0 { + break; + } + splits.push(convert(&haystack[..offset])); + count -= 1; + } + haystack = &haystack[offset + 1..]; + } + if !haystack.is_empty() { + splits.push(convert(haystack)); + } + splits + } + + fn py_rsplit_whitespace(&self, maxsplit: isize, convert: F) -> Vec + where + F: Fn(&Self) -> PyObjectRef, + { + // CPython rsplit_whitespace + let mut splits = Vec::new(); + let mut count = maxsplit; + let mut haystack = self; + while let Some(offset) = haystack.as_bytes().rfind_byteset(ASCII_WHITESPACES) { + if offset + 1 != haystack.len() { + if count == 0 { + break; + } + splits.push(convert(&haystack[offset + 1..])); + count -= 1; + } + haystack = &haystack[..offset]; + } + if !haystack.is_empty() { + splits.push(convert(haystack)); + } + splits + } +} + /// The unique reference of interned PyStr /// Always intended to be used as a static reference pub type PyStrInterned = PyInterned; @@ -1688,7 +2140,7 @@ impl PyStrInterned { impl std::fmt::Display for PyStrInterned { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - std::fmt::Display::fmt(self.as_str(), f) + self.data.fmt(f) } } diff --git a/vm/src/builtins/union.rs b/vm/src/builtins/union.rs index 2798724ae..165113e21 100644 --- a/vm/src/builtins/union.rs +++ b/vm/src/builtins/union.rs @@ -204,7 +204,7 @@ impl PyUnion { vm, )?; let mut res; - if new_args.len() == 0 { + if new_args.is_empty() { res = make_union(&new_args, vm); } else { res = new_args.fast_getitem(0); diff --git a/vm/src/bytesinner.rs b/vm/src/bytesinner.rs index bcbbd4fce..63d5148e0 100644 --- a/vm/src/bytesinner.rs +++ b/vm/src/bytesinner.rs @@ -355,13 +355,11 @@ impl PyBytesInner { } pub fn islower(&self) -> bool { - self.elements - .py_iscase(char::is_lowercase, char::is_uppercase) + self.elements.py_islower() } pub fn isupper(&self) -> bool { - self.elements - .py_iscase(char::is_uppercase, char::is_lowercase) + self.elements.py_isupper() } pub fn isspace(&self) -> bool { @@ -654,6 +652,7 @@ impl PyBytesInner { let elements = self.elements.py_split( options, vm, + || convert(&self.elements, vm), |v, s, vm| v.split_str(s).map(|v| convert(v, vm)).collect(), |v, s, n, vm| v.splitn_str(n, s).map(|v| convert(v, vm)).collect(), |v, n, vm| v.py_split_whitespace(n, |v| convert(v, vm)), @@ -673,6 +672,7 @@ impl PyBytesInner { let mut elements = self.elements.py_split( options, vm, + || convert(&self.elements, vm), |v, s, vm| v.rsplit_str(s).map(|v| convert(v, vm)).collect(), |v, s, n, vm| v.rsplitn_str(n, s).map(|v| convert(v, vm)).collect(), |v, n, vm| v.py_rsplit_whitespace(n, |v| convert(v, vm)), @@ -998,10 +998,12 @@ pub trait ByteOr: ToPrimitive { impl ByteOr for BigInt {} -impl AnyStrWrapper for PyBytesInner { - type Str = [u8]; - fn as_ref(&self) -> &[u8] { - &self.elements +impl AnyStrWrapper<[u8]> for PyBytesInner { + fn as_ref(&self) -> Option<&[u8]> { + Some(&self.elements) + } + fn is_empty(&self) -> bool { + self.elements.is_empty() } } @@ -1021,14 +1023,22 @@ impl AnyStrContainer<[u8]> for Vec { const ASCII_WHITESPACES: [u8; 6] = [0x20, 0x09, 0x0a, 0x0c, 0x0d, 0x0b]; +impl anystr::AnyChar for u8 { + fn is_lowercase(self) -> bool { + self.is_ascii_lowercase() + } + fn is_uppercase(self) -> bool { + self.is_ascii_uppercase() + } + fn bytes_len(self) -> usize { + 1 + } +} + impl AnyStr for [u8] { type Char = u8; type Container = Vec; - fn element_bytes_len(_: u8) -> usize { - 1 - } - fn to_container(&self) -> Self::Container { self.to_vec() } @@ -1037,10 +1047,6 @@ impl AnyStr for [u8] { self } - fn chars(&self) -> impl Iterator { - bstr::ByteSlice::chars(self) - } - fn elements(&self) -> impl Iterator { self.iter().copied() } diff --git a/vm/src/codecs.rs b/vm/src/codecs.rs index e10409741..8a690fcb5 100644 --- a/vm/src/codecs.rs +++ b/vm/src/codecs.rs @@ -1,3 +1,5 @@ +use rustpython_common::wtf8::{CodePoint, Wtf8Buf}; + use crate::{ AsObject, Context, PyObject, PyObjectRef, PyPayload, PyResult, TryFromObject, VirtualMachine, builtins::{PyBaseExceptionRef, PyBytesRef, PyStr, PyStrRef, PyTuple, PyTupleRef}, @@ -570,10 +572,11 @@ fn surrogatepass_errors(err: PyObjectRef, vm: &VirtualMachine) -> PyResult<(PyOb return Err(err.downcast().unwrap()); } let s_after_start = - crate::common::str::try_get_chars(s.as_str(), range.start..).unwrap_or(""); + crate::common::str::try_get_codepoints(s.as_wtf8(), range.start..).unwrap_or_default(); let num_chars = range.len(); let mut out: Vec = Vec::with_capacity(num_chars * 4); - for c in s_after_start.chars().take(num_chars).map(|x| x as u32) { + for c in s_after_start.code_points().take(num_chars) { + let c = c.to_u32(); if !(0xd800..=0xdfff).contains(&c) { // Not a surrogate, fail with original exception return Err(err.downcast().unwrap()); @@ -671,7 +674,7 @@ fn surrogatepass_errors(err: PyObjectRef, vm: &VirtualMachine) -> PyResult<(PyOb } Ok(( - vm.new_pyobj(format!("\\x{c:x?}")), + vm.new_pyobj(CodePoint::from_u32(c).unwrap()), range.start + byte_length, )) } else { @@ -683,11 +686,11 @@ fn surrogateescape_errors(err: PyObjectRef, vm: &VirtualMachine) -> PyResult<(Py if err.fast_isinstance(vm.ctx.exceptions.unicode_encode_error) { let range = extract_unicode_error_range(&err, vm)?; let object = PyStrRef::try_from_object(vm, err.get_attr("object", vm)?)?; - let s_after_start = - crate::common::str::try_get_chars(object.as_str(), range.start..).unwrap_or(""); + let s_after_start = crate::common::str::try_get_codepoints(object.as_wtf8(), range.start..) + .unwrap_or_default(); let mut out: Vec = Vec::with_capacity(range.len()); - for ch in s_after_start.chars().take(range.len()) { - let ch = ch as u32; + for ch in s_after_start.code_points().take(range.len()) { + let ch = ch.to_u32(); if !(0xdc80..=0xdcff).contains(&ch) { // Not a UTF-8b surrogate, fail with original exception return Err(err.downcast().unwrap()); @@ -702,14 +705,14 @@ fn surrogateescape_errors(err: PyObjectRef, vm: &VirtualMachine) -> PyResult<(Py let object = PyBytesRef::try_from_object(vm, object)?; let p = &object.as_bytes()[range.clone()]; let mut consumed = 0; - let mut replace = String::with_capacity(4 * range.len()); + let mut replace = Wtf8Buf::with_capacity(4 * range.len()); while consumed < 4 && consumed < range.len() { - let c = p[consumed] as u32; + let c = p[consumed] as u16; // Refuse to escape ASCII bytes if c < 128 { break; } - write!(replace, "#{}", 0xdc00 + c).unwrap(); + replace.push(CodePoint::from(0xdc00 + c)); consumed += 1; } if consumed == 0 { diff --git a/vm/src/frame.rs b/vm/src/frame.rs index cf695cd87..c4dc23c95 100644 --- a/vm/src/frame.rs +++ b/vm/src/frame.rs @@ -21,6 +21,7 @@ use crate::{ }; use indexmap::IndexMap; use itertools::Itertools; +use rustpython_common::wtf8::Wtf8Buf; #[cfg(feature = "threading")] use std::sync::atomic; use std::{fmt, iter::zip}; @@ -697,8 +698,8 @@ impl ExecutingFrame<'_> { .pop_multiple(size.get(arg) as usize) .as_slice() .iter() - .map(|pyobj| pyobj.payload::().unwrap().as_ref()) - .collect::(); + .map(|pyobj| pyobj.payload::().unwrap()) + .collect::(); let str_obj = vm.ctx.new_str(s); self.push_value(str_obj.into()); Ok(None) @@ -1468,7 +1469,7 @@ impl ExecutingFrame<'_> { let kwarg_names = kwarg_names .as_slice() .iter() - .map(|pyobj| pyobj.payload::().unwrap().as_ref().to_owned()); + .map(|pyobj| pyobj.payload::().unwrap().as_str().to_owned()); FuncArgs::with_kwargs_names(args, kwarg_names) } diff --git a/vm/src/import.rs b/vm/src/import.rs index 860f0b8a1..2d86e47d0 100644 --- a/vm/src/import.rs +++ b/vm/src/import.rs @@ -204,15 +204,15 @@ fn remove_importlib_frames_inner( // TODO: This function should do nothing on verbose mode. // TODO: Fix this function after making PyTraceback.next mutable -pub fn remove_importlib_frames( - vm: &VirtualMachine, - exc: &PyBaseExceptionRef, -) -> PyBaseExceptionRef { +pub fn remove_importlib_frames(vm: &VirtualMachine, exc: &PyBaseExceptionRef) { + if vm.state.settings.verbose != 0 { + return; + } + let always_trim = exc.fast_isinstance(vm.ctx.exceptions.import_error); if let Some(tb) = exc.traceback() { let trimmed_tb = remove_importlib_frames_inner(vm, Some(tb), always_trim).0; exc.set_traceback(trimmed_tb); } - exc.clone() } diff --git a/vm/src/stdlib/builtins.rs b/vm/src/stdlib/builtins.rs index 9c2826a1e..6cae1bb7d 100644 --- a/vm/src/stdlib/builtins.rs +++ b/vm/src/stdlib/builtins.rs @@ -31,7 +31,9 @@ mod builtins { stdlib::sys, types::PyComparisonOp, }; + use itertools::Itertools; use num_traits::{Signed, ToPrimitive}; + use rustpython_common::wtf8::CodePoint; #[cfg(not(feature = "rustpython-compiler"))] const CODEGEN_NOT_SUPPORTED: &str = @@ -85,13 +87,13 @@ mod builtins { } #[pyfunction] - fn chr(i: PyIntRef, vm: &VirtualMachine) -> PyResult { + fn chr(i: PyIntRef, vm: &VirtualMachine) -> PyResult { let value = i .try_to_primitive::(vm)? .to_u32() - .and_then(char::from_u32) + .and_then(CodePoint::from_u32) .ok_or_else(|| vm.new_value_error("chr() arg not in range(0x110000)".to_owned()))?; - Ok(value.to_string()) + Ok(value) } #[derive(FromArgs)] @@ -618,21 +620,15 @@ mod builtins { } Ok(u32::from(bytes[0])) }), - Either::B(string) => { - let string = string.as_str(); - let string_len = string.chars().count(); - if string_len != 1 { - return Err(vm.new_type_error(format!( + Either::B(string) => match string.as_wtf8().code_points().exactly_one() { + Ok(character) => Ok(character.to_u32()), + Err(_) => { + let string_len = string.char_len(); + Err(vm.new_type_error(format!( "ord() expected a character, but string of length {string_len} found" - ))); + ))) } - match string.chars().next() { - Some(character) => Ok(character as u32), - None => Err(vm.new_type_error( - "ord() could not guess the integer representing this character".to_owned(), - )), - } - } + }, } } diff --git a/vm/src/stdlib/codecs.rs b/vm/src/stdlib/codecs.rs index 976545f64..9f58c8693 100644 --- a/vm/src/stdlib/codecs.rs +++ b/vm/src/stdlib/codecs.rs @@ -2,6 +2,8 @@ pub(crate) use _codecs::make_module; #[pymodule] mod _codecs { + use rustpython_common::wtf8::Wtf8Buf; + use crate::common::encodings; use crate::{ AsObject, PyObject, PyObjectRef, PyResult, TryFromBorrowedObject, VirtualMachine, @@ -257,7 +259,7 @@ mod _codecs { } } - type DecodeResult = PyResult<(String, usize)>; + type DecodeResult = PyResult<(Wtf8Buf, usize)>; #[derive(FromArgs)] struct DecodeArgs { diff --git a/vm/src/vm/interpreter.rs b/vm/src/vm/interpreter.rs index a375dbedc..cc669e066 100644 --- a/vm/src/vm/interpreter.rs +++ b/vm/src/vm/interpreter.rs @@ -163,7 +163,7 @@ mod tests { let b = vm.new_pyobj(4_i32); let res = vm._mul(&a, &b).unwrap(); let value = res.payload::().unwrap(); - assert_eq!(value.as_ref(), "Hello Hello Hello Hello ") + assert_eq!(value.as_str(), "Hello Hello Hello Hello ") }) } } diff --git a/vm/src/vm/mod.rs b/vm/src/vm/mod.rs index dd647e36f..eb4e846de 100644 --- a/vm/src/vm/mod.rs +++ b/vm/src/vm/mod.rs @@ -610,7 +610,7 @@ impl VirtualMachine { let from_list = from_list.to_pyobject(self); import_func .call((module.to_owned(), globals, locals, from_list, level), self) - .map_err(|exc| import::remove_importlib_frames(self, &exc)) + .inspect_err(|exc| import::remove_importlib_frames(self, exc)) } } }