string_extract_if: initial implementation

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<char> = 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<F>(&mut self, filter: F) -> ExtractIf<'_, F>
+    where
+        F: FnMut(char) -> bool,
+    {
+        ExtractIf::new(self, filter)
+    }
 }
 
 impl FromUtf8Error {

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<u8>,
+
+    /// 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<F> Iterator for ExtractIf<'_, F>
+where
+    F: FnMut(char) -> bool,
+{
+    type Item = char;
+
+    fn next(&mut self) -> Option<char> {
+        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<usize>) {
+        (0, Some(self.old_len - self.valid_prefix.len() - self.bytes_removed))
+    }
+}
+
+#[unstable(feature = "string_extract_if", issue = "154318")]
+impl<F> 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<F> 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()
+    }
+}

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)]

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}你好！", .. }"#
+    );
+}