Add peekable_map_while and peekable_take_while functions and associated structs

library/core/src/iter/adapters/mod.rs

@@ -15,6 +15,8 @@ mod intersperse;
 mod map;
 mod map_while;
 mod peekable;
+mod peekable_map_while;
+mod peekable_take_while;
 mod rev;
 mod scan;
 mod skip;
@@ -57,6 +59,12 @@ pub use self::zip::TrustedRandomAccessNoCoerce;
 #[stable(feature = "iter_zip", since = "1.59.0")]
 pub use self::zip::zip;
 
+#[unstable(feature = "peekable_map_while", issue = "none")]
+pub use self::peekable_map_while::PeekableMapWhile;
+
+#[unstable(feature = "peekable_map_while", issue = "none")]
+pub use self::peekable_take_while::PeekableTakeWhile;
+
 /// This trait provides transitive access to source-stage in an iterator-adapter pipeline
 /// under the conditions that
 /// * the iterator source `S` itself implements `SourceIter<Source = S>`

library/core/src/iter/adapters/peekable.rs

@@ -1,4 +1,7 @@
-use crate::iter::{adapters::SourceIter, FusedIterator, TrustedLen};
+use crate::iter::{
+    adapters::{PeekableMapWhile, PeekableTakeWhile, SourceIter},
+    FusedIterator, TrustedLen,
+};
 use crate::ops::{ControlFlow, Try};
 
 /// An iterator with a `peek()` that returns an optional reference to the next
@@ -315,6 +318,74 @@ impl<I: Iterator> Peekable<I> {
     {
         self.next_if(|next| next == expected)
     }
+
+    /// Creates an iterator which yields elements based on a predicate. This differs from
+    /// [`map_while`] in that this will not advance the iterator if the predicate fails.
+    ///
+    /// [`map_while`]: Iterator::map_while
+    ///
+    /// # Example
+    /// Consumes numbers until parsing them fails.
+    /// ```
+    /// let vec = vec!["0", "1", "2", "three", "four"];
+    /// let mut xs = vec.iter().peekable();
+    /// let ys: Vec<u8> = xs.peekable_map_while(|x| x.parse().ok()).collect();
+    ///
+    /// assert_eq!(ys, vec![0, 1, 2]);
+    /// assert_eq!(xs.next(), Some(&"three"));
+    /// assert_eq!(xs.next(), Some(&"four"));
+    /// assert_eq!(xs.next(), None);
+    /// ```
+    #[unstable(feature = "peekable_map_while", issue = "none")]
+    pub fn peekable_map_while<'iter, B, P>(
+        &'iter mut self,
+        predicate: P,
+    ) -> PeekableMapWhile<'iter, I, P>
+    where
+        Self: Sized,
+        P: FnMut(&I::Item) -> Option<B>,
+    {
+        PeekableMapWhile::new(self, predicate)
+    }
+
+    /// Creates an iterator which yields elements based on a predicate. This differs from
+    /// [`take_while`] in that `peekable_take_while()` will not advance the iterator if the
+    /// predicate returns false. Like [`take_while`], the closure passed to `peekable_take_while()`
+    /// takes a reference, and if the iterator iterates over references, it can lead to a possibly
+    /// confusing case where the type of the closure is a double reference. More information can
+    /// be found in the documentation for [`take_while`].
+    ///
+    /// [`take_while`]: Iterator::take_while
+    ///
+    /// # Example
+    /// Takes numbers less than 15.
+    /// ```
+    /// let mut xs = [0, 1, 2, 3, 5, 13, 15, 16, 17, 19].iter().peekable();
+    /// let ys = [0, 1, 2, 3, 5, 13];
+    /// let it = xs.peekable_take_while(|&x| *x < 15);
+    /// let mut i = 0;
+    /// for x in it {
+    ///     assert_eq!(*x, ys[i]);
+    ///     i += 1;
+    /// }
+    /// assert_eq!(i, ys.len());
+    /// assert_eq!(xs.next(), Some(&15));
+    /// assert_eq!(xs.next(), Some(&16));
+    /// assert_eq!(xs.next(), Some(&17));
+    /// assert_eq!(xs.next(), Some(&19));
+    /// assert_eq!(xs.next(), None);
+    /// ```
+    #[unstable(feature = "peekable_map_while", issue = "none")]
+    pub fn peekable_take_while<'iter, P>(
+        &'iter mut self,
+        predicate: P,
+    ) -> PeekableTakeWhile<'iter, I, P>
+    where
+        Self: Sized,
+        P: FnMut(&I::Item) -> bool,
+    {
+        PeekableTakeWhile::new(self, predicate)
+    }
 }
 
 #[unstable(feature = "trusted_len", issue = "37572")]

library/core/src/iter/adapters/peekable_map_while.rs

@@ -0,0 +1,62 @@
+use crate::fmt;
+use crate::iter::Peekable;
+
+/// An iterator that only accepts elements while `predicate` and `peek` returns `Some(_)`.
+///
+/// This `struct` is created by the [`peekable_map_while`] method on [`Peekable`]. See its
+/// documentation for more.
+///
+/// [`peekable_map_while`]: Peekable::peekable_map_while
+/// [`Peekable`]: Peekable
+#[unstable(feature = "peekable_map_while", issue = "none")]
+#[must_use = "iterators are lazy and do nothing unless consumed"]
+pub struct PeekableMapWhile<'iter, I: Iterator, P> {
+    iter: &'iter mut Peekable<I>,
+    predicate: P,
+}
+
+#[unstable(feature = "peekable_map_while", issue = "none")]
+impl<'iter, I: Iterator, P> PeekableMapWhile<'iter, I, P> {
+    pub(in crate::iter) fn new(
+        iter: &'iter mut Peekable<I>,
+        predicate: P,
+    ) -> PeekableMapWhile<'iter, I, P> {
+        PeekableMapWhile { iter, predicate }
+    }
+}
+
+#[unstable(feature = "peekable_map_while", issue = "none")]
+impl<'iter, I, P> fmt::Debug for PeekableMapWhile<'iter, I, P>
+where
+    I: fmt::Debug + Iterator,
+    <I as Iterator>::Item: fmt::Debug,
+{
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("PeekableMapWhile").field("iter", &self.iter).finish()
+    }
+}
+
+#[unstable(feature = "peekable_map_while", issue = "none")]
+impl<'iter, B, I: Iterator, P> Iterator for PeekableMapWhile<'iter, I, P>
+where
+    P: FnMut(&I::Item) -> Option<B>,
+{
+    type Item = B;
+
+    #[inline]
+    fn next(&mut self) -> Option<B> {
+        let x = self.iter.peek()?;
+        if let Some(b) = (self.predicate)(x) {
+            self.iter.next();
+            Some(b)
+        } else {
+            None
+        }
+    }
+
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let (_, upper) = self.iter.size_hint();
+        (0, upper) // can't know a lower bound, due to the predicate
+    }
+}

library/core/src/iter/adapters/peekable_take_while.rs

@@ -0,0 +1,81 @@
+use crate::fmt;
+use crate::iter::{FusedIterator, Peekable};
+
+/// An iterator that only accepts elements while `next` is `Some` `predicate` returns `true`.
+///
+/// This `struct` is created by the [`peekable_take_while`] method on [`Peekable`]. See its
+/// documentation for more.
+///
+/// [`peekable_take_while`]: Peekable::peekable_take_while
+/// [`Peekable`]: Peekable
+#[unstable(feature = "peekable_map_while", issue = "none")]
+#[must_use = "iterators are lazy and do nothing unless consumed"]
+pub struct PeekableTakeWhile<'iter, I: Iterator, P> {
+    iter: &'iter mut Peekable<I>,
+    flag: bool,
+    predicate: P,
+}
+
+impl<'iter, I: Iterator, P> PeekableTakeWhile<'iter, I, P> {
+    pub(in crate::iter) fn new(
+        iter: &'iter mut Peekable<I>,
+        predicate: P,
+    ) -> PeekableTakeWhile<'iter, I, P> {
+        PeekableTakeWhile { iter, flag: false, predicate }
+    }
+}
+
+#[unstable(feature = "peekable_map_while", issue = "none")]
+impl<'iter, I, P> fmt::Debug for PeekableTakeWhile<'iter, I, P>
+where
+    I: fmt::Debug + Iterator,
+    <I as Iterator>::Item: fmt::Debug,
+{
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("PeekableTakeWhile")
+            .field("iter", &self.iter)
+            .field("flag", &self.flag)
+            .finish()
+    }
+}
+
+#[unstable(feature = "peekable_map_while", issue = "none")]
+impl<'iter, I: Iterator, P> Iterator for PeekableTakeWhile<'iter, I, P>
+where
+    P: FnMut(&I::Item) -> bool,
+{
+    type Item = I::Item;
+
+    #[inline]
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.flag {
+            None
+        } else {
+            let x = self.iter.peek()?;
+            if (self.predicate)(x) {
+                Some(self.iter.next().unwrap())
+            } else {
+                self.flag = true;
+                None
+            }
+        }
+    }
+
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        if self.flag {
+            (0, Some(0))
+        } else {
+            let (_, upper) = self.iter.size_hint();
+            (0, upper) // can't know a lower bound, due to the predicate
+        }
+    }
+}
+
+#[unstable(feature = "peekable_map_while", issue = "none")]
+impl<'iter, I, P> FusedIterator for PeekableTakeWhile<'iter, I, P>
+where
+    I: FusedIterator,
+    P: FnMut(&I::Item) -> bool,
+{
+}

library/core/tests/iter/adapters/mod.rs

@@ -12,6 +12,8 @@ mod inspect;
 mod intersperse;
 mod map;
 mod peekable;
+mod peekable_map_while;
+mod peekable_take_while;
 mod scan;
 mod skip;
 mod skip_while;

library/core/tests/iter/adapters/peekable_map_while.rs

@@ -0,0 +1,46 @@
+use core::iter::*;
+
+#[test]
+fn test_iterator_peekable_map_while() {
+    let vec = vec!["0", "1", "2", "three", "four"];
+    let mut xs = vec.iter().peekable();
+    let ys: Vec<u8> = xs.peekable_map_while(|x| x.parse().ok()).collect();
+
+    assert_eq!(ys, vec![0, 1, 2]);
+    assert_eq!(xs.next(), Some(&"three"));
+    assert_eq!(xs.next(), Some(&"four"));
+    assert_eq!(xs.next(), None);
+
+    let vec = vec!["0", "1", "2", "3", "4"];
+    let mut xs = vec.iter().peekable();
+    let ys: Vec<u8> = xs.peekable_map_while(|x| x.parse().ok()).collect();
+
+    assert_eq!(ys, vec![0, 1, 2, 3, 4]);
+    assert_eq!(xs.next(), None);
+
+    let vec = vec!["zero", "one", "two", "three", "four"];
+    let mut xs = vec.iter().peekable();
+    let ys: Vec<u8> = xs.peekable_map_while(|x| x.parse().ok()).collect();
+
+    assert_eq!(ys, vec![]);
+    assert_eq!(xs.next(), Some(&"zero"));
+    assert_eq!(xs.next(), Some(&"one"));
+    assert_eq!(xs.next(), Some(&"two"));
+    assert_eq!(xs.next(), Some(&"three"));
+    assert_eq!(xs.next(), Some(&"four"));
+    assert_eq!(xs.next(), None);
+}
+
+#[test]
+fn test_iterator_peekable_map_while_fold() {
+    let mut xs = ["0", "1", "2", "3", "4", "5", "six"].iter().peekable();
+    let ys = [0, 5, 10, 15, 20, 25];
+    let it = xs.peekable_map_while(|x| x.parse::<usize>().ok().map(|x| x * 5));
+    let i = it.fold(0, |i, x| {
+        assert_eq!(x, ys[i]);
+        i + 1
+    });
+    assert_eq!(i, ys.len());
+    assert_eq!(xs.next(), Some(&"six"));
+    assert_eq!(xs.next(), None);
+}

library/core/tests/iter/adapters/peekable_take_while.rs

@@ -0,0 +1,42 @@
+use core::iter::*;
+
+#[test]
+fn test_iterator_peekable_take_while() {
+    let mut xs = [0, 1, 2, 3, 5, 13, 15, 16, 17, 19].iter().peekable();
+    let ys = [0, 1, 2, 3, 5, 13];
+    let it = xs.peekable_take_while(|&x| *x < 15);
+    let mut i = 0;
+    for x in it {
+        assert_eq!(*x, ys[i]);
+        i += 1;
+    }
+    assert_eq!(i, ys.len());
+    assert_eq!(xs.next(), Some(&15));
+    assert_eq!(xs.next(), Some(&16));
+    assert_eq!(xs.next(), Some(&17));
+    assert_eq!(xs.next(), Some(&19));
+    assert_eq!(xs.next(), None);
+}
+
+#[test]
+fn test_iterator_peekable_take_while_folds() {
+    let f = &|acc, x| i32::checked_add(2 * acc, x);
+    assert_eq!(
+        (1..20).peekable().peekable_take_while(|&x| x != 10).try_fold(7, f),
+        (1..10).try_fold(7, f)
+    );
+    let mut xs = (1..20).peekable();
+    let mut iter = xs.peekable_take_while(|&x| x != 10);
+    assert_eq!(iter.try_fold(0, |x, y| Some(x + y)), Some((1..10).sum()));
+    assert_eq!(iter.next(), None, "flag should be set");
+    assert_eq!(xs.next(), Some(10));
+    let mut xs = (1..20).peekable();
+    let iter = xs.peekable_take_while(|&x| x != 10);
+    assert_eq!(iter.fold(0, |x, y| x + y), (1..10).sum());
+    assert_eq!(xs.next(), Some(10));
+
+    let mut xs = (10..50).peekable();
+    let mut iter = xs.peekable_take_while(|&x| x != 40);
+    assert_eq!(iter.try_fold(0, i8::checked_add), None);
+    assert_eq!(iter.next(), Some(20));
+}

library/core/tests/lib.rs

@@ -71,6 +71,7 @@
 #![feature(never_type)]
 #![feature(unwrap_infallible)]
 #![feature(result_into_ok_or_err)]
+#![feature(peekable_map_while)]
 #![feature(portable_simd)]
 #![feature(ptr_metadata)]
 #![feature(once_cell)]