diff --git a/library/alloc/src/string.rs b/library/alloc/src/string.rs index 30e52f3e1be46..963c5ab85d48b 100644 --- a/library/alloc/src/string.rs +++ b/library/alloc/src/string.rs @@ -67,6 +67,11 @@ use crate::str::{self, CharIndices, Chars, Utf8Error, from_utf8_unchecked_mut}; use crate::str::{FromStr, from_boxed_utf8_unchecked}; use crate::vec::{self, Vec}; +mod extract_if; + +#[unstable(feature = "string_extract_if", issue = "154318")] +pub use self::extract_if::ExtractIf; + /// A UTF-8–encoded, growable string. /// /// `String` is the most common string type. It has ownership over the contents @@ -2206,6 +2211,69 @@ impl String { let slice = self.vec.leak(); unsafe { from_utf8_unchecked_mut(slice) } } + + /// Creates an iterator which uses a closure to determine if a character should be removed. + /// + /// If the closure returns `true`, the character is removed from the string + /// and yielded. If the closure returns `false`, or panics, the character + /// remains in the string and will not be yielded. + /// + /// If the returned `ExtractIf` is not exhausted, e.g. because it is dropped without iterating + /// or the iteration short-circuits, then the remaining characters will be retained. + /// If you don't need the returned iterator, use `extract_if().for_each(drop)` or + /// [`retain`] with a negated predicate. + /// + /// [`retain`]: String::retain + /// + /// Using this method is equivalent to the following code: + /// + /// ``` + /// # #![feature(string_extract_if)] + /// # let some_predicate = |c: char| c.is_whitespace(); + /// # let mut s = "a\nb\u{3000}\td🦀".to_string(); + /// # let mut s2 = s.clone(); + /// let mut i = 0; + /// # let mut extracted = String::new(); + /// + /// while let Some(c) = s[i..].chars().next() { + /// if some_predicate(c) { + /// s.remove(i); + /// // your code here + /// # extracted.push(c); + /// } else { + /// i += c.len_utf8(); + /// } + /// } + /// + /// # let extracted2: String = s2.extract_if(some_predicate).collect(); + /// # assert_eq!(s, s2); + /// # assert_eq!(extracted, extracted2); + /// ``` + /// + /// But `extract_if` is easier to use. `extract_if` is also more efficient, + /// because it can backshift the characters of the string in bulk + /// and does not need to perform additional UTF-8 validity checks. + /// + /// # Examples + /// + /// Extracting control characters from a string into a separate buffer, reusing the original string: + /// + /// ``` + /// #![feature(string_extract_if)] + /// let mut string = "Hello,\u{0008} World!\r\n".to_string(); + /// + /// let extracted: Vec = string.extract_if(|c| c.is_control()).collect(); + /// + /// assert_eq!(string, "Hello, World!"); + /// assert_eq!(extracted, vec!['\u{0008}', '\r', '\n']); + /// ``` + #[unstable(feature = "string_extract_if", issue = "154318")] + pub fn extract_if(&mut self, filter: F) -> ExtractIf<'_, F> + where + F: FnMut(char) -> bool, + { + ExtractIf::new(self, filter) + } } impl FromUtf8Error { diff --git a/library/alloc/src/string/extract_if.rs b/library/alloc/src/string/extract_if.rs new file mode 100644 index 0000000000000..e5403e58038c6 --- /dev/null +++ b/library/alloc/src/string/extract_if.rs @@ -0,0 +1,159 @@ +use core::{fmt, ptr, slice}; + +use super::{String, Vec}; + +/// An iterator which uses a closure to determine if a character should be removed. +/// +/// This struct is created by [`String::extract_if`]. +/// See its documentation for more. +/// +/// # Example +/// +/// ``` +/// #![feature(string_extract_if)] +/// let mut s = "Hello! Привет! 你好!".to_string(); +/// let iter: std::string::ExtractIf<'_, _> = s.extract_if(|c| c.len_utf8() > 2); +/// ``` +#[unstable(feature = "string_extract_if", issue = "154318")] +#[must_use = "iterators are lazy and do nothing unless consumed; \ + use `retain` or `extract_if().for_each(drop)` to remove and discard characters"] +pub struct ExtractIf<'a, F> { + /// The underlying vector of the original [`String`]. We set its length to zero in [`ExtractIf::new`] + /// (to prevent invalid UTF-8 from being exposed if this struct is leaked before the iteration is complete) + /// and then gradually increase the `len` as this vector's prefix keeps filling with valid UTF-8. + /// The `len` is finally adjusted in [`ExtractIf::drop`]. + valid_prefix: &'a mut Vec, + + /// During the iteration, the underlying vector's consists of: + /// - A valid UTF-8 prefix (`valid_prefix.len()` bytes) + /// of characters that we iterated over and didn't extract. + /// - A middle portion of `bytes_removed` initialized bytes that might not be valid UTF-8. + /// - A valid UTF-8 suffix (`old_len - bytes_removed - valid_prefix.len()` bytes) + /// of characters that we have not iterated over yet. + /// - Potentially some spare capacity (`valid_prefix.capacity() - old_len`). We never touch it. + /// + /// The above (together with the fact that `valid_prefix.len() + bytes_removed <= old_len`, + /// that `old_len` is never changed, and that we never reallocate the vector) + /// is essentially this structure's invariant. + bytes_removed: usize, + + /// The (byte) length of the original [`String`] prior to draining. + /// Sadly, we need to hold on to a copy of it, + /// because we temporarily lower the [`String`]'s length during iteration + /// to prevent invalid UTF-8 from showing up. + /// This field is never changed. + old_len: usize, + + /// The filter test predicate. + pred: F, +} + +impl<'a, F> ExtractIf<'a, F> { + pub(super) fn new(string: &'a mut String, pred: F) -> Self { + let valid_prefix = &mut string.vec; + let old_len = valid_prefix.len(); + unsafe { valid_prefix.set_len(0) }; + ExtractIf { valid_prefix, bytes_removed: 0, old_len, pred } + } +} + +#[unstable(feature = "string_extract_if", issue = "154318")] +impl Iterator for ExtractIf<'_, F> +where + F: FnMut(char) -> bool, +{ + type Item = char; + + fn next(&mut self) -> Option { + loop { + // `valid_prefix` actually has `old_len` initialized bytes of memory in it, + // but sadly we can't access them using `get_unchecked` + // because it prohibits accesses outside of `0..valid_prefix.len()` range. + // So we have to do some pointer aritmetics here instead. + + // Have to hold on to a copy of this to not materialize any `&self.valid_prefx` + // while still using the pointer (and its derivatives) we got from `.as_mut_ptr()`. + let valid_prefix_len = self.valid_prefix.len(); + + let tail = { + let bytes = self.valid_prefix.as_mut_ptr(); + // SAFETY: all of these bytes were initialized by the original string. + let bytes = unsafe { slice::from_raw_parts_mut(bytes, self.old_len) }; + // SAFETY: by our invariant, `valid_prefix.len() <= old_len`. + unsafe { bytes.get_unchecked_mut(valid_prefix_len..) } + }; + + let c = unsafe { + // SAFETY: by our invariant, `bytes_removed <= old_len - valid_prefix.len()`. + let valid_tail = tail.get_unchecked(self.bytes_removed..); + // SAFETY: we have not touched these bytes before, so they remain valid UTF-8. + str::from_utf8_unchecked(valid_tail) + } + .chars() // FIXME(str_first_last_char): replace this with `first_char` + .next()?; + + let char_len = c.len_utf8(); + if (self.pred)(c) { + self.bytes_removed += char_len; + return Some(c); + } else { + let tail = tail.as_mut_ptr(); + unsafe { + // SAFETY: we just had a `&mut [u8]` covering both `src` and `dst`, so they are valid. + // `src` had a prefix of UTF-8-encoded `char` `c`, whose `.len_utf8()` is `char_len`, + // and `dst` precedes `src` in the slice mentioned above, so all of this is still in bounds. + ptr::copy(tail.add(self.bytes_removed), tail, char_len); + + // SAFETY: we just appended `char_len` more bytes of valid UTF-8 + // to the existing `valid_prefix_len` bytes of valid UTF-8. + self.valid_prefix.set_len(valid_prefix_len + char_len); + } + } + } + } + + fn size_hint(&self) -> (usize, Option) { + (0, Some(self.old_len - self.valid_prefix.len() - self.bytes_removed)) + } +} + +#[unstable(feature = "string_extract_if", issue = "154318")] +impl Drop for ExtractIf<'_, F> { + fn drop(&mut self) { + let hole_size = self.bytes_removed; + let valid_prefix_len = self.valid_prefix.len(); + let valid_tail_len = self.old_len - valid_prefix_len - hole_size; + if valid_tail_len > 0 { + unsafe { + let tail = self.valid_prefix.as_mut_ptr().add(valid_prefix_len); + ptr::copy(tail.add(hole_size), tail, valid_tail_len) + } + // SAFETY: we had a prefix of `valid_prefix_len` bytes of valid UTF-8, + // and a suffix of `valid_tail_len` bytes of valid UTF-8. + // We just `prt::copy`'ed the suffix to come right after the prefix, so it's safe to bump the length. + unsafe { self.valid_prefix.set_len(valid_prefix_len + valid_tail_len) } + } + } +} + +#[unstable(feature = "string_extract_if", issue = "154318")] +impl fmt::Debug for ExtractIf<'_, F> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let valid_prefix_len = self.valid_prefix.len(); + let bytes = self.valid_prefix.as_ptr(); + // SAFETY: all of this memory was initialized by the original string. + let bytes = unsafe { slice::from_raw_parts(bytes, self.old_len) }; + // SAFETY: by invariant, `valid_prefix_len <= old_len`. + let (valid_prefix, tail) = unsafe { bytes.split_at_unchecked(valid_prefix_len) }; + // SAFETY: by our invariant, `bytes_removed <= old_len - valid_prefix.len()`. + let (_hole, valid_suffix) = unsafe { tail.split_at_unchecked(self.bytes_removed) }; + // SAFETY: by our invariant, this prefix is valid UTF-8. + let valid_prefix = unsafe { str::from_utf8_unchecked(valid_prefix) }; + // SAFETY: by our invariant, this suffix is valid UTF-8. + let valid_suffix = unsafe { str::from_utf8_unchecked(valid_suffix) }; + f.debug_struct("ExtractIf") + .field("retained", &valid_prefix) + .field("remainder", &valid_suffix) + .finish_non_exhaustive() + } +} diff --git a/library/alloctests/tests/lib.rs b/library/alloctests/tests/lib.rs index 699a5010282b0..dbdc75fa1cef1 100644 --- a/library/alloctests/tests/lib.rs +++ b/library/alloctests/tests/lib.rs @@ -25,6 +25,7 @@ #![feature(iter_advance_by)] #![feature(iter_next_chunk)] #![feature(slice_partition_dedup)] +#![feature(string_extract_if)] #![feature(string_from_utf8_lossy_owned)] #![feature(string_remove_matches)] #![feature(const_btree_len)] diff --git a/library/alloctests/tests/string.rs b/library/alloctests/tests/string.rs index 08eb1855a4824..56df065d28e15 100644 --- a/library/alloctests/tests/string.rs +++ b/library/alloctests/tests/string.rs @@ -3,7 +3,7 @@ use std::cell::Cell; use std::collections::TryReserveErrorKind::*; use std::ops::Bound::*; use std::ops::{Bound, RangeBounds}; -use std::{assert_matches, panic, str}; +use std::{assert_matches, mem, panic, str}; pub trait IntoCow<'a, B: ?Sized> where @@ -956,3 +956,123 @@ fn test_str_concat() { let s: String = format!("{a}{b}"); assert_eq!(s.as_bytes()[9], 'd' as u8); } + +#[test] +fn extract_if_empty() { + let mut s = String::new(); + { + let mut iter = s.extract_if(|_| true); + assert_eq!(iter.size_hint(), (0, Some(0))); + assert_eq!(iter.next(), None); + assert_eq!(iter.size_hint(), (0, Some(0))); + assert_eq!(iter.next(), None); + assert_eq!(iter.size_hint(), (0, Some(0))); + } + assert_eq!(s.len(), 0); + assert_eq!(s, ""); +} + +#[test] +fn extract_if_false() { + let initial = "1234567890一二三四五六七八九十"; + let mut s = initial.to_string(); + + let initial_len = s.len(); + let mut count = 0; + { + let mut iter = s.extract_if(|_| false); + assert_eq!(iter.size_hint(), (0, Some(initial_len))); + for _ in iter.by_ref() { + count += 1; + } + assert_eq!(iter.size_hint(), (0, Some(0))); + assert_eq!(iter.next(), None); + assert_eq!(iter.size_hint(), (0, Some(0))); + } + + assert_eq!(count, 0); + assert_eq!(s.len(), initial_len); + assert_eq!(s, initial); +} + +#[test] +fn extract_if_true() { + let initial = "1234567890一二三四五六七八九十"; + + let mut s = initial.to_string(); + let extracted: String = s.extract_if(|_| true).collect(); + assert_eq!(extracted, initial); + assert_eq!(s, ""); + + let mut s = initial.to_string(); + let initial_len = s.len(); + let mut char_count = 0; + let mut byte_count = 0; + { + let mut iter = s.extract_if(|_| true); + assert_eq!(iter.size_hint(), (0, Some(initial_len))); + while let Some(c) = iter.next() { + char_count += 1; + byte_count += c.len_utf8(); + assert_eq!(iter.size_hint(), (0, Some(initial_len - byte_count))); + } + assert_eq!(iter.size_hint(), (0, Some(0))); + assert_eq!(iter.next(), None); + assert_eq!(iter.size_hint(), (0, Some(0))); + } + + assert_eq!(char_count, initial.chars().count()); + assert_eq!(byte_count, initial.len()); + assert_eq!(s.len(), 0); + assert_eq!(s, ""); +} + +#[test] +fn extract_if_unconsumed() { + let initial = "Hello, world!"; + let mut s = initial.to_string(); + + let drain = s.extract_if(|_| true); + drop(drain); + assert_eq!(s, initial); + + let mut i: usize = 0; + let mut drain = s.extract_if(|_| { + i += 1; + i % 2 == 0 + }); + for c in "el,w".chars() { + assert_eq!(drain.next(), Some(c)) + } + drop(drain); + assert_eq!(s, "Hlo orld!"); + + let mut s = initial.to_string(); + let mut i: usize = 0; + let mut drain = s.extract_if(|_| { + i += 1; + i % 2 == 0 + }); + for c in "el,w".chars() { + assert_eq!(drain.next(), Some(c)) + } + mem::forget(drain); + assert_eq!(s, "Hlo "); +} + +#[test] +fn extract_if_debug() { + let mut s = "Hello! Привет! 你好!".to_string(); + let mut drain = s.extract_if(|c| c.is_alphabetic()); + assert_eq!( + format!("{drain:?}"), + r#"ExtractIf { retained: "", remainder: "Hello! Привет!\u{3000}你好!", .. }"# + ); + for c in "HelloПривет".chars() { + assert_eq!(drain.next(), Some(c)) + } + assert_eq!( + format!("{drain:?}"), + r#"ExtractIf { retained: "! ", remainder: "!\u{3000}你好!", .. }"# + ); +}