diff --git a/library/core/src/array/mod.rs b/library/core/src/array/mod.rs index 32d344010aafd..9de241cd3ce55 100644 --- a/library/core/src/array/mod.rs +++ b/library/core/src/array/mod.rs @@ -454,7 +454,7 @@ where /// /// If `iter.next()` panicks, all items already yielded by the iterator are /// dropped. -fn collect_into_array(iter: &mut I) -> Option<[I::Item; N]> +pub(crate) fn collect_into_array(iter: &mut I) -> Option<[I::Item; N]> where I: Iterator, { diff --git a/library/core/src/iter/adapters/map_windows.rs b/library/core/src/iter/adapters/map_windows.rs new file mode 100644 index 0000000000000..862569d325d02 --- /dev/null +++ b/library/core/src/iter/adapters/map_windows.rs @@ -0,0 +1,159 @@ +use crate::{ + fmt, + mem::{self, MaybeUninit}, + ptr, +}; + +/// An iterator over the mapped windows of another iterator. +/// +/// This `struct` is created by the [`Iterator::map_windows`]. See its +/// documentation for more information. +#[must_use = "iterators are lazy and do nothing unless consumed"] +#[unstable(feature = "iter_map_windows", reason = "recently added", issue = "87155")] +pub struct MapWindows { + iter: I, + f: F, + + // The buffer is semantically `[MaybeUninit; 2 * N]`. However, due + // to limitations of const generics, we use this different type. Note that + // it has the same underlying memory layout. + // + // Invariant: if `buffer` is `Some`, `buffer[self.start..self.start + N]` is + // initialized, with all other elements being uninitialized. This also + // implies that `start <= N`. + buffer: Option<[[MaybeUninit; N]; 2]>, + start: usize, +} + +impl MapWindows { + pub(in crate::iter) fn new(mut iter: I, f: F) -> Self { + assert!(N > 0, "array in `Iterator::map_windows` must contain more than 0 elements"); + + let buffer = crate::array::collect_into_array(&mut iter).map(|first_half: [_; N]| { + // SAFETY: `MaybeUninit` is `repr(transparent)` and going from `T` to + // `MaybeUninit` is always safe. + let first_half = unsafe { + // FIXME(LukasKalbertodt): use `mem::transmute` once it works with arrays. + let copy: [MaybeUninit; N] = mem::transmute_copy(&first_half); + mem::forget(first_half); + copy + }; + [first_half, MaybeUninit::uninit_array()] + }); + + Self { iter, f, buffer, start: 0 } + } +} + +#[unstable(feature = "iter_map_windows", reason = "recently added", issue = "87155")] +impl Iterator for MapWindows +where + I: Iterator, + F: FnMut(&[I::Item; N]) -> R, +{ + type Item = R; + fn next(&mut self) -> Option { + let buffer_ptr = self.buffer.as_mut()?.as_mut_ptr().cast::>(); + + let out = { + debug_assert!(self.start + N <= 2 * N); + + // SAFETY: our invariant guarantees these elements are initialized. + let initialized_part = unsafe { + let ptr = buffer_ptr.add(self.start) as *const [I::Item; N]; + &*ptr + }; + (self.f)(initialized_part) + }; + + // Advance iterator. We first call `next` before changing our buffer at + // all. This means that if `next` panics, our invariant is upheld and + // our `Drop` impl drops the correct elements. + if let Some(next) = self.iter.next() { + if self.start == N { + // We have reached the end of our buffer and have to copy + // everything to the start. Example layout for N = 3. + // + // 0 1 2 3 4 5 0 1 2 3 4 5 + // ┌───┬───┬───┬───┬───┬───┐ ┌───┬───┬───┬───┬───┬───┐ + // │ - │ - │ - │ a │ b │ c │ -> │ b │ c │ n │ - │ - │ - │ + // └───┴───┴───┴───┴───┴───┘ └───┴───┴───┴───┴───┴───┘ + // ↑ ↑ + // start start + + // SAFETY: the two pointers are valid for reads/writes of N -1 + // elements because our array's size is semantically 2 * N. The + // regions also don't overlap for the same reason. + // + // We leave the old elements in place. As soon as `start` is set + // to 0, we treat them as uninitialized and treat their copies + // as initialized. + unsafe { + ptr::copy_nonoverlapping(buffer_ptr.add(N), buffer_ptr, N - 1); + (*buffer_ptr.add(N - 1)).write(next); + } + self.start = 0; + + // SAFETY: the index is valid and this is element `a` in the + // diagram above and has not been dropped yet. + unsafe { (*buffer_ptr.add(N)).assume_init_drop() }; + } else { + // SAFETY: `self.start` is < N as guaranteed by the invariant + // plus the check above. Even if the drop at the end panics, + // the invariant is upheld. + // + // Example layout for N = 3: + // + // 0 1 2 3 4 5 0 1 2 3 4 5 + // ┌───┬───┬───┬───┬───┬───┐ ┌───┬───┬───┬───┬───┬───┐ + // │ - │ a │ b │ c │ - │ - │ -> │ - │ - │ b │ c │ n │ - │ + // └───┴───┴───┴───┴───┴───┘ └───┴───┴───┴───┴───┴───┘ + // ↑ ↑ + // start start + // + unsafe { + (*buffer_ptr.add(self.start + N)).write(next); + self.start += 1; + (*buffer_ptr.add(self.start - 1)).assume_init_drop(); + } + } + } else { + // SAFETY: our invariant guarantees that N elements starting from + // `self.start` are initialized. We drop them here. + unsafe { + let initialized_part = crate::ptr::slice_from_raw_parts_mut( + buffer_ptr.add(self.start) as *mut I::Item, + N, + ); + crate::ptr::drop_in_place(initialized_part); + } + self.buffer = None; + } + + Some(out) + } +} + +#[unstable(feature = "iter_map_windows", reason = "recently added", issue = "87155")] +impl Drop for MapWindows { + fn drop(&mut self) { + if let Some(buffer) = self.buffer.as_mut() { + // SAFETY: our invariant guarantees that N elements starting from + // `self.start` are initialized. We drop them here. + unsafe { + let initialized_part = crate::ptr::slice_from_raw_parts_mut( + buffer.as_mut_ptr().cast::().add(self.start), + N, + ); + crate::ptr::drop_in_place(initialized_part); + } + } + } +} + +#[unstable(feature = "iter_map_windows", reason = "recently added", issue = "87155")] +impl fmt::Debug for MapWindows { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("MapWindows").field("iter", &self.iter).finish() + } +} diff --git a/library/core/src/iter/adapters/mod.rs b/library/core/src/iter/adapters/mod.rs index a3fbf4d9c38d8..13755e795f968 100644 --- a/library/core/src/iter/adapters/mod.rs +++ b/library/core/src/iter/adapters/mod.rs @@ -14,6 +14,7 @@ mod inspect; mod intersperse; mod map; mod map_while; +mod map_windows; mod peekable; mod rev; mod scan; @@ -48,6 +49,9 @@ pub use self::intersperse::{Intersperse, IntersperseWith}; #[unstable(feature = "iter_map_while", reason = "recently added", issue = "68537")] pub use self::map_while::MapWhile; +#[unstable(feature = "iter_map_windows", reason = "recently added", issue = "87155")] +pub use self::map_windows::MapWindows; + #[unstable(feature = "trusted_random_access", issue = "none")] pub use self::zip::TrustedRandomAccess; diff --git a/library/core/src/iter/mod.rs b/library/core/src/iter/mod.rs index bfb27da505eaf..143e69e0f092d 100644 --- a/library/core/src/iter/mod.rs +++ b/library/core/src/iter/mod.rs @@ -401,6 +401,8 @@ pub use self::adapters::Copied; pub use self::adapters::Flatten; #[unstable(feature = "iter_map_while", reason = "recently added", issue = "68537")] pub use self::adapters::MapWhile; +#[unstable(feature = "iter_map_windows", reason = "recently added", issue = "87155")] +pub use self::adapters::MapWindows; #[unstable(feature = "inplace_iteration", issue = "none")] pub use self::adapters::SourceIter; #[stable(feature = "iterator_step_by", since = "1.28.0")] diff --git a/library/core/src/iter/traits/iterator.rs b/library/core/src/iter/traits/iterator.rs index 8cb7aad28aa95..bf8aa5fd92d89 100644 --- a/library/core/src/iter/traits/iterator.rs +++ b/library/core/src/iter/traits/iterator.rs @@ -9,9 +9,8 @@ use super::super::TrustedRandomAccess; use super::super::{Chain, Cloned, Copied, Cycle, Enumerate, Filter, FilterMap, Fuse}; use super::super::{FlatMap, Flatten}; use super::super::{FromIterator, Intersperse, IntersperseWith, Product, Sum, Zip}; -use super::super::{ - Inspect, Map, MapWhile, Peekable, Rev, Scan, Skip, SkipWhile, StepBy, Take, TakeWhile, -}; +use super::super::{Inspect, Map, MapWhile, MapWindows, Peekable, Rev, Scan, Skip, SkipWhile}; +use super::super::{StepBy, Take, TakeWhile}; fn _assert_is_object_safe(_: &dyn Iterator) {} @@ -1449,6 +1448,90 @@ pub trait Iterator { Flatten::new(self) } + /// Calls the given function `f` for each contiguous window of size `N` over + /// `self` and returns an iterator over the outputs of `f`. + /// + /// In the following example, the closure is called three times with the + /// arguments `&['a', 'b']`, `&['b', 'c']` and `&['c', 'd']` respectively. + /// + /// ``` + /// #![feature(iter_map_windows)] + /// + /// let strings = "abcd".chars() + /// .map_windows(|[x, y]| format!("{}+{}", x, y)) + /// .collect::>(); + /// + /// assert_eq!(strings, vec!["a+b", "b+c", "c+d"]); + /// ``` + /// + /// Note that the const parameter `N` is usually inferred by the + /// destructured argument in the closure. + /// + /// The returned iterator yields 𝑘 − `N` + 1 items (where 𝑘 is the number of + /// items yielded by `self`). If `self` yields fewer than `N` items, the + /// iterator returned from this method is empty. + /// + /// + /// # Panics + /// + /// Panics if `N` is 0. This check will most probably get changed to a + /// compile time error before this method gets stabilized. + /// + /// + /// # Examples + /// + /// Building the sums of neighboring numbers. + /// + /// ``` + /// #![feature(iter_map_windows)] + /// + /// let mut it = [1, 3, 8, 1].iter().map_windows(|&[a, b]| a + b); + /// assert_eq!(it.next(), Some(4)); + /// assert_eq!(it.next(), Some(11)); + /// assert_eq!(it.next(), Some(9)); + /// assert_eq!(it.next(), None); + /// ``` + /// + /// Since the elements in the following example implement `Copy`, we can + /// just copy the array and get an iterator over the windows. + /// + /// ``` + /// #![feature(iter_map_windows)] + /// + /// let mut it = "ferris".chars().map_windows(|w: &[_; 3]| *w); + /// assert_eq!(it.next(), Some(['f', 'e', 'r'])); + /// assert_eq!(it.next(), Some(['e', 'r', 'r'])); + /// assert_eq!(it.next(), Some(['r', 'r', 'i'])); + /// assert_eq!(it.next(), Some(['r', 'i', 's'])); + /// assert_eq!(it.next(), None); + /// ``` + /// + /// You can also use this function to check the sortedness of an iterator. + /// For the simple case, rather use [`Iterator::is_sorted`]. + /// + /// ``` + /// #![feature(iter_map_windows)] + /// + /// let mut it = [0.5, 1.0, 3.5, 3.0, 8.5, 8.5, f32::NAN].iter() + /// .map_windows(|[a, b]| a <= b); + /// + /// assert_eq!(it.next(), Some(true)); // 0.5 <= 1.0 + /// assert_eq!(it.next(), Some(true)); // 1.0 <= 3.5 + /// assert_eq!(it.next(), Some(false)); // 3.5 <= 3.0 + /// assert_eq!(it.next(), Some(true)); // 3.0 <= 8.5 + /// assert_eq!(it.next(), Some(true)); // 8.5 <= 8.5 + /// assert_eq!(it.next(), Some(false)); // 8.5 <= NAN + /// assert_eq!(it.next(), None); + /// ``` + #[unstable(feature = "iter_map_windows", reason = "recently added", issue = "87155")] + fn map_windows(self, f: F) -> MapWindows + where + Self: Sized, + F: FnMut(&[Self::Item; N]) -> R, + { + MapWindows::new(self, f) + } + /// Creates an iterator which ends after the first [`None`]. /// /// After an iterator returns [`None`], future calls may or may not yield diff --git a/library/core/tests/iter/adapters/map_windows.rs b/library/core/tests/iter/adapters/map_windows.rs new file mode 100644 index 0000000000000..21e8d566ba1aa --- /dev/null +++ b/library/core/tests/iter/adapters/map_windows.rs @@ -0,0 +1,160 @@ +//! These tests mainly make sure the elements are correctly dropped. + +use std::sync::atomic::{AtomicBool, AtomicI64, Ordering::SeqCst}; + +#[derive(Debug)] +struct DropInfo { + dropped_twice: AtomicBool, + alive_count: AtomicI64, +} + +impl DropInfo { + const fn new() -> Self { + Self { dropped_twice: AtomicBool::new(false), alive_count: AtomicI64::new(0) } + } + + #[track_caller] + fn check(&self) { + assert!(!self.dropped_twice.load(SeqCst), "a value was dropped twice"); + assert_eq!(self.alive_count.load(SeqCst), 0); + } +} + +#[derive(Debug)] +struct DropCheck<'a> { + index: usize, + info: &'a DropInfo, + was_dropped: bool, +} + +impl<'a> DropCheck<'a> { + fn new(index: usize, info: &'a DropInfo) -> Self { + info.alive_count.fetch_add(1, SeqCst); + + Self { index, info, was_dropped: false } + } +} + +impl Drop for DropCheck<'_> { + fn drop(&mut self) { + if self.was_dropped { + self.info.dropped_twice.store(true, SeqCst); + } + self.was_dropped = true; + + self.info.alive_count.fetch_sub(1, SeqCst); + } +} + +fn iter(info: &DropInfo, len: usize, panic_at: usize) -> impl Iterator> { + (0..len).map(move |i| { + if i == panic_at { + panic!("intended panic"); + } + DropCheck::new(i, info) + }) +} + +#[track_caller] +fn check(len: usize, panic_at: usize) { + check_drops(|info| { + iter(info, len, panic_at).map_windows(|_: &[_; N]| {}).last(); + }); +} + +#[track_caller] +fn check_drops(f: impl FnOnce(&DropInfo)) { + let info = DropInfo::new(); + let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + f(&info); + })); + info.check(); +} + +#[test] +fn drop_check_no_iter_panic_n1() { + check::<1>(0, 100); + check::<1>(1, 100); + check::<1>(2, 100); + check::<1>(13, 100); +} + +#[test] +fn drop_check_no_iter_panic_n2() { + check::<2>(0, 100); + check::<2>(1, 100); + check::<2>(2, 100); + check::<2>(3, 100); + check::<2>(13, 100); +} + +#[test] +fn drop_check_no_iter_panic_n5() { + check::<5>(0, 100); + check::<5>(1, 100); + check::<5>(2, 100); + check::<5>(13, 100); + check::<5>(30, 100); +} + +#[test] +fn drop_check_panic_in_first_batch() { + check::<1>(7, 0); + + check::<2>(7, 0); + check::<2>(7, 1); + + check::<3>(7, 0); + check::<3>(7, 1); + check::<3>(7, 2); +} + +#[test] +fn drop_check_panic_in_middle() { + check::<1>(7, 1); + check::<1>(7, 5); + check::<1>(7, 6); + + check::<2>(7, 2); + check::<2>(7, 5); + check::<2>(7, 6); + + check::<5>(13, 5); + check::<5>(13, 8); + check::<5>(13, 12); +} + +#[test] +fn drop_check_len_equals_n() { + check::<1>(1, 100); + check::<1>(1, 0); + + check::<2>(2, 100); + check::<2>(2, 0); + check::<2>(2, 1); + + check::<5>(5, 100); + check::<5>(5, 0); + check::<5>(5, 1); + check::<5>(5, 4); +} + +#[test] +fn output_n1() { + assert_eq!("".chars().map_windows(|[c]| *c).collect::>(), vec![]); + assert_eq!("x".chars().map_windows(|[c]| *c).collect::>(), vec!['x']); + assert_eq!("abcd".chars().map_windows(|[c]| *c).collect::>(), vec!['a', 'b', 'c', 'd']); +} + +#[test] +fn output_n2() { + assert_eq!( + "".chars().map_windows(|a: &[_; 2]| *a).collect::>(), + >::new(), + ); + assert_eq!("ab".chars().map_windows(|a: &[_; 2]| *a).collect::>(), vec![['a', 'b']]); + assert_eq!( + "abcd".chars().map_windows(|a: &[_; 2]| *a).collect::>(), + vec![['a', 'b'], ['b', 'c'], ['c', 'd']], + ); +} diff --git a/library/core/tests/iter/adapters/mod.rs b/library/core/tests/iter/adapters/mod.rs index 96a53be1eaa50..9d719fc7d2443 100644 --- a/library/core/tests/iter/adapters/mod.rs +++ b/library/core/tests/iter/adapters/mod.rs @@ -11,6 +11,7 @@ mod fuse; mod inspect; mod intersperse; mod map; +mod map_windows; mod peekable; mod scan; mod skip; diff --git a/library/core/tests/lib.rs b/library/core/tests/lib.rs index cc4ff1be56314..f5bd6c867f4a4 100644 --- a/library/core/tests/lib.rs +++ b/library/core/tests/lib.rs @@ -28,6 +28,7 @@ #![feature(hashmap_internals)] #![feature(try_find)] #![feature(is_sorted)] +#![feature(iter_map_windows)] #![feature(pattern)] #![feature(sort_internals)] #![feature(slice_partition_at_index)]