Rollup merge of #143725 - kennytm:peekable_next_if_map, r=jhpratt

core: add Peekable::next_if_map

Implementation for rust-lang/rust#143702

3 files changed, 189 insertions, 0 deletions

diff --git a/library/core/src/iter/adapters/peekable.rs b/library/core/src/iter/adapters/peekable.rs
index a6522659620..a55de75d56c 100644
--- a/library/core/src/iter/adapters/peekable.rs
+++ b/library/core/src/iter/adapters/peekable.rs
@@ -317,6 +317,108 @@ impl<I: Iterator> Peekable<I> {
     {
         self.next_if(|next| next == expected)
     }
+
+    /// Consumes the next value of this iterator and applies a function `f` on it,
+    /// returning the result if the closure returns `Ok`.
+    ///
+    /// Otherwise if the closure returns `Err` the value is put back for the next iteration.
+    ///
+    /// The content of the `Err` variant is typically the original value of the closure,
+    /// but this is not required. If a different value is returned,
+    /// the next `peek()` or `next()` call will result in this new value.
+    /// This is similar to modifying the output of `peek_mut()`.
+    ///
+    /// If the closure panics, the next value will always be consumed and dropped
+    /// even if the panic is caught, because the closure never returned an `Err` value to put back.
+    ///
+    /// # Examples
+    ///
+    /// Parse the leading decimal number from an iterator of characters.
+    /// ```
+    /// #![feature(peekable_next_if_map)]
+    /// let mut iter = "125 GOTO 10".chars().peekable();
+    /// let mut line_num = 0_u32;
+    /// while let Some(digit) = iter.next_if_map(|c| c.to_digit(10).ok_or(c)) {
+    ///     line_num = line_num * 10 + digit;
+    /// }
+    /// assert_eq!(line_num, 125);
+    /// assert_eq!(iter.collect::<String>(), " GOTO 10");
+    /// ```
+    ///
+    /// Matching custom types.
+    /// ```
+    /// #![feature(peekable_next_if_map)]
+    ///
+    /// #[derive(Debug, PartialEq, Eq)]
+    /// enum Node {
+    ///     Comment(String),
+    ///     Red(String),
+    ///     Green(String),
+    ///     Blue(String),
+    /// }
+    ///
+    /// /// Combines all consecutive `Comment` nodes into a single one.
+    /// fn combine_comments(nodes: Vec<Node>) -> Vec<Node> {
+    ///     let mut result = Vec::with_capacity(nodes.len());
+    ///     let mut iter = nodes.into_iter().peekable();
+    ///     let mut comment_text = None::<String>;
+    ///     loop {
+    ///         // Typically the closure in .next_if_map() matches on the input,
+    ///         //  extracts the desired pattern into an `Ok`,
+    ///         //  and puts the rest into an `Err`.
+    ///         while let Some(text) = iter.next_if_map(|node| match node {
+    ///             Node::Comment(text) => Ok(text),
+    ///             other => Err(other),
+    ///         }) {
+    ///             comment_text.get_or_insert_default().push_str(&text);
+    ///         }
+    ///
+    ///         if let Some(text) = comment_text.take() {
+    ///             result.push(Node::Comment(text));
+    ///         }
+    ///         if let Some(node) = iter.next() {
+    ///             result.push(node);
+    ///         } else {
+    ///             break;
+    ///         }
+    ///     }
+    ///     result
+    /// }
+    ///# assert_eq!( // hiding the test to avoid cluttering the documentation.
+    ///#     combine_comments(vec![
+    ///#         Node::Comment("The".to_owned()),
+    ///#         Node::Comment("Quick".to_owned()),
+    ///#         Node::Comment("Brown".to_owned()),
+    ///#         Node::Red("Fox".to_owned()),
+    ///#         Node::Green("Jumped".to_owned()),
+    ///#         Node::Comment("Over".to_owned()),
+    ///#         Node::Blue("The".to_owned()),
+    ///#         Node::Comment("Lazy".to_owned()),
+    ///#         Node::Comment("Dog".to_owned()),
+    ///#     ]),
+    ///#     vec![
+    ///#         Node::Comment("TheQuickBrown".to_owned()),
+    ///#         Node::Red("Fox".to_owned()),
+    ///#         Node::Green("Jumped".to_owned()),
+    ///#         Node::Comment("Over".to_owned()),
+    ///#         Node::Blue("The".to_owned()),
+    ///#         Node::Comment("LazyDog".to_owned()),
+    ///#     ],
+    ///# )
+    /// ```
+    #[unstable(feature = "peekable_next_if_map", issue = "143702")]
+    pub fn next_if_map<R>(&mut self, f: impl FnOnce(I::Item) -> Result<R, I::Item>) -> Option<R> {
+        let unpeek = if let Some(item) = self.next() {
+            match f(item) {
+                Ok(result) => return Some(result),
+                Err(item) => Some(item),
+            }
+        } else {
+            None
+        };
+        self.peeked = Some(unpeek);
+        None
+    }
 }
 
 #[unstable(feature = "trusted_len", issue = "37572")]
diff --git a/library/coretests/tests/iter/adapters/peekable.rs b/library/coretests/tests/iter/adapters/peekable.rs
index 7f4341b8902..f0549e8d6c2 100644
--- a/library/coretests/tests/iter/adapters/peekable.rs
+++ b/library/coretests/tests/iter/adapters/peekable.rs
@@ -271,3 +271,89 @@ fn test_peekable_non_fused() {
     assert_eq!(iter.peek(), None);
     assert_eq!(iter.next_back(), None);
 }
+
+#[test]
+fn test_peekable_next_if_map_mutation() {
+    fn collatz((mut num, mut len): (u64, u32)) -> Result<u32, (u64, u32)> {
+        let jump = num.trailing_zeros();
+        num >>= jump;
+        len += jump;
+        if num == 1 { Ok(len) } else { Err((3 * num + 1, len + 1)) }
+    }
+
+    let mut iter = once((3, 0)).peekable();
+    assert_eq!(iter.peek(), Some(&(3, 0)));
+    assert_eq!(iter.next_if_map(collatz), None);
+    assert_eq!(iter.peek(), Some(&(10, 1)));
+    assert_eq!(iter.next_if_map(collatz), None);
+    assert_eq!(iter.peek(), Some(&(16, 3)));
+    assert_eq!(iter.next_if_map(collatz), Some(7));
+    assert_eq!(iter.peek(), None);
+    assert_eq!(iter.next_if_map(collatz), None);
+}
+
+#[test]
+#[cfg_attr(not(panic = "unwind"), ignore = "test requires unwinding support")]
+fn test_peekable_next_if_map_panic() {
+    use core::cell::Cell;
+    use std::panic::{AssertUnwindSafe, catch_unwind};
+
+    struct BitsetOnDrop<'a> {
+        value: u32,
+        cell: &'a Cell<u32>,
+    }
+    impl<'a> Drop for BitsetOnDrop<'a> {
+        fn drop(&mut self) {
+            self.cell.update(|v| v | self.value);
+        }
+    }
+
+    let cell = &Cell::new(0);
+    let mut it = [
+        BitsetOnDrop { value: 1, cell },
+        BitsetOnDrop { value: 2, cell },
+        BitsetOnDrop { value: 4, cell },
+        BitsetOnDrop { value: 8, cell },
+    ]
+    .into_iter()
+    .peekable();
+
+    // sanity check, .peek() won't consume the value, .next() will transfer ownership.
+    let item = it.peek().unwrap();
+    assert_eq!(item.value, 1);
+    assert_eq!(cell.get(), 0);
+    let item = it.next().unwrap();
+    assert_eq!(item.value, 1);
+    assert_eq!(cell.get(), 0);
+    drop(item);
+    assert_eq!(cell.get(), 1);
+
+    // next_if_map returning Ok should transfer the value out.
+    let item = it.next_if_map(Ok).unwrap();
+    assert_eq!(item.value, 2);
+    assert_eq!(cell.get(), 1);
+    drop(item);
+    assert_eq!(cell.get(), 3);
+
+    // next_if_map returning Err should not drop anything.
+    assert_eq!(it.next_if_map::<()>(Err), None);
+    assert_eq!(cell.get(), 3);
+    assert_eq!(it.peek().unwrap().value, 4);
+    assert_eq!(cell.get(), 3);
+
+    // next_if_map panicking should consume and drop the item.
+    let result = catch_unwind({
+        let mut it = AssertUnwindSafe(&mut it);
+        move || it.next_if_map::<()>(|_| panic!())
+    });
+    assert!(result.is_err());
+    assert_eq!(cell.get(), 7);
+    assert_eq!(it.next().unwrap().value, 8);
+    assert_eq!(cell.get(), 15);
+    assert!(it.peek().is_none());
+
+    // next_if_map should *not* execute the closure if the iterator is exhausted.
+    assert!(it.next_if_map::<()>(|_| panic!()).is_none());
+    assert!(it.peek().is_none());
+    assert_eq!(cell.get(), 15);
+}
diff --git a/library/coretests/tests/lib.rs b/library/coretests/tests/lib.rs
index ffb96b5a7cf..642d28759ae 100644
--- a/library/coretests/tests/lib.rs
+++ b/library/coretests/tests/lib.rs
@@ -82,6 +82,7 @@
 #![feature(numfmt)]
 #![feature(option_reduce)]
 #![feature(pattern)]
+#![feature(peekable_next_if_map)]
 #![feature(pointer_is_aligned_to)]
 #![feature(portable_simd)]
 #![feature(ptr_metadata)]