From 2261154549996eb7737a895e5df5fa894dc03e42 Mon Sep 17 00:00:00 2001 From: kennytm Date: Thu, 10 Jul 2025 18:22:18 +0800 Subject: [PATCH] core: add Peekable::next_if_map --- library/core/src/iter/adapters/peekable.rs | 102 ++++++++++++++++++ .../coretests/tests/iter/adapters/peekable.rs | 86 +++++++++++++++ library/coretests/tests/lib.rs | 1 + 3 files changed, 189 insertions(+) diff --git a/library/core/src/iter/adapters/peekable.rs b/library/core/src/iter/adapters/peekable.rs index a6522659620a0..a55de75d56c6e 100644 --- a/library/core/src/iter/adapters/peekable.rs +++ b/library/core/src/iter/adapters/peekable.rs @@ -317,6 +317,108 @@ impl Peekable { { 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::(), " 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) -> Vec { + /// let mut result = Vec::with_capacity(nodes.len()); + /// let mut iter = nodes.into_iter().peekable(); + /// let mut comment_text = None::; + /// 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(&mut self, f: impl FnOnce(I::Item) -> Result) -> Option { + 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 7f4341b8902c8..f0549e8d6c2c3 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 { + 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, + } + 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 fdef736c0c0f7..18e6e96c235f0 100644 --- a/library/coretests/tests/lib.rs +++ b/library/coretests/tests/lib.rs @@ -78,6 +78,7 @@ #![feature(next_index)] #![feature(numfmt)] #![feature(pattern)] +#![feature(peekable_next_if_map)] #![feature(pointer_is_aligned_to)] #![feature(portable_simd)] #![feature(ptr_metadata)]