diff options
Diffstat (limited to 'library/core/src/iter/adapters')
24 files changed, 5426 insertions, 0 deletions
diff --git a/library/core/src/iter/adapters/by_ref_sized.rs b/library/core/src/iter/adapters/by_ref_sized.rs new file mode 100644 index 000000000..cc1e8e8a2 --- /dev/null +++ b/library/core/src/iter/adapters/by_ref_sized.rs @@ -0,0 +1,86 @@ +use crate::ops::Try; + +/// Like `Iterator::by_ref`, but requiring `Sized` so it can forward generics. +/// +/// Ideally this will no longer be required, eventually, but as can be seen in +/// the benchmarks (as of Feb 2022 at least) `by_ref` can have performance cost. +#[unstable(feature = "std_internals", issue = "none")] +#[derive(Debug)] +pub struct ByRefSized<'a, I>(pub &'a mut I); + +#[unstable(feature = "std_internals", issue = "none")] +impl<I: Iterator> Iterator for ByRefSized<'_, I> { + type Item = I::Item; + + #[inline] + fn next(&mut self) -> Option<Self::Item> { + self.0.next() + } + + #[inline] + fn size_hint(&self) -> (usize, Option<usize>) { + self.0.size_hint() + } + + #[inline] + fn advance_by(&mut self, n: usize) -> Result<(), usize> { + self.0.advance_by(n) + } + + #[inline] + fn nth(&mut self, n: usize) -> Option<Self::Item> { + self.0.nth(n) + } + + #[inline] + fn fold<B, F>(self, init: B, f: F) -> B + where + F: FnMut(B, Self::Item) -> B, + { + self.0.fold(init, f) + } + + #[inline] + fn try_fold<B, F, R>(&mut self, init: B, f: F) -> R + where + F: FnMut(B, Self::Item) -> R, + R: Try<Output = B>, + { + self.0.try_fold(init, f) + } +} + +#[unstable(feature = "std_internals", issue = "none")] +impl<I: DoubleEndedIterator> DoubleEndedIterator for ByRefSized<'_, I> { + #[inline] + fn next_back(&mut self) -> Option<Self::Item> { + self.0.next_back() + } + + #[inline] + fn advance_back_by(&mut self, n: usize) -> Result<(), usize> { + self.0.advance_back_by(n) + } + + #[inline] + fn nth_back(&mut self, n: usize) -> Option<Self::Item> { + self.0.nth_back(n) + } + + #[inline] + fn rfold<B, F>(self, init: B, f: F) -> B + where + F: FnMut(B, Self::Item) -> B, + { + self.0.rfold(init, f) + } + + #[inline] + fn try_rfold<B, F, R>(&mut self, init: B, f: F) -> R + where + F: FnMut(B, Self::Item) -> R, + R: Try<Output = B>, + { + self.0.try_rfold(init, f) + } +} diff --git a/library/core/src/iter/adapters/chain.rs b/library/core/src/iter/adapters/chain.rs new file mode 100644 index 000000000..60eb3a6da --- /dev/null +++ b/library/core/src/iter/adapters/chain.rs @@ -0,0 +1,292 @@ +use crate::iter::{DoubleEndedIterator, FusedIterator, Iterator, TrustedLen}; +use crate::ops::Try; + +/// An iterator that links two iterators together, in a chain. +/// +/// This `struct` is created by [`Iterator::chain`]. See its documentation +/// for more. +/// +/// # Examples +/// +/// ``` +/// use std::iter::Chain; +/// use std::slice::Iter; +/// +/// let a1 = [1, 2, 3]; +/// let a2 = [4, 5, 6]; +/// let iter: Chain<Iter<_>, Iter<_>> = a1.iter().chain(a2.iter()); +/// ``` +#[derive(Clone, Debug)] +#[must_use = "iterators are lazy and do nothing unless consumed"] +#[stable(feature = "rust1", since = "1.0.0")] +pub struct Chain<A, B> { + // These are "fused" with `Option` so we don't need separate state to track which part is + // already exhausted, and we may also get niche layout for `None`. We don't use the real `Fuse` + // adapter because its specialization for `FusedIterator` unconditionally descends into the + // iterator, and that could be expensive to keep revisiting stuff like nested chains. It also + // hurts compiler performance to add more iterator layers to `Chain`. + // + // Only the "first" iterator is actually set `None` when exhausted, depending on whether you + // iterate forward or backward. If you mix directions, then both sides may be `None`. + a: Option<A>, + b: Option<B>, +} +impl<A, B> Chain<A, B> { + pub(in super::super) fn new(a: A, b: B) -> Chain<A, B> { + Chain { a: Some(a), b: Some(b) } + } +} + +#[stable(feature = "rust1", since = "1.0.0")] +impl<A, B> Iterator for Chain<A, B> +where + A: Iterator, + B: Iterator<Item = A::Item>, +{ + type Item = A::Item; + + #[inline] + fn next(&mut self) -> Option<A::Item> { + and_then_or_clear(&mut self.a, Iterator::next).or_else(|| self.b.as_mut()?.next()) + } + + #[inline] + #[rustc_inherit_overflow_checks] + fn count(self) -> usize { + let a_count = match self.a { + Some(a) => a.count(), + None => 0, + }; + let b_count = match self.b { + Some(b) => b.count(), + None => 0, + }; + a_count + b_count + } + + fn try_fold<Acc, F, R>(&mut self, mut acc: Acc, mut f: F) -> R + where + Self: Sized, + F: FnMut(Acc, Self::Item) -> R, + R: Try<Output = Acc>, + { + if let Some(ref mut a) = self.a { + acc = a.try_fold(acc, &mut f)?; + self.a = None; + } + if let Some(ref mut b) = self.b { + acc = b.try_fold(acc, f)?; + // we don't fuse the second iterator + } + try { acc } + } + + fn fold<Acc, F>(self, mut acc: Acc, mut f: F) -> Acc + where + F: FnMut(Acc, Self::Item) -> Acc, + { + if let Some(a) = self.a { + acc = a.fold(acc, &mut f); + } + if let Some(b) = self.b { + acc = b.fold(acc, f); + } + acc + } + + #[inline] + fn advance_by(&mut self, n: usize) -> Result<(), usize> { + let mut rem = n; + + if let Some(ref mut a) = self.a { + match a.advance_by(rem) { + Ok(()) => return Ok(()), + Err(k) => rem -= k, + } + self.a = None; + } + + if let Some(ref mut b) = self.b { + match b.advance_by(rem) { + Ok(()) => return Ok(()), + Err(k) => rem -= k, + } + // we don't fuse the second iterator + } + + if rem == 0 { Ok(()) } else { Err(n - rem) } + } + + #[inline] + fn nth(&mut self, mut n: usize) -> Option<Self::Item> { + if let Some(ref mut a) = self.a { + match a.advance_by(n) { + Ok(()) => match a.next() { + None => n = 0, + x => return x, + }, + Err(k) => n -= k, + } + + self.a = None; + } + + self.b.as_mut()?.nth(n) + } + + #[inline] + fn find<P>(&mut self, mut predicate: P) -> Option<Self::Item> + where + P: FnMut(&Self::Item) -> bool, + { + and_then_or_clear(&mut self.a, |a| a.find(&mut predicate)) + .or_else(|| self.b.as_mut()?.find(predicate)) + } + + #[inline] + fn last(self) -> Option<A::Item> { + // Must exhaust a before b. + let a_last = self.a.and_then(Iterator::last); + let b_last = self.b.and_then(Iterator::last); + b_last.or(a_last) + } + + #[inline] + fn size_hint(&self) -> (usize, Option<usize>) { + match self { + Chain { a: Some(a), b: Some(b) } => { + let (a_lower, a_upper) = a.size_hint(); + let (b_lower, b_upper) = b.size_hint(); + + let lower = a_lower.saturating_add(b_lower); + + let upper = match (a_upper, b_upper) { + (Some(x), Some(y)) => x.checked_add(y), + _ => None, + }; + + (lower, upper) + } + Chain { a: Some(a), b: None } => a.size_hint(), + Chain { a: None, b: Some(b) } => b.size_hint(), + Chain { a: None, b: None } => (0, Some(0)), + } + } +} + +#[stable(feature = "rust1", since = "1.0.0")] +impl<A, B> DoubleEndedIterator for Chain<A, B> +where + A: DoubleEndedIterator, + B: DoubleEndedIterator<Item = A::Item>, +{ + #[inline] + fn next_back(&mut self) -> Option<A::Item> { + and_then_or_clear(&mut self.b, |b| b.next_back()).or_else(|| self.a.as_mut()?.next_back()) + } + + #[inline] + fn advance_back_by(&mut self, n: usize) -> Result<(), usize> { + let mut rem = n; + + if let Some(ref mut b) = self.b { + match b.advance_back_by(rem) { + Ok(()) => return Ok(()), + Err(k) => rem -= k, + } + self.b = None; + } + + if let Some(ref mut a) = self.a { + match a.advance_back_by(rem) { + Ok(()) => return Ok(()), + Err(k) => rem -= k, + } + // we don't fuse the second iterator + } + + if rem == 0 { Ok(()) } else { Err(n - rem) } + } + + #[inline] + fn nth_back(&mut self, mut n: usize) -> Option<Self::Item> { + if let Some(ref mut b) = self.b { + match b.advance_back_by(n) { + Ok(()) => match b.next_back() { + None => n = 0, + x => return x, + }, + Err(k) => n -= k, + } + + self.b = None; + } + + self.a.as_mut()?.nth_back(n) + } + + #[inline] + fn rfind<P>(&mut self, mut predicate: P) -> Option<Self::Item> + where + P: FnMut(&Self::Item) -> bool, + { + and_then_or_clear(&mut self.b, |b| b.rfind(&mut predicate)) + .or_else(|| self.a.as_mut()?.rfind(predicate)) + } + + fn try_rfold<Acc, F, R>(&mut self, mut acc: Acc, mut f: F) -> R + where + Self: Sized, + F: FnMut(Acc, Self::Item) -> R, + R: Try<Output = Acc>, + { + if let Some(ref mut b) = self.b { + acc = b.try_rfold(acc, &mut f)?; + self.b = None; + } + if let Some(ref mut a) = self.a { + acc = a.try_rfold(acc, f)?; + // we don't fuse the second iterator + } + try { acc } + } + + fn rfold<Acc, F>(self, mut acc: Acc, mut f: F) -> Acc + where + F: FnMut(Acc, Self::Item) -> Acc, + { + if let Some(b) = self.b { + acc = b.rfold(acc, &mut f); + } + if let Some(a) = self.a { + acc = a.rfold(acc, f); + } + acc + } +} + +// Note: *both* must be fused to handle double-ended iterators. +#[stable(feature = "fused", since = "1.26.0")] +impl<A, B> FusedIterator for Chain<A, B> +where + A: FusedIterator, + B: FusedIterator<Item = A::Item>, +{ +} + +#[unstable(feature = "trusted_len", issue = "37572")] +unsafe impl<A, B> TrustedLen for Chain<A, B> +where + A: TrustedLen, + B: TrustedLen<Item = A::Item>, +{ +} + +#[inline] +fn and_then_or_clear<T, U>(opt: &mut Option<T>, f: impl FnOnce(&mut T) -> Option<U>) -> Option<U> { + let x = f(opt.as_mut()?); + if x.is_none() { + *opt = None; + } + x +} diff --git a/library/core/src/iter/adapters/cloned.rs b/library/core/src/iter/adapters/cloned.rs new file mode 100644 index 000000000..aba24a79d --- /dev/null +++ b/library/core/src/iter/adapters/cloned.rs @@ -0,0 +1,142 @@ +use crate::iter::adapters::{ + zip::try_get_unchecked, TrustedRandomAccess, TrustedRandomAccessNoCoerce, +}; +use crate::iter::{FusedIterator, TrustedLen}; +use crate::ops::Try; + +/// An iterator that clones the elements of an underlying iterator. +/// +/// This `struct` is created by the [`cloned`] method on [`Iterator`]. See its +/// documentation for more. +/// +/// [`cloned`]: Iterator::cloned +/// [`Iterator`]: trait.Iterator.html +#[stable(feature = "iter_cloned", since = "1.1.0")] +#[must_use = "iterators are lazy and do nothing unless consumed"] +#[derive(Clone, Debug)] +pub struct Cloned<I> { + it: I, +} + +impl<I> Cloned<I> { + pub(in crate::iter) fn new(it: I) -> Cloned<I> { + Cloned { it } + } +} + +fn clone_try_fold<T: Clone, Acc, R>(mut f: impl FnMut(Acc, T) -> R) -> impl FnMut(Acc, &T) -> R { + move |acc, elt| f(acc, elt.clone()) +} + +#[stable(feature = "iter_cloned", since = "1.1.0")] +impl<'a, I, T: 'a> Iterator for Cloned<I> +where + I: Iterator<Item = &'a T>, + T: Clone, +{ + type Item = T; + + fn next(&mut self) -> Option<T> { + self.it.next().cloned() + } + + fn size_hint(&self) -> (usize, Option<usize>) { + self.it.size_hint() + } + + fn try_fold<B, F, R>(&mut self, init: B, f: F) -> R + where + Self: Sized, + F: FnMut(B, Self::Item) -> R, + R: Try<Output = B>, + { + self.it.try_fold(init, clone_try_fold(f)) + } + + fn fold<Acc, F>(self, init: Acc, f: F) -> Acc + where + F: FnMut(Acc, Self::Item) -> Acc, + { + self.it.map(T::clone).fold(init, f) + } + + unsafe fn __iterator_get_unchecked(&mut self, idx: usize) -> T + where + Self: TrustedRandomAccessNoCoerce, + { + // SAFETY: the caller must uphold the contract for + // `Iterator::__iterator_get_unchecked`. + unsafe { try_get_unchecked(&mut self.it, idx).clone() } + } +} + +#[stable(feature = "iter_cloned", since = "1.1.0")] +impl<'a, I, T: 'a> DoubleEndedIterator for Cloned<I> +where + I: DoubleEndedIterator<Item = &'a T>, + T: Clone, +{ + fn next_back(&mut self) -> Option<T> { + self.it.next_back().cloned() + } + + fn try_rfold<B, F, R>(&mut self, init: B, f: F) -> R + where + Self: Sized, + F: FnMut(B, Self::Item) -> R, + R: Try<Output = B>, + { + self.it.try_rfold(init, clone_try_fold(f)) + } + + fn rfold<Acc, F>(self, init: Acc, f: F) -> Acc + where + F: FnMut(Acc, Self::Item) -> Acc, + { + self.it.map(T::clone).rfold(init, f) + } +} + +#[stable(feature = "iter_cloned", since = "1.1.0")] +impl<'a, I, T: 'a> ExactSizeIterator for Cloned<I> +where + I: ExactSizeIterator<Item = &'a T>, + T: Clone, +{ + fn len(&self) -> usize { + self.it.len() + } + + fn is_empty(&self) -> bool { + self.it.is_empty() + } +} + +#[stable(feature = "fused", since = "1.26.0")] +impl<'a, I, T: 'a> FusedIterator for Cloned<I> +where + I: FusedIterator<Item = &'a T>, + T: Clone, +{ +} + +#[doc(hidden)] +#[unstable(feature = "trusted_random_access", issue = "none")] +unsafe impl<I> TrustedRandomAccess for Cloned<I> where I: TrustedRandomAccess {} + +#[doc(hidden)] +#[unstable(feature = "trusted_random_access", issue = "none")] +unsafe impl<I> TrustedRandomAccessNoCoerce for Cloned<I> +where + I: TrustedRandomAccessNoCoerce, +{ + const MAY_HAVE_SIDE_EFFECT: bool = true; +} + +#[unstable(feature = "trusted_len", issue = "37572")] +unsafe impl<'a, I, T: 'a> TrustedLen for Cloned<I> +where + I: TrustedLen<Item = &'a T>, + T: Clone, +{ +} diff --git a/library/core/src/iter/adapters/copied.rs b/library/core/src/iter/adapters/copied.rs new file mode 100644 index 000000000..f9bfd77d7 --- /dev/null +++ b/library/core/src/iter/adapters/copied.rs @@ -0,0 +1,168 @@ +use crate::iter::adapters::{ + zip::try_get_unchecked, TrustedRandomAccess, TrustedRandomAccessNoCoerce, +}; +use crate::iter::{FusedIterator, TrustedLen}; +use crate::ops::Try; + +/// An iterator that copies the elements of an underlying iterator. +/// +/// This `struct` is created by the [`copied`] method on [`Iterator`]. See its +/// documentation for more. +/// +/// [`copied`]: Iterator::copied +/// [`Iterator`]: trait.Iterator.html +#[stable(feature = "iter_copied", since = "1.36.0")] +#[must_use = "iterators are lazy and do nothing unless consumed"] +#[derive(Clone, Debug)] +pub struct Copied<I> { + it: I, +} + +impl<I> Copied<I> { + pub(in crate::iter) fn new(it: I) -> Copied<I> { + Copied { it } + } +} + +fn copy_fold<T: Copy, Acc>(mut f: impl FnMut(Acc, T) -> Acc) -> impl FnMut(Acc, &T) -> Acc { + move |acc, &elt| f(acc, elt) +} + +fn copy_try_fold<T: Copy, Acc, R>(mut f: impl FnMut(Acc, T) -> R) -> impl FnMut(Acc, &T) -> R { + move |acc, &elt| f(acc, elt) +} + +#[stable(feature = "iter_copied", since = "1.36.0")] +impl<'a, I, T: 'a> Iterator for Copied<I> +where + I: Iterator<Item = &'a T>, + T: Copy, +{ + type Item = T; + + fn next(&mut self) -> Option<T> { + self.it.next().copied() + } + + fn size_hint(&self) -> (usize, Option<usize>) { + self.it.size_hint() + } + + fn try_fold<B, F, R>(&mut self, init: B, f: F) -> R + where + Self: Sized, + F: FnMut(B, Self::Item) -> R, + R: Try<Output = B>, + { + self.it.try_fold(init, copy_try_fold(f)) + } + + fn fold<Acc, F>(self, init: Acc, f: F) -> Acc + where + F: FnMut(Acc, Self::Item) -> Acc, + { + self.it.fold(init, copy_fold(f)) + } + + fn nth(&mut self, n: usize) -> Option<T> { + self.it.nth(n).copied() + } + + fn last(self) -> Option<T> { + self.it.last().copied() + } + + fn count(self) -> usize { + self.it.count() + } + + #[inline] + fn advance_by(&mut self, n: usize) -> Result<(), usize> { + self.it.advance_by(n) + } + + unsafe fn __iterator_get_unchecked(&mut self, idx: usize) -> T + where + Self: TrustedRandomAccessNoCoerce, + { + // SAFETY: the caller must uphold the contract for + // `Iterator::__iterator_get_unchecked`. + *unsafe { try_get_unchecked(&mut self.it, idx) } + } +} + +#[stable(feature = "iter_copied", since = "1.36.0")] +impl<'a, I, T: 'a> DoubleEndedIterator for Copied<I> +where + I: DoubleEndedIterator<Item = &'a T>, + T: Copy, +{ + fn next_back(&mut self) -> Option<T> { + self.it.next_back().copied() + } + + fn try_rfold<B, F, R>(&mut self, init: B, f: F) -> R + where + Self: Sized, + F: FnMut(B, Self::Item) -> R, + R: Try<Output = B>, + { + self.it.try_rfold(init, copy_try_fold(f)) + } + + fn rfold<Acc, F>(self, init: Acc, f: F) -> Acc + where + F: FnMut(Acc, Self::Item) -> Acc, + { + self.it.rfold(init, copy_fold(f)) + } + + #[inline] + fn advance_back_by(&mut self, n: usize) -> Result<(), usize> { + self.it.advance_back_by(n) + } +} + +#[stable(feature = "iter_copied", since = "1.36.0")] +impl<'a, I, T: 'a> ExactSizeIterator for Copied<I> +where + I: ExactSizeIterator<Item = &'a T>, + T: Copy, +{ + fn len(&self) -> usize { + self.it.len() + } + + fn is_empty(&self) -> bool { + self.it.is_empty() + } +} + +#[stable(feature = "iter_copied", since = "1.36.0")] +impl<'a, I, T: 'a> FusedIterator for Copied<I> +where + I: FusedIterator<Item = &'a T>, + T: Copy, +{ +} + +#[doc(hidden)] +#[unstable(feature = "trusted_random_access", issue = "none")] +unsafe impl<I> TrustedRandomAccess for Copied<I> where I: TrustedRandomAccess {} + +#[doc(hidden)] +#[unstable(feature = "trusted_random_access", issue = "none")] +unsafe impl<I> TrustedRandomAccessNoCoerce for Copied<I> +where + I: TrustedRandomAccessNoCoerce, +{ + const MAY_HAVE_SIDE_EFFECT: bool = I::MAY_HAVE_SIDE_EFFECT; +} + +#[stable(feature = "iter_copied", since = "1.36.0")] +unsafe impl<'a, I, T: 'a> TrustedLen for Copied<I> +where + I: TrustedLen<Item = &'a T>, + T: Copy, +{ +} diff --git a/library/core/src/iter/adapters/cycle.rs b/library/core/src/iter/adapters/cycle.rs new file mode 100644 index 000000000..02b593907 --- /dev/null +++ b/library/core/src/iter/adapters/cycle.rs @@ -0,0 +1,108 @@ +use crate::{iter::FusedIterator, ops::Try}; + +/// An iterator that repeats endlessly. +/// +/// This `struct` is created by the [`cycle`] method on [`Iterator`]. See its +/// documentation for more. +/// +/// [`cycle`]: Iterator::cycle +/// [`Iterator`]: trait.Iterator.html +#[derive(Clone, Debug)] +#[must_use = "iterators are lazy and do nothing unless consumed"] +#[stable(feature = "rust1", since = "1.0.0")] +pub struct Cycle<I> { + orig: I, + iter: I, +} + +impl<I: Clone> Cycle<I> { + pub(in crate::iter) fn new(iter: I) -> Cycle<I> { + Cycle { orig: iter.clone(), iter } + } +} + +#[stable(feature = "rust1", since = "1.0.0")] +impl<I> Iterator for Cycle<I> +where + I: Clone + Iterator, +{ + type Item = <I as Iterator>::Item; + + #[inline] + fn next(&mut self) -> Option<<I as Iterator>::Item> { + match self.iter.next() { + None => { + self.iter = self.orig.clone(); + self.iter.next() + } + y => y, + } + } + + #[inline] + fn size_hint(&self) -> (usize, Option<usize>) { + // the cycle iterator is either empty or infinite + match self.orig.size_hint() { + sz @ (0, Some(0)) => sz, + (0, _) => (0, None), + _ => (usize::MAX, None), + } + } + + #[inline] + fn try_fold<Acc, F, R>(&mut self, mut acc: Acc, mut f: F) -> R + where + F: FnMut(Acc, Self::Item) -> R, + R: Try<Output = Acc>, + { + // fully iterate the current iterator. this is necessary because + // `self.iter` may be empty even when `self.orig` isn't + acc = self.iter.try_fold(acc, &mut f)?; + self.iter = self.orig.clone(); + + // complete a full cycle, keeping track of whether the cycled + // iterator is empty or not. we need to return early in case + // of an empty iterator to prevent an infinite loop + let mut is_empty = true; + acc = self.iter.try_fold(acc, |acc, x| { + is_empty = false; + f(acc, x) + })?; + + if is_empty { + return try { acc }; + } + + loop { + self.iter = self.orig.clone(); + acc = self.iter.try_fold(acc, &mut f)?; + } + } + + #[inline] + #[rustc_inherit_overflow_checks] + fn advance_by(&mut self, n: usize) -> Result<(), usize> { + let mut rem = n; + match self.iter.advance_by(rem) { + ret @ Ok(_) => return ret, + Err(advanced) => rem -= advanced, + } + + while rem > 0 { + self.iter = self.orig.clone(); + match self.iter.advance_by(rem) { + ret @ Ok(_) => return ret, + Err(0) => return Err(n - rem), + Err(advanced) => rem -= advanced, + } + } + + Ok(()) + } + + // No `fold` override, because `fold` doesn't make much sense for `Cycle`, + // and we can't do anything better than the default. +} + +#[stable(feature = "fused", since = "1.26.0")] +impl<I> FusedIterator for Cycle<I> where I: Clone + Iterator {} diff --git a/library/core/src/iter/adapters/enumerate.rs b/library/core/src/iter/adapters/enumerate.rs new file mode 100644 index 000000000..14a126951 --- /dev/null +++ b/library/core/src/iter/adapters/enumerate.rs @@ -0,0 +1,266 @@ +use crate::iter::adapters::{ + zip::try_get_unchecked, SourceIter, TrustedRandomAccess, TrustedRandomAccessNoCoerce, +}; +use crate::iter::{FusedIterator, InPlaceIterable, TrustedLen}; +use crate::ops::Try; + +/// An iterator that yields the current count and the element during iteration. +/// +/// This `struct` is created by the [`enumerate`] method on [`Iterator`]. See its +/// documentation for more. +/// +/// [`enumerate`]: Iterator::enumerate +/// [`Iterator`]: trait.Iterator.html +#[derive(Clone, Debug)] +#[must_use = "iterators are lazy and do nothing unless consumed"] +#[stable(feature = "rust1", since = "1.0.0")] +pub struct Enumerate<I> { + iter: I, + count: usize, +} +impl<I> Enumerate<I> { + pub(in crate::iter) fn new(iter: I) -> Enumerate<I> { + Enumerate { iter, count: 0 } + } +} + +#[stable(feature = "rust1", since = "1.0.0")] +impl<I> Iterator for Enumerate<I> +where + I: Iterator, +{ + type Item = (usize, <I as Iterator>::Item); + + /// # Overflow Behavior + /// + /// The method does no guarding against overflows, so enumerating more than + /// `usize::MAX` elements either produces the wrong result or panics. If + /// debug assertions are enabled, a panic is guaranteed. + /// + /// # Panics + /// + /// Might panic if the index of the element overflows a `usize`. + #[inline] + #[rustc_inherit_overflow_checks] + fn next(&mut self) -> Option<(usize, <I as Iterator>::Item)> { + let a = self.iter.next()?; + let i = self.count; + self.count += 1; + Some((i, a)) + } + + #[inline] + fn size_hint(&self) -> (usize, Option<usize>) { + self.iter.size_hint() + } + + #[inline] + #[rustc_inherit_overflow_checks] + fn nth(&mut self, n: usize) -> Option<(usize, I::Item)> { + let a = self.iter.nth(n)?; + let i = self.count + n; + self.count = i + 1; + Some((i, a)) + } + + #[inline] + fn count(self) -> usize { + self.iter.count() + } + + #[inline] + fn try_fold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> R, + R: Try<Output = Acc>, + { + #[inline] + fn enumerate<'a, T, Acc, R>( + count: &'a mut usize, + mut fold: impl FnMut(Acc, (usize, T)) -> R + 'a, + ) -> impl FnMut(Acc, T) -> R + 'a { + #[rustc_inherit_overflow_checks] + move |acc, item| { + let acc = fold(acc, (*count, item)); + *count += 1; + acc + } + } + + self.iter.try_fold(init, enumerate(&mut self.count, fold)) + } + + #[inline] + fn fold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc + where + Fold: FnMut(Acc, Self::Item) -> Acc, + { + #[inline] + fn enumerate<T, Acc>( + mut count: usize, + mut fold: impl FnMut(Acc, (usize, T)) -> Acc, + ) -> impl FnMut(Acc, T) -> Acc { + #[rustc_inherit_overflow_checks] + move |acc, item| { + let acc = fold(acc, (count, item)); + count += 1; + acc + } + } + + self.iter.fold(init, enumerate(self.count, fold)) + } + + #[inline] + #[rustc_inherit_overflow_checks] + fn advance_by(&mut self, n: usize) -> Result<(), usize> { + match self.iter.advance_by(n) { + ret @ Ok(_) => { + self.count += n; + ret + } + ret @ Err(advanced) => { + self.count += advanced; + ret + } + } + } + + #[rustc_inherit_overflow_checks] + #[inline] + unsafe fn __iterator_get_unchecked(&mut self, idx: usize) -> <Self as Iterator>::Item + where + Self: TrustedRandomAccessNoCoerce, + { + // SAFETY: the caller must uphold the contract for + // `Iterator::__iterator_get_unchecked`. + let value = unsafe { try_get_unchecked(&mut self.iter, idx) }; + (self.count + idx, value) + } +} + +#[stable(feature = "rust1", since = "1.0.0")] +impl<I> DoubleEndedIterator for Enumerate<I> +where + I: ExactSizeIterator + DoubleEndedIterator, +{ + #[inline] + fn next_back(&mut self) -> Option<(usize, <I as Iterator>::Item)> { + let a = self.iter.next_back()?; + let len = self.iter.len(); + // Can safely add, `ExactSizeIterator` promises that the number of + // elements fits into a `usize`. + Some((self.count + len, a)) + } + + #[inline] + fn nth_back(&mut self, n: usize) -> Option<(usize, <I as Iterator>::Item)> { + let a = self.iter.nth_back(n)?; + let len = self.iter.len(); + // Can safely add, `ExactSizeIterator` promises that the number of + // elements fits into a `usize`. + Some((self.count + len, a)) + } + + #[inline] + fn try_rfold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> R, + R: Try<Output = Acc>, + { + // Can safely add and subtract the count, as `ExactSizeIterator` promises + // that the number of elements fits into a `usize`. + fn enumerate<T, Acc, R>( + mut count: usize, + mut fold: impl FnMut(Acc, (usize, T)) -> R, + ) -> impl FnMut(Acc, T) -> R { + move |acc, item| { + count -= 1; + fold(acc, (count, item)) + } + } + + let count = self.count + self.iter.len(); + self.iter.try_rfold(init, enumerate(count, fold)) + } + + #[inline] + fn rfold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc + where + Fold: FnMut(Acc, Self::Item) -> Acc, + { + // Can safely add and subtract the count, as `ExactSizeIterator` promises + // that the number of elements fits into a `usize`. + fn enumerate<T, Acc>( + mut count: usize, + mut fold: impl FnMut(Acc, (usize, T)) -> Acc, + ) -> impl FnMut(Acc, T) -> Acc { + move |acc, item| { + count -= 1; + fold(acc, (count, item)) + } + } + + let count = self.count + self.iter.len(); + self.iter.rfold(init, enumerate(count, fold)) + } + + #[inline] + fn advance_back_by(&mut self, n: usize) -> Result<(), usize> { + // we do not need to update the count since that only tallies the number of items + // consumed from the front. consuming items from the back can never reduce that. + self.iter.advance_back_by(n) + } +} + +#[stable(feature = "rust1", since = "1.0.0")] +impl<I> ExactSizeIterator for Enumerate<I> +where + I: ExactSizeIterator, +{ + fn len(&self) -> usize { + self.iter.len() + } + + fn is_empty(&self) -> bool { + self.iter.is_empty() + } +} + +#[doc(hidden)] +#[unstable(feature = "trusted_random_access", issue = "none")] +unsafe impl<I> TrustedRandomAccess for Enumerate<I> where I: TrustedRandomAccess {} + +#[doc(hidden)] +#[unstable(feature = "trusted_random_access", issue = "none")] +unsafe impl<I> TrustedRandomAccessNoCoerce for Enumerate<I> +where + I: TrustedRandomAccessNoCoerce, +{ + const MAY_HAVE_SIDE_EFFECT: bool = I::MAY_HAVE_SIDE_EFFECT; +} + +#[stable(feature = "fused", since = "1.26.0")] +impl<I> FusedIterator for Enumerate<I> where I: FusedIterator {} + +#[unstable(feature = "trusted_len", issue = "37572")] +unsafe impl<I> TrustedLen for Enumerate<I> where I: TrustedLen {} + +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl<I> SourceIter for Enumerate<I> +where + I: SourceIter, +{ + type Source = I::Source; + + #[inline] + unsafe fn as_inner(&mut self) -> &mut I::Source { + // SAFETY: unsafe function forwarding to unsafe function with the same requirements + unsafe { SourceIter::as_inner(&mut self.iter) } + } +} + +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl<I: InPlaceIterable> InPlaceIterable for Enumerate<I> {} diff --git a/library/core/src/iter/adapters/filter.rs b/library/core/src/iter/adapters/filter.rs new file mode 100644 index 000000000..a0afaa326 --- /dev/null +++ b/library/core/src/iter/adapters/filter.rs @@ -0,0 +1,152 @@ +use crate::fmt; +use crate::iter::{adapters::SourceIter, FusedIterator, InPlaceIterable}; +use crate::ops::Try; + +/// An iterator that filters the elements of `iter` with `predicate`. +/// +/// This `struct` is created by the [`filter`] method on [`Iterator`]. See its +/// documentation for more. +/// +/// [`filter`]: Iterator::filter +/// [`Iterator`]: trait.Iterator.html +#[must_use = "iterators are lazy and do nothing unless consumed"] +#[stable(feature = "rust1", since = "1.0.0")] +#[derive(Clone)] +pub struct Filter<I, P> { + // Used for `SplitWhitespace` and `SplitAsciiWhitespace` `as_str` methods + pub(crate) iter: I, + predicate: P, +} +impl<I, P> Filter<I, P> { + pub(in crate::iter) fn new(iter: I, predicate: P) -> Filter<I, P> { + Filter { iter, predicate } + } +} + +#[stable(feature = "core_impl_debug", since = "1.9.0")] +impl<I: fmt::Debug, P> fmt::Debug for Filter<I, P> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Filter").field("iter", &self.iter).finish() + } +} + +fn filter_fold<T, Acc>( + mut predicate: impl FnMut(&T) -> bool, + mut fold: impl FnMut(Acc, T) -> Acc, +) -> impl FnMut(Acc, T) -> Acc { + move |acc, item| if predicate(&item) { fold(acc, item) } else { acc } +} + +fn filter_try_fold<'a, T, Acc, R: Try<Output = Acc>>( + predicate: &'a mut impl FnMut(&T) -> bool, + mut fold: impl FnMut(Acc, T) -> R + 'a, +) -> impl FnMut(Acc, T) -> R + 'a { + move |acc, item| if predicate(&item) { fold(acc, item) } else { try { acc } } +} + +#[stable(feature = "rust1", since = "1.0.0")] +impl<I: Iterator, P> Iterator for Filter<I, P> +where + P: FnMut(&I::Item) -> bool, +{ + type Item = I::Item; + + #[inline] + fn next(&mut self) -> Option<I::Item> { + self.iter.find(&mut self.predicate) + } + + #[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 + } + + // this special case allows the compiler to make `.filter(_).count()` + // branchless. Barring perfect branch prediction (which is unattainable in + // the general case), this will be much faster in >90% of cases (containing + // virtually all real workloads) and only a tiny bit slower in the rest. + // + // Having this specialization thus allows us to write `.filter(p).count()` + // where we would otherwise write `.map(|x| p(x) as usize).sum()`, which is + // less readable and also less backwards-compatible to Rust before 1.10. + // + // Using the branchless version will also simplify the LLVM byte code, thus + // leaving more budget for LLVM optimizations. + #[inline] + fn count(self) -> usize { + #[inline] + fn to_usize<T>(mut predicate: impl FnMut(&T) -> bool) -> impl FnMut(T) -> usize { + move |x| predicate(&x) as usize + } + + self.iter.map(to_usize(self.predicate)).sum() + } + + #[inline] + fn try_fold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> R, + R: Try<Output = Acc>, + { + self.iter.try_fold(init, filter_try_fold(&mut self.predicate, fold)) + } + + #[inline] + fn fold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc + where + Fold: FnMut(Acc, Self::Item) -> Acc, + { + self.iter.fold(init, filter_fold(self.predicate, fold)) + } +} + +#[stable(feature = "rust1", since = "1.0.0")] +impl<I: DoubleEndedIterator, P> DoubleEndedIterator for Filter<I, P> +where + P: FnMut(&I::Item) -> bool, +{ + #[inline] + fn next_back(&mut self) -> Option<I::Item> { + self.iter.rfind(&mut self.predicate) + } + + #[inline] + fn try_rfold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> R, + R: Try<Output = Acc>, + { + self.iter.try_rfold(init, filter_try_fold(&mut self.predicate, fold)) + } + + #[inline] + fn rfold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc + where + Fold: FnMut(Acc, Self::Item) -> Acc, + { + self.iter.rfold(init, filter_fold(self.predicate, fold)) + } +} + +#[stable(feature = "fused", since = "1.26.0")] +impl<I: FusedIterator, P> FusedIterator for Filter<I, P> where P: FnMut(&I::Item) -> bool {} + +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl<P, I> SourceIter for Filter<I, P> +where + I: SourceIter, +{ + type Source = I::Source; + + #[inline] + unsafe fn as_inner(&mut self) -> &mut I::Source { + // SAFETY: unsafe function forwarding to unsafe function with the same requirements + unsafe { SourceIter::as_inner(&mut self.iter) } + } +} + +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl<I: InPlaceIterable, P> InPlaceIterable for Filter<I, P> where P: FnMut(&I::Item) -> bool {} diff --git a/library/core/src/iter/adapters/filter_map.rs b/library/core/src/iter/adapters/filter_map.rs new file mode 100644 index 000000000..e0d665c9e --- /dev/null +++ b/library/core/src/iter/adapters/filter_map.rs @@ -0,0 +1,149 @@ +use crate::fmt; +use crate::iter::{adapters::SourceIter, FusedIterator, InPlaceIterable}; +use crate::ops::{ControlFlow, Try}; + +/// An iterator that uses `f` to both filter and map elements from `iter`. +/// +/// This `struct` is created by the [`filter_map`] method on [`Iterator`]. See its +/// documentation for more. +/// +/// [`filter_map`]: Iterator::filter_map +/// [`Iterator`]: trait.Iterator.html +#[must_use = "iterators are lazy and do nothing unless consumed"] +#[stable(feature = "rust1", since = "1.0.0")] +#[derive(Clone)] +pub struct FilterMap<I, F> { + iter: I, + f: F, +} +impl<I, F> FilterMap<I, F> { + pub(in crate::iter) fn new(iter: I, f: F) -> FilterMap<I, F> { + FilterMap { iter, f } + } +} + +#[stable(feature = "core_impl_debug", since = "1.9.0")] +impl<I: fmt::Debug, F> fmt::Debug for FilterMap<I, F> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("FilterMap").field("iter", &self.iter).finish() + } +} + +fn filter_map_fold<T, B, Acc>( + mut f: impl FnMut(T) -> Option<B>, + mut fold: impl FnMut(Acc, B) -> Acc, +) -> impl FnMut(Acc, T) -> Acc { + move |acc, item| match f(item) { + Some(x) => fold(acc, x), + None => acc, + } +} + +fn filter_map_try_fold<'a, T, B, Acc, R: Try<Output = Acc>>( + f: &'a mut impl FnMut(T) -> Option<B>, + mut fold: impl FnMut(Acc, B) -> R + 'a, +) -> impl FnMut(Acc, T) -> R + 'a { + move |acc, item| match f(item) { + Some(x) => fold(acc, x), + None => try { acc }, + } +} + +#[stable(feature = "rust1", since = "1.0.0")] +impl<B, I: Iterator, F> Iterator for FilterMap<I, F> +where + F: FnMut(I::Item) -> Option<B>, +{ + type Item = B; + + #[inline] + fn next(&mut self) -> Option<B> { + self.iter.find_map(&mut self.f) + } + + #[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 + } + + #[inline] + fn try_fold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> R, + R: Try<Output = Acc>, + { + self.iter.try_fold(init, filter_map_try_fold(&mut self.f, fold)) + } + + #[inline] + fn fold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc + where + Fold: FnMut(Acc, Self::Item) -> Acc, + { + self.iter.fold(init, filter_map_fold(self.f, fold)) + } +} + +#[stable(feature = "rust1", since = "1.0.0")] +impl<B, I: DoubleEndedIterator, F> DoubleEndedIterator for FilterMap<I, F> +where + F: FnMut(I::Item) -> Option<B>, +{ + #[inline] + fn next_back(&mut self) -> Option<B> { + #[inline] + fn find<T, B>( + f: &mut impl FnMut(T) -> Option<B>, + ) -> impl FnMut((), T) -> ControlFlow<B> + '_ { + move |(), x| match f(x) { + Some(x) => ControlFlow::Break(x), + None => ControlFlow::CONTINUE, + } + } + + self.iter.try_rfold((), find(&mut self.f)).break_value() + } + + #[inline] + fn try_rfold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> R, + R: Try<Output = Acc>, + { + self.iter.try_rfold(init, filter_map_try_fold(&mut self.f, fold)) + } + + #[inline] + fn rfold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc + where + Fold: FnMut(Acc, Self::Item) -> Acc, + { + self.iter.rfold(init, filter_map_fold(self.f, fold)) + } +} + +#[stable(feature = "fused", since = "1.26.0")] +impl<B, I: FusedIterator, F> FusedIterator for FilterMap<I, F> where F: FnMut(I::Item) -> Option<B> {} + +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl<I, F> SourceIter for FilterMap<I, F> +where + I: SourceIter, +{ + type Source = I::Source; + + #[inline] + unsafe fn as_inner(&mut self) -> &mut I::Source { + // SAFETY: unsafe function forwarding to unsafe function with the same requirements + unsafe { SourceIter::as_inner(&mut self.iter) } + } +} + +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl<B, I: InPlaceIterable, F> InPlaceIterable for FilterMap<I, F> where + F: FnMut(I::Item) -> Option<B> +{ +} diff --git a/library/core/src/iter/adapters/flatten.rs b/library/core/src/iter/adapters/flatten.rs new file mode 100644 index 000000000..15a120e35 --- /dev/null +++ b/library/core/src/iter/adapters/flatten.rs @@ -0,0 +1,599 @@ +use crate::fmt; +use crate::iter::{DoubleEndedIterator, Fuse, FusedIterator, Iterator, Map, TrustedLen}; +use crate::ops::Try; + +/// An iterator that maps each element to an iterator, and yields the elements +/// of the produced iterators. +/// +/// This `struct` is created by [`Iterator::flat_map`]. See its documentation +/// for more. +#[must_use = "iterators are lazy and do nothing unless consumed"] +#[stable(feature = "rust1", since = "1.0.0")] +pub struct FlatMap<I, U: IntoIterator, F> { + inner: FlattenCompat<Map<I, F>, <U as IntoIterator>::IntoIter>, +} + +impl<I: Iterator, U: IntoIterator, F: FnMut(I::Item) -> U> FlatMap<I, U, F> { + pub(in crate::iter) fn new(iter: I, f: F) -> FlatMap<I, U, F> { + FlatMap { inner: FlattenCompat::new(iter.map(f)) } + } +} + +#[stable(feature = "rust1", since = "1.0.0")] +impl<I: Clone, U, F: Clone> Clone for FlatMap<I, U, F> +where + U: Clone + IntoIterator<IntoIter: Clone>, +{ + fn clone(&self) -> Self { + FlatMap { inner: self.inner.clone() } + } +} + +#[stable(feature = "core_impl_debug", since = "1.9.0")] +impl<I: fmt::Debug, U, F> fmt::Debug for FlatMap<I, U, F> +where + U: IntoIterator<IntoIter: fmt::Debug>, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("FlatMap").field("inner", &self.inner).finish() + } +} + +#[stable(feature = "rust1", since = "1.0.0")] +impl<I: Iterator, U: IntoIterator, F> Iterator for FlatMap<I, U, F> +where + F: FnMut(I::Item) -> U, +{ + type Item = U::Item; + + #[inline] + fn next(&mut self) -> Option<U::Item> { + self.inner.next() + } + + #[inline] + fn size_hint(&self) -> (usize, Option<usize>) { + self.inner.size_hint() + } + + #[inline] + fn try_fold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> R, + R: Try<Output = Acc>, + { + self.inner.try_fold(init, fold) + } + + #[inline] + fn fold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc + where + Fold: FnMut(Acc, Self::Item) -> Acc, + { + self.inner.fold(init, fold) + } +} + +#[stable(feature = "rust1", since = "1.0.0")] +impl<I: DoubleEndedIterator, U, F> DoubleEndedIterator for FlatMap<I, U, F> +where + F: FnMut(I::Item) -> U, + U: IntoIterator<IntoIter: DoubleEndedIterator>, +{ + #[inline] + fn next_back(&mut self) -> Option<U::Item> { + self.inner.next_back() + } + + #[inline] + fn try_rfold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> R, + R: Try<Output = Acc>, + { + self.inner.try_rfold(init, fold) + } + + #[inline] + fn rfold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc + where + Fold: FnMut(Acc, Self::Item) -> Acc, + { + self.inner.rfold(init, fold) + } +} + +#[stable(feature = "fused", since = "1.26.0")] +impl<I, U, F> FusedIterator for FlatMap<I, U, F> +where + I: FusedIterator, + U: IntoIterator, + F: FnMut(I::Item) -> U, +{ +} + +#[unstable(feature = "trusted_len", issue = "37572")] +unsafe impl<T, I, F, const N: usize> TrustedLen for FlatMap<I, [T; N], F> +where + I: TrustedLen, + F: FnMut(I::Item) -> [T; N], +{ +} + +#[unstable(feature = "trusted_len", issue = "37572")] +unsafe impl<'a, T, I, F, const N: usize> TrustedLen for FlatMap<I, &'a [T; N], F> +where + I: TrustedLen, + F: FnMut(I::Item) -> &'a [T; N], +{ +} + +#[unstable(feature = "trusted_len", issue = "37572")] +unsafe impl<'a, T, I, F, const N: usize> TrustedLen for FlatMap<I, &'a mut [T; N], F> +where + I: TrustedLen, + F: FnMut(I::Item) -> &'a mut [T; N], +{ +} + +/// An iterator that flattens one level of nesting in an iterator of things +/// that can be turned into iterators. +/// +/// This `struct` is created by the [`flatten`] method on [`Iterator`]. See its +/// documentation for more. +/// +/// [`flatten`]: Iterator::flatten() +#[must_use = "iterators are lazy and do nothing unless consumed"] +#[stable(feature = "iterator_flatten", since = "1.29.0")] +pub struct Flatten<I: Iterator<Item: IntoIterator>> { + inner: FlattenCompat<I, <I::Item as IntoIterator>::IntoIter>, +} + +impl<I: Iterator<Item: IntoIterator>> Flatten<I> { + pub(in super::super) fn new(iter: I) -> Flatten<I> { + Flatten { inner: FlattenCompat::new(iter) } + } +} + +#[stable(feature = "iterator_flatten", since = "1.29.0")] +impl<I, U> fmt::Debug for Flatten<I> +where + I: fmt::Debug + Iterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>, + U: fmt::Debug + Iterator, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Flatten").field("inner", &self.inner).finish() + } +} + +#[stable(feature = "iterator_flatten", since = "1.29.0")] +impl<I, U> Clone for Flatten<I> +where + I: Clone + Iterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>, + U: Clone + Iterator, +{ + fn clone(&self) -> Self { + Flatten { inner: self.inner.clone() } + } +} + +#[stable(feature = "iterator_flatten", since = "1.29.0")] +impl<I, U> Iterator for Flatten<I> +where + I: Iterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>, + U: Iterator, +{ + type Item = U::Item; + + #[inline] + fn next(&mut self) -> Option<U::Item> { + self.inner.next() + } + + #[inline] + fn size_hint(&self) -> (usize, Option<usize>) { + self.inner.size_hint() + } + + #[inline] + fn try_fold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> R, + R: Try<Output = Acc>, + { + self.inner.try_fold(init, fold) + } + + #[inline] + fn fold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc + where + Fold: FnMut(Acc, Self::Item) -> Acc, + { + self.inner.fold(init, fold) + } +} + +#[stable(feature = "iterator_flatten", since = "1.29.0")] +impl<I, U> DoubleEndedIterator for Flatten<I> +where + I: DoubleEndedIterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>, + U: DoubleEndedIterator, +{ + #[inline] + fn next_back(&mut self) -> Option<U::Item> { + self.inner.next_back() + } + + #[inline] + fn try_rfold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> R, + R: Try<Output = Acc>, + { + self.inner.try_rfold(init, fold) + } + + #[inline] + fn rfold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc + where + Fold: FnMut(Acc, Self::Item) -> Acc, + { + self.inner.rfold(init, fold) + } +} + +#[stable(feature = "iterator_flatten", since = "1.29.0")] +impl<I, U> FusedIterator for Flatten<I> +where + I: FusedIterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>, + U: Iterator, +{ +} + +#[unstable(feature = "trusted_len", issue = "37572")] +unsafe impl<I> TrustedLen for Flatten<I> +where + I: TrustedLen, + <I as Iterator>::Item: TrustedConstSize, +{ +} + +/// Real logic of both `Flatten` and `FlatMap` which simply delegate to +/// this type. +#[derive(Clone, Debug)] +struct FlattenCompat<I, U> { + iter: Fuse<I>, + frontiter: Option<U>, + backiter: Option<U>, +} +impl<I, U> FlattenCompat<I, U> +where + I: Iterator, +{ + /// Adapts an iterator by flattening it, for use in `flatten()` and `flat_map()`. + fn new(iter: I) -> FlattenCompat<I, U> { + FlattenCompat { iter: iter.fuse(), frontiter: None, backiter: None } + } +} + +impl<I, U> Iterator for FlattenCompat<I, U> +where + I: Iterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>, + U: Iterator, +{ + type Item = U::Item; + + #[inline] + fn next(&mut self) -> Option<U::Item> { + loop { + if let elt @ Some(_) = and_then_or_clear(&mut self.frontiter, Iterator::next) { + return elt; + } + match self.iter.next() { + None => return and_then_or_clear(&mut self.backiter, Iterator::next), + Some(inner) => self.frontiter = Some(inner.into_iter()), + } + } + } + + #[inline] + fn size_hint(&self) -> (usize, Option<usize>) { + let (flo, fhi) = self.frontiter.as_ref().map_or((0, Some(0)), U::size_hint); + let (blo, bhi) = self.backiter.as_ref().map_or((0, Some(0)), U::size_hint); + let lo = flo.saturating_add(blo); + + if let Some(fixed_size) = <<I as Iterator>::Item as ConstSizeIntoIterator>::size() { + let (lower, upper) = self.iter.size_hint(); + + let lower = lower.saturating_mul(fixed_size).saturating_add(lo); + let upper = + try { fhi?.checked_add(bhi?)?.checked_add(fixed_size.checked_mul(upper?)?)? }; + + return (lower, upper); + } + + match (self.iter.size_hint(), fhi, bhi) { + ((0, Some(0)), Some(a), Some(b)) => (lo, a.checked_add(b)), + _ => (lo, None), + } + } + + #[inline] + fn try_fold<Acc, Fold, R>(&mut self, mut init: Acc, mut fold: Fold) -> R + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> R, + R: Try<Output = Acc>, + { + #[inline] + fn flatten<'a, T: IntoIterator, Acc, R: Try<Output = Acc>>( + frontiter: &'a mut Option<T::IntoIter>, + fold: &'a mut impl FnMut(Acc, T::Item) -> R, + ) -> impl FnMut(Acc, T) -> R + 'a { + move |acc, x| { + let mut mid = x.into_iter(); + let r = mid.try_fold(acc, &mut *fold); + *frontiter = Some(mid); + r + } + } + + if let Some(ref mut front) = self.frontiter { + init = front.try_fold(init, &mut fold)?; + } + self.frontiter = None; + + init = self.iter.try_fold(init, flatten(&mut self.frontiter, &mut fold))?; + self.frontiter = None; + + if let Some(ref mut back) = self.backiter { + init = back.try_fold(init, &mut fold)?; + } + self.backiter = None; + + try { init } + } + + #[inline] + fn fold<Acc, Fold>(self, mut init: Acc, mut fold: Fold) -> Acc + where + Fold: FnMut(Acc, Self::Item) -> Acc, + { + #[inline] + fn flatten<T: IntoIterator, Acc>( + fold: &mut impl FnMut(Acc, T::Item) -> Acc, + ) -> impl FnMut(Acc, T) -> Acc + '_ { + move |acc, x| x.into_iter().fold(acc, &mut *fold) + } + + if let Some(front) = self.frontiter { + init = front.fold(init, &mut fold); + } + + init = self.iter.fold(init, flatten(&mut fold)); + + if let Some(back) = self.backiter { + init = back.fold(init, &mut fold); + } + + init + } + + #[inline] + #[rustc_inherit_overflow_checks] + fn advance_by(&mut self, n: usize) -> Result<(), usize> { + let mut rem = n; + loop { + if let Some(ref mut front) = self.frontiter { + match front.advance_by(rem) { + ret @ Ok(_) => return ret, + Err(advanced) => rem -= advanced, + } + } + self.frontiter = match self.iter.next() { + Some(iterable) => Some(iterable.into_iter()), + _ => break, + } + } + + self.frontiter = None; + + if let Some(ref mut back) = self.backiter { + match back.advance_by(rem) { + ret @ Ok(_) => return ret, + Err(advanced) => rem -= advanced, + } + } + + if rem > 0 { + return Err(n - rem); + } + + self.backiter = None; + + Ok(()) + } +} + +impl<I, U> DoubleEndedIterator for FlattenCompat<I, U> +where + I: DoubleEndedIterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>, + U: DoubleEndedIterator, +{ + #[inline] + fn next_back(&mut self) -> Option<U::Item> { + loop { + if let elt @ Some(_) = and_then_or_clear(&mut self.backiter, |b| b.next_back()) { + return elt; + } + match self.iter.next_back() { + None => return and_then_or_clear(&mut self.frontiter, |f| f.next_back()), + Some(inner) => self.backiter = Some(inner.into_iter()), + } + } + } + + #[inline] + fn try_rfold<Acc, Fold, R>(&mut self, mut init: Acc, mut fold: Fold) -> R + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> R, + R: Try<Output = Acc>, + { + #[inline] + fn flatten<'a, T: IntoIterator, Acc, R: Try<Output = Acc>>( + backiter: &'a mut Option<T::IntoIter>, + fold: &'a mut impl FnMut(Acc, T::Item) -> R, + ) -> impl FnMut(Acc, T) -> R + 'a + where + T::IntoIter: DoubleEndedIterator, + { + move |acc, x| { + let mut mid = x.into_iter(); + let r = mid.try_rfold(acc, &mut *fold); + *backiter = Some(mid); + r + } + } + + if let Some(ref mut back) = self.backiter { + init = back.try_rfold(init, &mut fold)?; + } + self.backiter = None; + + init = self.iter.try_rfold(init, flatten(&mut self.backiter, &mut fold))?; + self.backiter = None; + + if let Some(ref mut front) = self.frontiter { + init = front.try_rfold(init, &mut fold)?; + } + self.frontiter = None; + + try { init } + } + + #[inline] + fn rfold<Acc, Fold>(self, mut init: Acc, mut fold: Fold) -> Acc + where + Fold: FnMut(Acc, Self::Item) -> Acc, + { + #[inline] + fn flatten<T: IntoIterator, Acc>( + fold: &mut impl FnMut(Acc, T::Item) -> Acc, + ) -> impl FnMut(Acc, T) -> Acc + '_ + where + T::IntoIter: DoubleEndedIterator, + { + move |acc, x| x.into_iter().rfold(acc, &mut *fold) + } + + if let Some(back) = self.backiter { + init = back.rfold(init, &mut fold); + } + + init = self.iter.rfold(init, flatten(&mut fold)); + + if let Some(front) = self.frontiter { + init = front.rfold(init, &mut fold); + } + + init + } + + #[inline] + #[rustc_inherit_overflow_checks] + fn advance_back_by(&mut self, n: usize) -> Result<(), usize> { + let mut rem = n; + loop { + if let Some(ref mut back) = self.backiter { + match back.advance_back_by(rem) { + ret @ Ok(_) => return ret, + Err(advanced) => rem -= advanced, + } + } + match self.iter.next_back() { + Some(iterable) => self.backiter = Some(iterable.into_iter()), + _ => break, + } + } + + self.backiter = None; + + if let Some(ref mut front) = self.frontiter { + match front.advance_back_by(rem) { + ret @ Ok(_) => return ret, + Err(advanced) => rem -= advanced, + } + } + + if rem > 0 { + return Err(n - rem); + } + + self.frontiter = None; + + Ok(()) + } +} + +trait ConstSizeIntoIterator: IntoIterator { + // FIXME(#31844): convert to an associated const once specialization supports that + fn size() -> Option<usize>; +} + +impl<T> ConstSizeIntoIterator for T +where + T: IntoIterator, +{ + #[inline] + default fn size() -> Option<usize> { + None + } +} + +impl<T, const N: usize> ConstSizeIntoIterator for [T; N] { + #[inline] + fn size() -> Option<usize> { + Some(N) + } +} + +impl<T, const N: usize> ConstSizeIntoIterator for &[T; N] { + #[inline] + fn size() -> Option<usize> { + Some(N) + } +} + +impl<T, const N: usize> ConstSizeIntoIterator for &mut [T; N] { + #[inline] + fn size() -> Option<usize> { + Some(N) + } +} + +#[doc(hidden)] +#[unstable(feature = "std_internals", issue = "none")] +// FIXME(#20400): Instead of this helper trait there should be multiple impl TrustedLen for Flatten<> +// blocks with different bounds on Iterator::Item but the compiler erroneously considers them overlapping +pub unsafe trait TrustedConstSize: IntoIterator {} + +#[unstable(feature = "std_internals", issue = "none")] +unsafe impl<T, const N: usize> TrustedConstSize for [T; N] {} +#[unstable(feature = "std_internals", issue = "none")] +unsafe impl<T, const N: usize> TrustedConstSize for &'_ [T; N] {} +#[unstable(feature = "std_internals", issue = "none")] +unsafe impl<T, const N: usize> TrustedConstSize for &'_ mut [T; N] {} + +#[inline] +fn and_then_or_clear<T, U>(opt: &mut Option<T>, f: impl FnOnce(&mut T) -> Option<U>) -> Option<U> { + let x = f(opt.as_mut()?); + if x.is_none() { + *opt = None; + } + x +} diff --git a/library/core/src/iter/adapters/fuse.rs b/library/core/src/iter/adapters/fuse.rs new file mode 100644 index 000000000..c93144542 --- /dev/null +++ b/library/core/src/iter/adapters/fuse.rs @@ -0,0 +1,413 @@ +use crate::intrinsics; +use crate::iter::adapters::zip::try_get_unchecked; +use crate::iter::{ + DoubleEndedIterator, ExactSizeIterator, FusedIterator, TrustedLen, TrustedRandomAccess, + TrustedRandomAccessNoCoerce, +}; +use crate::ops::Try; + +/// An iterator that yields `None` forever after the underlying iterator +/// yields `None` once. +/// +/// This `struct` is created by [`Iterator::fuse`]. See its documentation +/// for more. +#[derive(Clone, Debug)] +#[must_use = "iterators are lazy and do nothing unless consumed"] +#[stable(feature = "rust1", since = "1.0.0")] +pub struct Fuse<I> { + // NOTE: for `I: FusedIterator`, we never bother setting `None`, but + // we still have to be prepared for that state due to variance. + // See rust-lang/rust#85863 + iter: Option<I>, +} +impl<I> Fuse<I> { + pub(in crate::iter) fn new(iter: I) -> Fuse<I> { + Fuse { iter: Some(iter) } + } +} + +#[stable(feature = "fused", since = "1.26.0")] +impl<I> FusedIterator for Fuse<I> where I: Iterator {} + +// Any specialized implementation here is made internal +// to avoid exposing default fns outside this trait. +#[stable(feature = "rust1", since = "1.0.0")] +impl<I> Iterator for Fuse<I> +where + I: Iterator, +{ + type Item = <I as Iterator>::Item; + + #[inline] + fn next(&mut self) -> Option<Self::Item> { + FuseImpl::next(self) + } + + #[inline] + fn nth(&mut self, n: usize) -> Option<I::Item> { + FuseImpl::nth(self, n) + } + + #[inline] + fn last(self) -> Option<Self::Item> { + match self.iter { + Some(iter) => iter.last(), + None => None, + } + } + + #[inline] + fn count(self) -> usize { + match self.iter { + Some(iter) => iter.count(), + None => 0, + } + } + + #[inline] + fn size_hint(&self) -> (usize, Option<usize>) { + match self.iter { + Some(ref iter) => iter.size_hint(), + None => (0, Some(0)), + } + } + + #[inline] + fn try_fold<Acc, Fold, R>(&mut self, acc: Acc, fold: Fold) -> R + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> R, + R: Try<Output = Acc>, + { + FuseImpl::try_fold(self, acc, fold) + } + + #[inline] + fn fold<Acc, Fold>(self, mut acc: Acc, fold: Fold) -> Acc + where + Fold: FnMut(Acc, Self::Item) -> Acc, + { + if let Some(iter) = self.iter { + acc = iter.fold(acc, fold); + } + acc + } + + #[inline] + fn find<P>(&mut self, predicate: P) -> Option<Self::Item> + where + P: FnMut(&Self::Item) -> bool, + { + FuseImpl::find(self, predicate) + } + + #[inline] + unsafe fn __iterator_get_unchecked(&mut self, idx: usize) -> Self::Item + where + Self: TrustedRandomAccessNoCoerce, + { + match self.iter { + // SAFETY: the caller must uphold the contract for + // `Iterator::__iterator_get_unchecked`. + Some(ref mut iter) => unsafe { try_get_unchecked(iter, idx) }, + // SAFETY: the caller asserts there is an item at `i`, so we're not exhausted. + None => unsafe { intrinsics::unreachable() }, + } + } +} + +#[stable(feature = "rust1", since = "1.0.0")] +impl<I> DoubleEndedIterator for Fuse<I> +where + I: DoubleEndedIterator, +{ + #[inline] + fn next_back(&mut self) -> Option<<I as Iterator>::Item> { + FuseImpl::next_back(self) + } + + #[inline] + fn nth_back(&mut self, n: usize) -> Option<<I as Iterator>::Item> { + FuseImpl::nth_back(self, n) + } + + #[inline] + fn try_rfold<Acc, Fold, R>(&mut self, acc: Acc, fold: Fold) -> R + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> R, + R: Try<Output = Acc>, + { + FuseImpl::try_rfold(self, acc, fold) + } + + #[inline] + fn rfold<Acc, Fold>(self, mut acc: Acc, fold: Fold) -> Acc + where + Fold: FnMut(Acc, Self::Item) -> Acc, + { + if let Some(iter) = self.iter { + acc = iter.rfold(acc, fold); + } + acc + } + + #[inline] + fn rfind<P>(&mut self, predicate: P) -> Option<Self::Item> + where + P: FnMut(&Self::Item) -> bool, + { + FuseImpl::rfind(self, predicate) + } +} + +#[stable(feature = "rust1", since = "1.0.0")] +impl<I> ExactSizeIterator for Fuse<I> +where + I: ExactSizeIterator, +{ + fn len(&self) -> usize { + match self.iter { + Some(ref iter) => iter.len(), + None => 0, + } + } + + fn is_empty(&self) -> bool { + match self.iter { + Some(ref iter) => iter.is_empty(), + None => true, + } + } +} + +#[unstable(feature = "trusted_len", issue = "37572")] +// SAFETY: `TrustedLen` requires that an accurate length is reported via `size_hint()`. As `Fuse` +// is just forwarding this to the wrapped iterator `I` this property is preserved and it is safe to +// implement `TrustedLen` here. +unsafe impl<I> TrustedLen for Fuse<I> where I: TrustedLen {} + +#[doc(hidden)] +#[unstable(feature = "trusted_random_access", issue = "none")] +// SAFETY: `TrustedRandomAccess` requires that `size_hint()` must be exact and cheap to call, and +// `Iterator::__iterator_get_unchecked()` must be implemented accordingly. +// +// This is safe to implement as `Fuse` is just forwarding these to the wrapped iterator `I`, which +// preserves these properties. +unsafe impl<I> TrustedRandomAccess for Fuse<I> where I: TrustedRandomAccess {} + +#[doc(hidden)] +#[unstable(feature = "trusted_random_access", issue = "none")] +unsafe impl<I> TrustedRandomAccessNoCoerce for Fuse<I> +where + I: TrustedRandomAccessNoCoerce, +{ + const MAY_HAVE_SIDE_EFFECT: bool = I::MAY_HAVE_SIDE_EFFECT; +} + +/// Fuse specialization trait +/// +/// We only need to worry about `&mut self` methods, which +/// may exhaust the iterator without consuming it. +#[doc(hidden)] +trait FuseImpl<I> { + type Item; + + // Functions specific to any normal Iterators + fn next(&mut self) -> Option<Self::Item>; + fn nth(&mut self, n: usize) -> Option<Self::Item>; + fn try_fold<Acc, Fold, R>(&mut self, acc: Acc, fold: Fold) -> R + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> R, + R: Try<Output = Acc>; + fn find<P>(&mut self, predicate: P) -> Option<Self::Item> + where + P: FnMut(&Self::Item) -> bool; + + // Functions specific to DoubleEndedIterators + fn next_back(&mut self) -> Option<Self::Item> + where + I: DoubleEndedIterator; + fn nth_back(&mut self, n: usize) -> Option<Self::Item> + where + I: DoubleEndedIterator; + fn try_rfold<Acc, Fold, R>(&mut self, acc: Acc, fold: Fold) -> R + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> R, + R: Try<Output = Acc>, + I: DoubleEndedIterator; + fn rfind<P>(&mut self, predicate: P) -> Option<Self::Item> + where + P: FnMut(&Self::Item) -> bool, + I: DoubleEndedIterator; +} + +/// General `Fuse` impl which sets `iter = None` when exhausted. +#[doc(hidden)] +impl<I> FuseImpl<I> for Fuse<I> +where + I: Iterator, +{ + type Item = <I as Iterator>::Item; + + #[inline] + default fn next(&mut self) -> Option<<I as Iterator>::Item> { + and_then_or_clear(&mut self.iter, Iterator::next) + } + + #[inline] + default fn nth(&mut self, n: usize) -> Option<I::Item> { + and_then_or_clear(&mut self.iter, |iter| iter.nth(n)) + } + + #[inline] + default fn try_fold<Acc, Fold, R>(&mut self, mut acc: Acc, fold: Fold) -> R + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> R, + R: Try<Output = Acc>, + { + if let Some(ref mut iter) = self.iter { + acc = iter.try_fold(acc, fold)?; + self.iter = None; + } + try { acc } + } + + #[inline] + default fn find<P>(&mut self, predicate: P) -> Option<Self::Item> + where + P: FnMut(&Self::Item) -> bool, + { + and_then_or_clear(&mut self.iter, |iter| iter.find(predicate)) + } + + #[inline] + default fn next_back(&mut self) -> Option<<I as Iterator>::Item> + where + I: DoubleEndedIterator, + { + and_then_or_clear(&mut self.iter, |iter| iter.next_back()) + } + + #[inline] + default fn nth_back(&mut self, n: usize) -> Option<<I as Iterator>::Item> + where + I: DoubleEndedIterator, + { + and_then_or_clear(&mut self.iter, |iter| iter.nth_back(n)) + } + + #[inline] + default fn try_rfold<Acc, Fold, R>(&mut self, mut acc: Acc, fold: Fold) -> R + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> R, + R: Try<Output = Acc>, + I: DoubleEndedIterator, + { + if let Some(ref mut iter) = self.iter { + acc = iter.try_rfold(acc, fold)?; + self.iter = None; + } + try { acc } + } + + #[inline] + default fn rfind<P>(&mut self, predicate: P) -> Option<Self::Item> + where + P: FnMut(&Self::Item) -> bool, + I: DoubleEndedIterator, + { + and_then_or_clear(&mut self.iter, |iter| iter.rfind(predicate)) + } +} + +/// Specialized `Fuse` impl which doesn't bother clearing `iter` when exhausted. +/// However, we must still be prepared for the possibility that it was already cleared! +#[doc(hidden)] +impl<I> FuseImpl<I> for Fuse<I> +where + I: FusedIterator, +{ + #[inline] + fn next(&mut self) -> Option<<I as Iterator>::Item> { + self.iter.as_mut()?.next() + } + + #[inline] + fn nth(&mut self, n: usize) -> Option<I::Item> { + self.iter.as_mut()?.nth(n) + } + + #[inline] + fn try_fold<Acc, Fold, R>(&mut self, mut acc: Acc, fold: Fold) -> R + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> R, + R: Try<Output = Acc>, + { + if let Some(ref mut iter) = self.iter { + acc = iter.try_fold(acc, fold)?; + } + try { acc } + } + + #[inline] + fn find<P>(&mut self, predicate: P) -> Option<Self::Item> + where + P: FnMut(&Self::Item) -> bool, + { + self.iter.as_mut()?.find(predicate) + } + + #[inline] + fn next_back(&mut self) -> Option<<I as Iterator>::Item> + where + I: DoubleEndedIterator, + { + self.iter.as_mut()?.next_back() + } + + #[inline] + fn nth_back(&mut self, n: usize) -> Option<<I as Iterator>::Item> + where + I: DoubleEndedIterator, + { + self.iter.as_mut()?.nth_back(n) + } + + #[inline] + fn try_rfold<Acc, Fold, R>(&mut self, mut acc: Acc, fold: Fold) -> R + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> R, + R: Try<Output = Acc>, + I: DoubleEndedIterator, + { + if let Some(ref mut iter) = self.iter { + acc = iter.try_rfold(acc, fold)?; + } + try { acc } + } + + #[inline] + fn rfind<P>(&mut self, predicate: P) -> Option<Self::Item> + where + P: FnMut(&Self::Item) -> bool, + I: DoubleEndedIterator, + { + self.iter.as_mut()?.rfind(predicate) + } +} + +#[inline] +fn and_then_or_clear<T, U>(opt: &mut Option<T>, f: impl FnOnce(&mut T) -> Option<U>) -> Option<U> { + let x = f(opt.as_mut()?); + if x.is_none() { + *opt = None; + } + x +} diff --git a/library/core/src/iter/adapters/inspect.rs b/library/core/src/iter/adapters/inspect.rs new file mode 100644 index 000000000..19839fdfe --- /dev/null +++ b/library/core/src/iter/adapters/inspect.rs @@ -0,0 +1,166 @@ +use crate::fmt; +use crate::iter::{adapters::SourceIter, FusedIterator, InPlaceIterable}; +use crate::ops::Try; + +/// An iterator that calls a function with a reference to each element before +/// yielding it. +/// +/// This `struct` is created by the [`inspect`] method on [`Iterator`]. See its +/// documentation for more. +/// +/// [`inspect`]: Iterator::inspect +/// [`Iterator`]: trait.Iterator.html +#[must_use = "iterators are lazy and do nothing unless consumed"] +#[stable(feature = "rust1", since = "1.0.0")] +#[derive(Clone)] +pub struct Inspect<I, F> { + iter: I, + f: F, +} +impl<I, F> Inspect<I, F> { + pub(in crate::iter) fn new(iter: I, f: F) -> Inspect<I, F> { + Inspect { iter, f } + } +} + +#[stable(feature = "core_impl_debug", since = "1.9.0")] +impl<I: fmt::Debug, F> fmt::Debug for Inspect<I, F> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Inspect").field("iter", &self.iter).finish() + } +} + +impl<I: Iterator, F> Inspect<I, F> +where + F: FnMut(&I::Item), +{ + #[inline] + fn do_inspect(&mut self, elt: Option<I::Item>) -> Option<I::Item> { + if let Some(ref a) = elt { + (self.f)(a); + } + + elt + } +} + +fn inspect_fold<T, Acc>( + mut f: impl FnMut(&T), + mut fold: impl FnMut(Acc, T) -> Acc, +) -> impl FnMut(Acc, T) -> Acc { + move |acc, item| { + f(&item); + fold(acc, item) + } +} + +fn inspect_try_fold<'a, T, Acc, R>( + f: &'a mut impl FnMut(&T), + mut fold: impl FnMut(Acc, T) -> R + 'a, +) -> impl FnMut(Acc, T) -> R + 'a { + move |acc, item| { + f(&item); + fold(acc, item) + } +} + +#[stable(feature = "rust1", since = "1.0.0")] +impl<I: Iterator, F> Iterator for Inspect<I, F> +where + F: FnMut(&I::Item), +{ + type Item = I::Item; + + #[inline] + fn next(&mut self) -> Option<I::Item> { + let next = self.iter.next(); + self.do_inspect(next) + } + + #[inline] + fn size_hint(&self) -> (usize, Option<usize>) { + self.iter.size_hint() + } + + #[inline] + fn try_fold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> R, + R: Try<Output = Acc>, + { + self.iter.try_fold(init, inspect_try_fold(&mut self.f, fold)) + } + + #[inline] + fn fold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc + where + Fold: FnMut(Acc, Self::Item) -> Acc, + { + self.iter.fold(init, inspect_fold(self.f, fold)) + } +} + +#[stable(feature = "rust1", since = "1.0.0")] +impl<I: DoubleEndedIterator, F> DoubleEndedIterator for Inspect<I, F> +where + F: FnMut(&I::Item), +{ + #[inline] + fn next_back(&mut self) -> Option<I::Item> { + let next = self.iter.next_back(); + self.do_inspect(next) + } + + #[inline] + fn try_rfold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> R, + R: Try<Output = Acc>, + { + self.iter.try_rfold(init, inspect_try_fold(&mut self.f, fold)) + } + + #[inline] + fn rfold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc + where + Fold: FnMut(Acc, Self::Item) -> Acc, + { + self.iter.rfold(init, inspect_fold(self.f, fold)) + } +} + +#[stable(feature = "rust1", since = "1.0.0")] +impl<I: ExactSizeIterator, F> ExactSizeIterator for Inspect<I, F> +where + F: FnMut(&I::Item), +{ + fn len(&self) -> usize { + self.iter.len() + } + + fn is_empty(&self) -> bool { + self.iter.is_empty() + } +} + +#[stable(feature = "fused", since = "1.26.0")] +impl<I: FusedIterator, F> FusedIterator for Inspect<I, F> where F: FnMut(&I::Item) {} + +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl<I, F> SourceIter for Inspect<I, F> +where + I: SourceIter, +{ + type Source = I::Source; + + #[inline] + unsafe fn as_inner(&mut self) -> &mut I::Source { + // SAFETY: unsafe function forwarding to unsafe function with the same requirements + unsafe { SourceIter::as_inner(&mut self.iter) } + } +} + +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl<I: InPlaceIterable, F> InPlaceIterable for Inspect<I, F> where F: FnMut(&I::Item) {} diff --git a/library/core/src/iter/adapters/intersperse.rs b/library/core/src/iter/adapters/intersperse.rs new file mode 100644 index 000000000..d8bbd424c --- /dev/null +++ b/library/core/src/iter/adapters/intersperse.rs @@ -0,0 +1,187 @@ +use super::Peekable; + +/// An iterator adapter that places a separator between all elements. +/// +/// This `struct` is created by [`Iterator::intersperse`]. See its documentation +/// for more information. +#[unstable(feature = "iter_intersperse", reason = "recently added", issue = "79524")] +#[derive(Debug, Clone)] +pub struct Intersperse<I: Iterator> +where + I::Item: Clone, +{ + separator: I::Item, + iter: Peekable<I>, + needs_sep: bool, +} + +impl<I: Iterator> Intersperse<I> +where + I::Item: Clone, +{ + pub(in crate::iter) fn new(iter: I, separator: I::Item) -> Self { + Self { iter: iter.peekable(), separator, needs_sep: false } + } +} + +#[unstable(feature = "iter_intersperse", reason = "recently added", issue = "79524")] +impl<I> Iterator for Intersperse<I> +where + I: Iterator, + I::Item: Clone, +{ + type Item = I::Item; + + #[inline] + fn next(&mut self) -> Option<I::Item> { + if self.needs_sep && self.iter.peek().is_some() { + self.needs_sep = false; + Some(self.separator.clone()) + } else { + self.needs_sep = true; + self.iter.next() + } + } + + fn fold<B, F>(self, init: B, f: F) -> B + where + Self: Sized, + F: FnMut(B, Self::Item) -> B, + { + let separator = self.separator; + intersperse_fold(self.iter, init, f, move || separator.clone(), self.needs_sep) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + intersperse_size_hint(&self.iter, self.needs_sep) + } +} + +/// An iterator adapter that places a separator between all elements. +/// +/// This `struct` is created by [`Iterator::intersperse_with`]. See its +/// documentation for more information. +#[unstable(feature = "iter_intersperse", reason = "recently added", issue = "79524")] +pub struct IntersperseWith<I, G> +where + I: Iterator, +{ + separator: G, + iter: Peekable<I>, + needs_sep: bool, +} + +#[unstable(feature = "iter_intersperse", reason = "recently added", issue = "79524")] +impl<I, G> crate::fmt::Debug for IntersperseWith<I, G> +where + I: Iterator + crate::fmt::Debug, + I::Item: crate::fmt::Debug, + G: crate::fmt::Debug, +{ + fn fmt(&self, f: &mut crate::fmt::Formatter<'_>) -> crate::fmt::Result { + f.debug_struct("IntersperseWith") + .field("separator", &self.separator) + .field("iter", &self.iter) + .field("needs_sep", &self.needs_sep) + .finish() + } +} + +#[unstable(feature = "iter_intersperse", reason = "recently added", issue = "79524")] +impl<I, G> crate::clone::Clone for IntersperseWith<I, G> +where + I: Iterator + crate::clone::Clone, + I::Item: crate::clone::Clone, + G: Clone, +{ + fn clone(&self) -> Self { + IntersperseWith { + separator: self.separator.clone(), + iter: self.iter.clone(), + needs_sep: self.needs_sep.clone(), + } + } +} + +impl<I, G> IntersperseWith<I, G> +where + I: Iterator, + G: FnMut() -> I::Item, +{ + pub(in crate::iter) fn new(iter: I, separator: G) -> Self { + Self { iter: iter.peekable(), separator, needs_sep: false } + } +} + +#[unstable(feature = "iter_intersperse", reason = "recently added", issue = "79524")] +impl<I, G> Iterator for IntersperseWith<I, G> +where + I: Iterator, + G: FnMut() -> I::Item, +{ + type Item = I::Item; + + #[inline] + fn next(&mut self) -> Option<I::Item> { + if self.needs_sep && self.iter.peek().is_some() { + self.needs_sep = false; + Some((self.separator)()) + } else { + self.needs_sep = true; + self.iter.next() + } + } + + fn fold<B, F>(self, init: B, f: F) -> B + where + Self: Sized, + F: FnMut(B, Self::Item) -> B, + { + intersperse_fold(self.iter, init, f, self.separator, self.needs_sep) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + intersperse_size_hint(&self.iter, self.needs_sep) + } +} + +fn intersperse_size_hint<I>(iter: &I, needs_sep: bool) -> (usize, Option<usize>) +where + I: Iterator, +{ + let (lo, hi) = iter.size_hint(); + let next_is_elem = !needs_sep; + ( + lo.saturating_sub(next_is_elem as usize).saturating_add(lo), + hi.and_then(|hi| hi.saturating_sub(next_is_elem as usize).checked_add(hi)), + ) +} + +fn intersperse_fold<I, B, F, G>( + mut iter: I, + init: B, + mut f: F, + mut separator: G, + needs_sep: bool, +) -> B +where + I: Iterator, + F: FnMut(B, I::Item) -> B, + G: FnMut() -> I::Item, +{ + let mut accum = init; + + if !needs_sep { + if let Some(x) = iter.next() { + accum = f(accum, x); + } else { + return accum; + } + } + + iter.fold(accum, |mut accum, x| { + accum = f(accum, separator()); + accum = f(accum, x); + accum + }) +} diff --git a/library/core/src/iter/adapters/map.rs b/library/core/src/iter/adapters/map.rs new file mode 100644 index 000000000..9e25dbe46 --- /dev/null +++ b/library/core/src/iter/adapters/map.rs @@ -0,0 +1,218 @@ +use crate::fmt; +use crate::iter::adapters::{ + zip::try_get_unchecked, SourceIter, TrustedRandomAccess, TrustedRandomAccessNoCoerce, +}; +use crate::iter::{FusedIterator, InPlaceIterable, TrustedLen}; +use crate::ops::Try; + +/// An iterator that maps the values of `iter` with `f`. +/// +/// This `struct` is created by the [`map`] method on [`Iterator`]. See its +/// documentation for more. +/// +/// [`map`]: Iterator::map +/// [`Iterator`]: trait.Iterator.html +/// +/// # Notes about side effects +/// +/// The [`map`] iterator implements [`DoubleEndedIterator`], meaning that +/// you can also [`map`] backwards: +/// +/// ```rust +/// let v: Vec<i32> = [1, 2, 3].into_iter().map(|x| x + 1).rev().collect(); +/// +/// assert_eq!(v, [4, 3, 2]); +/// ``` +/// +/// [`DoubleEndedIterator`]: trait.DoubleEndedIterator.html +/// +/// But if your closure has state, iterating backwards may act in a way you do +/// not expect. Let's go through an example. First, in the forward direction: +/// +/// ```rust +/// let mut c = 0; +/// +/// for pair in ['a', 'b', 'c'].into_iter() +/// .map(|letter| { c += 1; (letter, c) }) { +/// println!("{pair:?}"); +/// } +/// ``` +/// +/// This will print `('a', 1), ('b', 2), ('c', 3)`. +/// +/// Now consider this twist where we add a call to `rev`. This version will +/// print `('c', 1), ('b', 2), ('a', 3)`. Note that the letters are reversed, +/// but the values of the counter still go in order. This is because `map()` is +/// still being called lazily on each item, but we are popping items off the +/// back of the vector now, instead of shifting them from the front. +/// +/// ```rust +/// let mut c = 0; +/// +/// for pair in ['a', 'b', 'c'].into_iter() +/// .map(|letter| { c += 1; (letter, c) }) +/// .rev() { +/// println!("{pair:?}"); +/// } +/// ``` +#[must_use = "iterators are lazy and do nothing unless consumed"] +#[stable(feature = "rust1", since = "1.0.0")] +#[derive(Clone)] +pub struct Map<I, F> { + // Used for `SplitWhitespace` and `SplitAsciiWhitespace` `as_str` methods + pub(crate) iter: I, + f: F, +} + +impl<I, F> Map<I, F> { + pub(in crate::iter) fn new(iter: I, f: F) -> Map<I, F> { + Map { iter, f } + } +} + +#[stable(feature = "core_impl_debug", since = "1.9.0")] +impl<I: fmt::Debug, F> fmt::Debug for Map<I, F> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Map").field("iter", &self.iter).finish() + } +} + +fn map_fold<T, B, Acc>( + mut f: impl FnMut(T) -> B, + mut g: impl FnMut(Acc, B) -> Acc, +) -> impl FnMut(Acc, T) -> Acc { + move |acc, elt| g(acc, f(elt)) +} + +fn map_try_fold<'a, T, B, Acc, R>( + f: &'a mut impl FnMut(T) -> B, + mut g: impl FnMut(Acc, B) -> R + 'a, +) -> impl FnMut(Acc, T) -> R + 'a { + move |acc, elt| g(acc, f(elt)) +} + +#[stable(feature = "rust1", since = "1.0.0")] +impl<B, I: Iterator, F> Iterator for Map<I, F> +where + F: FnMut(I::Item) -> B, +{ + type Item = B; + + #[inline] + fn next(&mut self) -> Option<B> { + self.iter.next().map(&mut self.f) + } + + #[inline] + fn size_hint(&self) -> (usize, Option<usize>) { + self.iter.size_hint() + } + + fn try_fold<Acc, G, R>(&mut self, init: Acc, g: G) -> R + where + Self: Sized, + G: FnMut(Acc, Self::Item) -> R, + R: Try<Output = Acc>, + { + self.iter.try_fold(init, map_try_fold(&mut self.f, g)) + } + + fn fold<Acc, G>(self, init: Acc, g: G) -> Acc + where + G: FnMut(Acc, Self::Item) -> Acc, + { + self.iter.fold(init, map_fold(self.f, g)) + } + + #[inline] + unsafe fn __iterator_get_unchecked(&mut self, idx: usize) -> B + where + Self: TrustedRandomAccessNoCoerce, + { + // SAFETY: the caller must uphold the contract for + // `Iterator::__iterator_get_unchecked`. + unsafe { (self.f)(try_get_unchecked(&mut self.iter, idx)) } + } +} + +#[stable(feature = "rust1", since = "1.0.0")] +impl<B, I: DoubleEndedIterator, F> DoubleEndedIterator for Map<I, F> +where + F: FnMut(I::Item) -> B, +{ + #[inline] + fn next_back(&mut self) -> Option<B> { + self.iter.next_back().map(&mut self.f) + } + + fn try_rfold<Acc, G, R>(&mut self, init: Acc, g: G) -> R + where + Self: Sized, + G: FnMut(Acc, Self::Item) -> R, + R: Try<Output = Acc>, + { + self.iter.try_rfold(init, map_try_fold(&mut self.f, g)) + } + + fn rfold<Acc, G>(self, init: Acc, g: G) -> Acc + where + G: FnMut(Acc, Self::Item) -> Acc, + { + self.iter.rfold(init, map_fold(self.f, g)) + } +} + +#[stable(feature = "rust1", since = "1.0.0")] +impl<B, I: ExactSizeIterator, F> ExactSizeIterator for Map<I, F> +where + F: FnMut(I::Item) -> B, +{ + fn len(&self) -> usize { + self.iter.len() + } + + fn is_empty(&self) -> bool { + self.iter.is_empty() + } +} + +#[stable(feature = "fused", since = "1.26.0")] +impl<B, I: FusedIterator, F> FusedIterator for Map<I, F> where F: FnMut(I::Item) -> B {} + +#[unstable(feature = "trusted_len", issue = "37572")] +unsafe impl<B, I, F> TrustedLen for Map<I, F> +where + I: TrustedLen, + F: FnMut(I::Item) -> B, +{ +} + +#[doc(hidden)] +#[unstable(feature = "trusted_random_access", issue = "none")] +unsafe impl<I, F> TrustedRandomAccess for Map<I, F> where I: TrustedRandomAccess {} + +#[doc(hidden)] +#[unstable(feature = "trusted_random_access", issue = "none")] +unsafe impl<I, F> TrustedRandomAccessNoCoerce for Map<I, F> +where + I: TrustedRandomAccessNoCoerce, +{ + const MAY_HAVE_SIDE_EFFECT: bool = true; +} + +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl<I, F> SourceIter for Map<I, F> +where + I: SourceIter, +{ + type Source = I::Source; + + #[inline] + unsafe fn as_inner(&mut self) -> &mut I::Source { + // SAFETY: unsafe function forwarding to unsafe function with the same requirements + unsafe { SourceIter::as_inner(&mut self.iter) } + } +} + +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl<B, I: InPlaceIterable, F> InPlaceIterable for Map<I, F> where F: FnMut(I::Item) -> B {} diff --git a/library/core/src/iter/adapters/map_while.rs b/library/core/src/iter/adapters/map_while.rs new file mode 100644 index 000000000..1e8d6bf3e --- /dev/null +++ b/library/core/src/iter/adapters/map_while.rs @@ -0,0 +1,100 @@ +use crate::fmt; +use crate::iter::{adapters::SourceIter, InPlaceIterable}; +use crate::ops::{ControlFlow, Try}; + +/// An iterator that only accepts elements while `predicate` returns `Some(_)`. +/// +/// This `struct` is created by the [`map_while`] method on [`Iterator`]. See its +/// documentation for more. +/// +/// [`map_while`]: Iterator::map_while +/// [`Iterator`]: trait.Iterator.html +#[must_use = "iterators are lazy and do nothing unless consumed"] +#[stable(feature = "iter_map_while", since = "1.57.0")] +#[derive(Clone)] +pub struct MapWhile<I, P> { + iter: I, + predicate: P, +} + +impl<I, P> MapWhile<I, P> { + pub(in crate::iter) fn new(iter: I, predicate: P) -> MapWhile<I, P> { + MapWhile { iter, predicate } + } +} + +#[stable(feature = "iter_map_while", since = "1.57.0")] +impl<I: fmt::Debug, P> fmt::Debug for MapWhile<I, P> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("MapWhile").field("iter", &self.iter).finish() + } +} + +#[stable(feature = "iter_map_while", since = "1.57.0")] +impl<B, I: Iterator, P> Iterator for MapWhile<I, P> +where + P: FnMut(I::Item) -> Option<B>, +{ + type Item = B; + + #[inline] + fn next(&mut self) -> Option<B> { + let x = self.iter.next()?; + (self.predicate)(x) + } + + #[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 + } + + #[inline] + fn try_fold<Acc, Fold, R>(&mut self, init: Acc, mut fold: Fold) -> R + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> R, + R: Try<Output = Acc>, + { + let Self { iter, predicate } = self; + iter.try_fold(init, |acc, x| match predicate(x) { + Some(item) => ControlFlow::from_try(fold(acc, item)), + None => ControlFlow::Break(try { acc }), + }) + .into_try() + } + + #[inline] + fn fold<Acc, Fold>(mut self, init: Acc, fold: Fold) -> Acc + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> Acc, + { + #[inline] + fn ok<B, T>(mut f: impl FnMut(B, T) -> B) -> impl FnMut(B, T) -> Result<B, !> { + move |acc, x| Ok(f(acc, x)) + } + + self.try_fold(init, ok(fold)).unwrap() + } +} + +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl<I, P> SourceIter for MapWhile<I, P> +where + I: SourceIter, +{ + type Source = I::Source; + + #[inline] + unsafe fn as_inner(&mut self) -> &mut I::Source { + // SAFETY: unsafe function forwarding to unsafe function with the same requirements + unsafe { SourceIter::as_inner(&mut self.iter) } + } +} + +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl<B, I: InPlaceIterable, P> InPlaceIterable for MapWhile<I, P> where + P: FnMut(I::Item) -> Option<B> +{ +} diff --git a/library/core/src/iter/adapters/mod.rs b/library/core/src/iter/adapters/mod.rs new file mode 100644 index 000000000..916a26e24 --- /dev/null +++ b/library/core/src/iter/adapters/mod.rs @@ -0,0 +1,232 @@ +use crate::iter::{InPlaceIterable, Iterator}; +use crate::ops::{ChangeOutputType, ControlFlow, FromResidual, NeverShortCircuit, Residual, Try}; + +mod by_ref_sized; +mod chain; +mod cloned; +mod copied; +mod cycle; +mod enumerate; +mod filter; +mod filter_map; +mod flatten; +mod fuse; +mod inspect; +mod intersperse; +mod map; +mod map_while; +mod peekable; +mod rev; +mod scan; +mod skip; +mod skip_while; +mod step_by; +mod take; +mod take_while; +mod zip; + +#[stable(feature = "rust1", since = "1.0.0")] +pub use self::{ + chain::Chain, cycle::Cycle, enumerate::Enumerate, filter::Filter, filter_map::FilterMap, + flatten::FlatMap, fuse::Fuse, inspect::Inspect, map::Map, peekable::Peekable, rev::Rev, + scan::Scan, skip::Skip, skip_while::SkipWhile, take::Take, take_while::TakeWhile, zip::Zip, +}; + +#[unstable(feature = "std_internals", issue = "none")] +pub use self::by_ref_sized::ByRefSized; + +#[stable(feature = "iter_cloned", since = "1.1.0")] +pub use self::cloned::Cloned; + +#[stable(feature = "iterator_step_by", since = "1.28.0")] +pub use self::step_by::StepBy; + +#[stable(feature = "iterator_flatten", since = "1.29.0")] +pub use self::flatten::Flatten; + +#[stable(feature = "iter_copied", since = "1.36.0")] +pub use self::copied::Copied; + +#[unstable(feature = "iter_intersperse", reason = "recently added", issue = "79524")] +pub use self::intersperse::{Intersperse, IntersperseWith}; + +#[stable(feature = "iter_map_while", since = "1.57.0")] +pub use self::map_while::MapWhile; + +#[unstable(feature = "trusted_random_access", issue = "none")] +pub use self::zip::TrustedRandomAccess; + +#[unstable(feature = "trusted_random_access", issue = "none")] +pub use self::zip::TrustedRandomAccessNoCoerce; + +#[stable(feature = "iter_zip", since = "1.59.0")] +pub use self::zip::zip; + +/// 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>` +/// * there is a delegating implementation of this trait for each adapter in the pipeline between +/// the source and the pipeline consumer. +/// +/// When the source is an owning iterator struct (commonly called `IntoIter`) then +/// this can be useful for specializing [`FromIterator`] implementations or recovering the +/// remaining elements after an iterator has been partially exhausted. +/// +/// Note that implementations do not necessarily have to provide access to the inner-most +/// source of a pipeline. A stateful intermediate adapter might eagerly evaluate a part +/// of the pipeline and expose its internal storage as source. +/// +/// The trait is unsafe because implementers must uphold additional safety properties. +/// See [`as_inner`] for details. +/// +/// The primary use of this trait is in-place iteration. Refer to the [`vec::in_place_collect`] +/// module documentation for more information. +/// +/// [`vec::in_place_collect`]: ../../../../alloc/vec/in_place_collect/index.html +/// +/// # Examples +/// +/// Retrieving a partially consumed source: +/// +/// ``` +/// # #![feature(inplace_iteration)] +/// # use std::iter::SourceIter; +/// +/// let mut iter = vec![9, 9, 9].into_iter().map(|i| i * i); +/// let _ = iter.next(); +/// let mut remainder = std::mem::replace(unsafe { iter.as_inner() }, Vec::new().into_iter()); +/// println!("n = {} elements remaining", remainder.len()); +/// ``` +/// +/// [`FromIterator`]: crate::iter::FromIterator +/// [`as_inner`]: SourceIter::as_inner +#[unstable(issue = "none", feature = "inplace_iteration")] +#[doc(hidden)] +#[rustc_specialization_trait] +pub unsafe trait SourceIter { + /// A source stage in an iterator pipeline. + type Source; + + /// Retrieve the source of an iterator pipeline. + /// + /// # Safety + /// + /// Implementations of must return the same mutable reference for their lifetime, unless + /// replaced by a caller. + /// Callers may only replace the reference when they stopped iteration and drop the + /// iterator pipeline after extracting the source. + /// + /// This means iterator adapters can rely on the source not changing during + /// iteration but they cannot rely on it in their Drop implementations. + /// + /// Implementing this method means adapters relinquish private-only access to their + /// source and can only rely on guarantees made based on method receiver types. + /// The lack of restricted access also requires that adapters must uphold the source's + /// public API even when they have access to its internals. + /// + /// Callers in turn must expect the source to be in any state that is consistent with + /// its public API since adapters sitting between it and the source have the same + /// access. In particular an adapter may have consumed more elements than strictly necessary. + /// + /// The overall goal of these requirements is to let the consumer of a pipeline use + /// * whatever remains in the source after iteration has stopped + /// * the memory that has become unused by advancing a consuming iterator + /// + /// [`next()`]: Iterator::next() + unsafe fn as_inner(&mut self) -> &mut Self::Source; +} + +/// An iterator adapter that produces output as long as the underlying +/// iterator produces values where `Try::branch` says to `ControlFlow::Continue`. +/// +/// If a `ControlFlow::Break` is encountered, the iterator stops and the +/// residual is stored. +pub(crate) struct GenericShunt<'a, I, R> { + iter: I, + residual: &'a mut Option<R>, +} + +/// Process the given iterator as if it yielded a the item's `Try::Output` +/// type instead. Any `Try::Residual`s encountered will stop the inner iterator +/// and be propagated back to the overall result. +pub(crate) fn try_process<I, T, R, F, U>(iter: I, mut f: F) -> ChangeOutputType<I::Item, U> +where + I: Iterator<Item: Try<Output = T, Residual = R>>, + for<'a> F: FnMut(GenericShunt<'a, I, R>) -> U, + R: Residual<U>, +{ + let mut residual = None; + let shunt = GenericShunt { iter, residual: &mut residual }; + let value = f(shunt); + match residual { + Some(r) => FromResidual::from_residual(r), + None => Try::from_output(value), + } +} + +impl<I, R> Iterator for GenericShunt<'_, I, R> +where + I: Iterator<Item: Try<Residual = R>>, +{ + type Item = <I::Item as Try>::Output; + + fn next(&mut self) -> Option<Self::Item> { + self.try_for_each(ControlFlow::Break).break_value() + } + + fn size_hint(&self) -> (usize, Option<usize>) { + if self.residual.is_some() { + (0, Some(0)) + } else { + let (_, upper) = self.iter.size_hint(); + (0, upper) + } + } + + fn try_fold<B, F, T>(&mut self, init: B, mut f: F) -> T + where + F: FnMut(B, Self::Item) -> T, + T: Try<Output = B>, + { + self.iter + .try_fold(init, |acc, x| match Try::branch(x) { + ControlFlow::Continue(x) => ControlFlow::from_try(f(acc, x)), + ControlFlow::Break(r) => { + *self.residual = Some(r); + ControlFlow::Break(try { acc }) + } + }) + .into_try() + } + + fn fold<B, F>(mut self, init: B, fold: F) -> B + where + Self: Sized, + F: FnMut(B, Self::Item) -> B, + { + self.try_fold(init, NeverShortCircuit::wrap_mut_2(fold)).0 + } +} + +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl<I, R> SourceIter for GenericShunt<'_, I, R> +where + I: SourceIter, +{ + type Source = I::Source; + + #[inline] + unsafe fn as_inner(&mut self) -> &mut Self::Source { + // SAFETY: unsafe function forwarding to unsafe function with the same requirements + unsafe { SourceIter::as_inner(&mut self.iter) } + } +} + +// SAFETY: GenericShunt::next calls `I::try_for_each`, which has to advance `iter` +// in order to return `Some(_)`. Since `iter` has type `I: InPlaceIterable` it's +// guaranteed that at least one item will be moved out from the underlying source. +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl<I, T, R> InPlaceIterable for GenericShunt<'_, I, R> where + I: Iterator<Item: Try<Output = T, Residual = R>> + InPlaceIterable +{ +} diff --git a/library/core/src/iter/adapters/peekable.rs b/library/core/src/iter/adapters/peekable.rs new file mode 100644 index 000000000..20aca323b --- /dev/null +++ b/library/core/src/iter/adapters/peekable.rs @@ -0,0 +1,335 @@ +use crate::iter::{adapters::SourceIter, FusedIterator, TrustedLen}; +use crate::ops::{ControlFlow, Try}; + +/// An iterator with a `peek()` that returns an optional reference to the next +/// element. +/// +/// This `struct` is created by the [`peekable`] method on [`Iterator`]. See its +/// documentation for more. +/// +/// [`peekable`]: Iterator::peekable +/// [`Iterator`]: trait.Iterator.html +#[derive(Clone, Debug)] +#[must_use = "iterators are lazy and do nothing unless consumed"] +#[stable(feature = "rust1", since = "1.0.0")] +pub struct Peekable<I: Iterator> { + iter: I, + /// Remember a peeked value, even if it was None. + peeked: Option<Option<I::Item>>, +} + +impl<I: Iterator> Peekable<I> { + pub(in crate::iter) fn new(iter: I) -> Peekable<I> { + Peekable { iter, peeked: None } + } +} + +// Peekable must remember if a None has been seen in the `.peek()` method. +// It ensures that `.peek(); .peek();` or `.peek(); .next();` only advances the +// underlying iterator at most once. This does not by itself make the iterator +// fused. +#[stable(feature = "rust1", since = "1.0.0")] +impl<I: Iterator> Iterator for Peekable<I> { + type Item = I::Item; + + #[inline] + fn next(&mut self) -> Option<I::Item> { + match self.peeked.take() { + Some(v) => v, + None => self.iter.next(), + } + } + + #[inline] + #[rustc_inherit_overflow_checks] + fn count(mut self) -> usize { + match self.peeked.take() { + Some(None) => 0, + Some(Some(_)) => 1 + self.iter.count(), + None => self.iter.count(), + } + } + + #[inline] + fn nth(&mut self, n: usize) -> Option<I::Item> { + match self.peeked.take() { + Some(None) => None, + Some(v @ Some(_)) if n == 0 => v, + Some(Some(_)) => self.iter.nth(n - 1), + None => self.iter.nth(n), + } + } + + #[inline] + fn last(mut self) -> Option<I::Item> { + let peek_opt = match self.peeked.take() { + Some(None) => return None, + Some(v) => v, + None => None, + }; + self.iter.last().or(peek_opt) + } + + #[inline] + fn size_hint(&self) -> (usize, Option<usize>) { + let peek_len = match self.peeked { + Some(None) => return (0, Some(0)), + Some(Some(_)) => 1, + None => 0, + }; + let (lo, hi) = self.iter.size_hint(); + let lo = lo.saturating_add(peek_len); + let hi = match hi { + Some(x) => x.checked_add(peek_len), + None => None, + }; + (lo, hi) + } + + #[inline] + fn try_fold<B, F, R>(&mut self, init: B, mut f: F) -> R + where + Self: Sized, + F: FnMut(B, Self::Item) -> R, + R: Try<Output = B>, + { + let acc = match self.peeked.take() { + Some(None) => return try { init }, + Some(Some(v)) => f(init, v)?, + None => init, + }; + self.iter.try_fold(acc, f) + } + + #[inline] + fn fold<Acc, Fold>(self, init: Acc, mut fold: Fold) -> Acc + where + Fold: FnMut(Acc, Self::Item) -> Acc, + { + let acc = match self.peeked { + Some(None) => return init, + Some(Some(v)) => fold(init, v), + None => init, + }; + self.iter.fold(acc, fold) + } +} + +#[stable(feature = "double_ended_peek_iterator", since = "1.38.0")] +impl<I> DoubleEndedIterator for Peekable<I> +where + I: DoubleEndedIterator, +{ + #[inline] + fn next_back(&mut self) -> Option<Self::Item> { + match self.peeked.as_mut() { + Some(v @ Some(_)) => self.iter.next_back().or_else(|| v.take()), + Some(None) => None, + None => self.iter.next_back(), + } + } + + #[inline] + fn try_rfold<B, F, R>(&mut self, init: B, mut f: F) -> R + where + Self: Sized, + F: FnMut(B, Self::Item) -> R, + R: Try<Output = B>, + { + match self.peeked.take() { + Some(None) => try { init }, + Some(Some(v)) => match self.iter.try_rfold(init, &mut f).branch() { + ControlFlow::Continue(acc) => f(acc, v), + ControlFlow::Break(r) => { + self.peeked = Some(Some(v)); + R::from_residual(r) + } + }, + None => self.iter.try_rfold(init, f), + } + } + + #[inline] + fn rfold<Acc, Fold>(self, init: Acc, mut fold: Fold) -> Acc + where + Fold: FnMut(Acc, Self::Item) -> Acc, + { + match self.peeked { + Some(None) => init, + Some(Some(v)) => { + let acc = self.iter.rfold(init, &mut fold); + fold(acc, v) + } + None => self.iter.rfold(init, fold), + } + } +} + +#[stable(feature = "rust1", since = "1.0.0")] +impl<I: ExactSizeIterator> ExactSizeIterator for Peekable<I> {} + +#[stable(feature = "fused", since = "1.26.0")] +impl<I: FusedIterator> FusedIterator for Peekable<I> {} + +impl<I: Iterator> Peekable<I> { + /// Returns a reference to the next() value without advancing the iterator. + /// + /// Like [`next`], if there is a value, it is wrapped in a `Some(T)`. + /// But if the iteration is over, `None` is returned. + /// + /// [`next`]: Iterator::next + /// + /// Because `peek()` returns a reference, and many iterators iterate over + /// references, there can be a possibly confusing situation where the + /// return value is a double reference. You can see this effect in the + /// examples below. + /// + /// # Examples + /// + /// Basic usage: + /// + /// ``` + /// let xs = [1, 2, 3]; + /// + /// let mut iter = xs.iter().peekable(); + /// + /// // peek() lets us see into the future + /// assert_eq!(iter.peek(), Some(&&1)); + /// assert_eq!(iter.next(), Some(&1)); + /// + /// assert_eq!(iter.next(), Some(&2)); + /// + /// // The iterator does not advance even if we `peek` multiple times + /// assert_eq!(iter.peek(), Some(&&3)); + /// assert_eq!(iter.peek(), Some(&&3)); + /// + /// assert_eq!(iter.next(), Some(&3)); + /// + /// // After the iterator is finished, so is `peek()` + /// assert_eq!(iter.peek(), None); + /// assert_eq!(iter.next(), None); + /// ``` + #[inline] + #[stable(feature = "rust1", since = "1.0.0")] + pub fn peek(&mut self) -> Option<&I::Item> { + let iter = &mut self.iter; + self.peeked.get_or_insert_with(|| iter.next()).as_ref() + } + + /// Returns a mutable reference to the next() value without advancing the iterator. + /// + /// Like [`next`], if there is a value, it is wrapped in a `Some(T)`. + /// But if the iteration is over, `None` is returned. + /// + /// Because `peek_mut()` returns a reference, and many iterators iterate over + /// references, there can be a possibly confusing situation where the + /// return value is a double reference. You can see this effect in the examples + /// below. + /// + /// [`next`]: Iterator::next + /// + /// # Examples + /// + /// Basic usage: + /// + /// ``` + /// let mut iter = [1, 2, 3].iter().peekable(); + /// + /// // Like with `peek()`, we can see into the future without advancing the iterator. + /// assert_eq!(iter.peek_mut(), Some(&mut &1)); + /// assert_eq!(iter.peek_mut(), Some(&mut &1)); + /// assert_eq!(iter.next(), Some(&1)); + /// + /// // Peek into the iterator and set the value behind the mutable reference. + /// if let Some(p) = iter.peek_mut() { + /// assert_eq!(*p, &2); + /// *p = &5; + /// } + /// + /// // The value we put in reappears as the iterator continues. + /// assert_eq!(iter.collect::<Vec<_>>(), vec![&5, &3]); + /// ``` + #[inline] + #[stable(feature = "peekable_peek_mut", since = "1.53.0")] + pub fn peek_mut(&mut self) -> Option<&mut I::Item> { + let iter = &mut self.iter; + self.peeked.get_or_insert_with(|| iter.next()).as_mut() + } + + /// Consume and return the next value of this iterator if a condition is true. + /// + /// If `func` returns `true` for the next value of this iterator, consume and return it. + /// Otherwise, return `None`. + /// + /// # Examples + /// Consume a number if it's equal to 0. + /// ``` + /// let mut iter = (0..5).peekable(); + /// // The first item of the iterator is 0; consume it. + /// assert_eq!(iter.next_if(|&x| x == 0), Some(0)); + /// // The next item returned is now 1, so `consume` will return `false`. + /// assert_eq!(iter.next_if(|&x| x == 0), None); + /// // `next_if` saves the value of the next item if it was not equal to `expected`. + /// assert_eq!(iter.next(), Some(1)); + /// ``` + /// + /// Consume any number less than 10. + /// ``` + /// let mut iter = (1..20).peekable(); + /// // Consume all numbers less than 10 + /// while iter.next_if(|&x| x < 10).is_some() {} + /// // The next value returned will be 10 + /// assert_eq!(iter.next(), Some(10)); + /// ``` + #[stable(feature = "peekable_next_if", since = "1.51.0")] + pub fn next_if(&mut self, func: impl FnOnce(&I::Item) -> bool) -> Option<I::Item> { + match self.next() { + Some(matched) if func(&matched) => Some(matched), + other => { + // Since we called `self.next()`, we consumed `self.peeked`. + assert!(self.peeked.is_none()); + self.peeked = Some(other); + None + } + } + } + + /// Consume and return the next item if it is equal to `expected`. + /// + /// # Example + /// Consume a number if it's equal to 0. + /// ``` + /// let mut iter = (0..5).peekable(); + /// // The first item of the iterator is 0; consume it. + /// assert_eq!(iter.next_if_eq(&0), Some(0)); + /// // The next item returned is now 1, so `consume` will return `false`. + /// assert_eq!(iter.next_if_eq(&0), None); + /// // `next_if_eq` saves the value of the next item if it was not equal to `expected`. + /// assert_eq!(iter.next(), Some(1)); + /// ``` + #[stable(feature = "peekable_next_if", since = "1.51.0")] + pub fn next_if_eq<T>(&mut self, expected: &T) -> Option<I::Item> + where + T: ?Sized, + I::Item: PartialEq<T>, + { + self.next_if(|next| next == expected) + } +} + +#[unstable(feature = "trusted_len", issue = "37572")] +unsafe impl<I> TrustedLen for Peekable<I> where I: TrustedLen {} + +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl<I: Iterator> SourceIter for Peekable<I> +where + I: SourceIter, +{ + type Source = I::Source; + + #[inline] + unsafe fn as_inner(&mut self) -> &mut I::Source { + // SAFETY: unsafe function forwarding to unsafe function with the same requirements + unsafe { SourceIter::as_inner(&mut self.iter) } + } +} diff --git a/library/core/src/iter/adapters/rev.rs b/library/core/src/iter/adapters/rev.rs new file mode 100644 index 000000000..139fb7bbd --- /dev/null +++ b/library/core/src/iter/adapters/rev.rs @@ -0,0 +1,137 @@ +use crate::iter::{FusedIterator, TrustedLen}; +use crate::ops::Try; + +/// A double-ended iterator with the direction inverted. +/// +/// This `struct` is created by the [`rev`] method on [`Iterator`]. See its +/// documentation for more. +/// +/// [`rev`]: Iterator::rev +/// [`Iterator`]: trait.Iterator.html +#[derive(Clone, Debug)] +#[must_use = "iterators are lazy and do nothing unless consumed"] +#[stable(feature = "rust1", since = "1.0.0")] +pub struct Rev<T> { + iter: T, +} + +impl<T> Rev<T> { + pub(in crate::iter) fn new(iter: T) -> Rev<T> { + Rev { iter } + } +} + +#[stable(feature = "rust1", since = "1.0.0")] +impl<I> Iterator for Rev<I> +where + I: DoubleEndedIterator, +{ + type Item = <I as Iterator>::Item; + + #[inline] + fn next(&mut self) -> Option<<I as Iterator>::Item> { + self.iter.next_back() + } + #[inline] + fn size_hint(&self) -> (usize, Option<usize>) { + self.iter.size_hint() + } + + #[inline] + fn advance_by(&mut self, n: usize) -> Result<(), usize> { + self.iter.advance_back_by(n) + } + + #[inline] + fn nth(&mut self, n: usize) -> Option<<I as Iterator>::Item> { + self.iter.nth_back(n) + } + + fn try_fold<B, F, R>(&mut self, init: B, f: F) -> R + where + Self: Sized, + F: FnMut(B, Self::Item) -> R, + R: Try<Output = B>, + { + self.iter.try_rfold(init, f) + } + + fn fold<Acc, F>(self, init: Acc, f: F) -> Acc + where + F: FnMut(Acc, Self::Item) -> Acc, + { + self.iter.rfold(init, f) + } + + #[inline] + fn find<P>(&mut self, predicate: P) -> Option<Self::Item> + where + P: FnMut(&Self::Item) -> bool, + { + self.iter.rfind(predicate) + } +} + +#[stable(feature = "rust1", since = "1.0.0")] +impl<I> DoubleEndedIterator for Rev<I> +where + I: DoubleEndedIterator, +{ + #[inline] + fn next_back(&mut self) -> Option<<I as Iterator>::Item> { + self.iter.next() + } + + #[inline] + fn advance_back_by(&mut self, n: usize) -> Result<(), usize> { + self.iter.advance_by(n) + } + + #[inline] + fn nth_back(&mut self, n: usize) -> Option<<I as Iterator>::Item> { + self.iter.nth(n) + } + + fn try_rfold<B, F, R>(&mut self, init: B, f: F) -> R + where + Self: Sized, + F: FnMut(B, Self::Item) -> R, + R: Try<Output = B>, + { + self.iter.try_fold(init, f) + } + + fn rfold<Acc, F>(self, init: Acc, f: F) -> Acc + where + F: FnMut(Acc, Self::Item) -> Acc, + { + self.iter.fold(init, f) + } + + fn rfind<P>(&mut self, predicate: P) -> Option<Self::Item> + where + P: FnMut(&Self::Item) -> bool, + { + self.iter.find(predicate) + } +} + +#[stable(feature = "rust1", since = "1.0.0")] +impl<I> ExactSizeIterator for Rev<I> +where + I: ExactSizeIterator + DoubleEndedIterator, +{ + fn len(&self) -> usize { + self.iter.len() + } + + fn is_empty(&self) -> bool { + self.iter.is_empty() + } +} + +#[stable(feature = "fused", since = "1.26.0")] +impl<I> FusedIterator for Rev<I> where I: FusedIterator + DoubleEndedIterator {} + +#[unstable(feature = "trusted_len", issue = "37572")] +unsafe impl<I> TrustedLen for Rev<I> where I: TrustedLen + DoubleEndedIterator {} diff --git a/library/core/src/iter/adapters/scan.rs b/library/core/src/iter/adapters/scan.rs new file mode 100644 index 000000000..80bfd2231 --- /dev/null +++ b/library/core/src/iter/adapters/scan.rs @@ -0,0 +1,110 @@ +use crate::fmt; +use crate::iter::{adapters::SourceIter, InPlaceIterable}; +use crate::ops::{ControlFlow, Try}; + +/// An iterator to maintain state while iterating another iterator. +/// +/// This `struct` is created by the [`scan`] method on [`Iterator`]. See its +/// documentation for more. +/// +/// [`scan`]: Iterator::scan +/// [`Iterator`]: trait.Iterator.html +#[must_use = "iterators are lazy and do nothing unless consumed"] +#[stable(feature = "rust1", since = "1.0.0")] +#[derive(Clone)] +pub struct Scan<I, St, F> { + iter: I, + f: F, + state: St, +} + +impl<I, St, F> Scan<I, St, F> { + pub(in crate::iter) fn new(iter: I, state: St, f: F) -> Scan<I, St, F> { + Scan { iter, state, f } + } +} + +#[stable(feature = "core_impl_debug", since = "1.9.0")] +impl<I: fmt::Debug, St: fmt::Debug, F> fmt::Debug for Scan<I, St, F> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Scan").field("iter", &self.iter).field("state", &self.state).finish() + } +} + +#[stable(feature = "rust1", since = "1.0.0")] +impl<B, I, St, F> Iterator for Scan<I, St, F> +where + I: Iterator, + F: FnMut(&mut St, I::Item) -> Option<B>, +{ + type Item = B; + + #[inline] + fn next(&mut self) -> Option<B> { + let a = self.iter.next()?; + (self.f)(&mut self.state, a) + } + + #[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 scan function + } + + #[inline] + fn try_fold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> R, + R: Try<Output = Acc>, + { + fn scan<'a, T, St, B, Acc, R: Try<Output = Acc>>( + state: &'a mut St, + f: &'a mut impl FnMut(&mut St, T) -> Option<B>, + mut fold: impl FnMut(Acc, B) -> R + 'a, + ) -> impl FnMut(Acc, T) -> ControlFlow<R, Acc> + 'a { + move |acc, x| match f(state, x) { + None => ControlFlow::Break(try { acc }), + Some(x) => ControlFlow::from_try(fold(acc, x)), + } + } + + let state = &mut self.state; + let f = &mut self.f; + self.iter.try_fold(init, scan(state, f, fold)).into_try() + } + + #[inline] + fn fold<Acc, Fold>(mut self, init: Acc, fold: Fold) -> Acc + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> Acc, + { + #[inline] + fn ok<B, T>(mut f: impl FnMut(B, T) -> B) -> impl FnMut(B, T) -> Result<B, !> { + move |acc, x| Ok(f(acc, x)) + } + + self.try_fold(init, ok(fold)).unwrap() + } +} + +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl<St, F, I> SourceIter for Scan<I, St, F> +where + I: SourceIter, +{ + type Source = I::Source; + + #[inline] + unsafe fn as_inner(&mut self) -> &mut I::Source { + // SAFETY: unsafe function forwarding to unsafe function with the same requirements + unsafe { SourceIter::as_inner(&mut self.iter) } + } +} + +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl<St, F, B, I: InPlaceIterable> InPlaceIterable for Scan<I, St, F> where + F: FnMut(&mut St, I::Item) -> Option<B> +{ +} diff --git a/library/core/src/iter/adapters/skip.rs b/library/core/src/iter/adapters/skip.rs new file mode 100644 index 000000000..2c283100f --- /dev/null +++ b/library/core/src/iter/adapters/skip.rs @@ -0,0 +1,239 @@ +use crate::intrinsics::unlikely; +use crate::iter::{adapters::SourceIter, FusedIterator, InPlaceIterable}; +use crate::ops::{ControlFlow, Try}; + +/// An iterator that skips over `n` elements of `iter`. +/// +/// This `struct` is created by the [`skip`] method on [`Iterator`]. See its +/// documentation for more. +/// +/// [`skip`]: Iterator::skip +/// [`Iterator`]: trait.Iterator.html +#[derive(Clone, Debug)] +#[must_use = "iterators are lazy and do nothing unless consumed"] +#[stable(feature = "rust1", since = "1.0.0")] +pub struct Skip<I> { + iter: I, + n: usize, +} + +impl<I> Skip<I> { + pub(in crate::iter) fn new(iter: I, n: usize) -> Skip<I> { + Skip { iter, n } + } +} + +#[stable(feature = "rust1", since = "1.0.0")] +impl<I> Iterator for Skip<I> +where + I: Iterator, +{ + type Item = <I as Iterator>::Item; + + #[inline] + fn next(&mut self) -> Option<I::Item> { + if unlikely(self.n > 0) { + self.iter.nth(crate::mem::take(&mut self.n) - 1)?; + } + self.iter.next() + } + + #[inline] + fn nth(&mut self, n: usize) -> Option<I::Item> { + // Can't just add n + self.n due to overflow. + if self.n > 0 { + let to_skip = self.n; + self.n = 0; + // nth(n) skips n+1 + self.iter.nth(to_skip - 1)?; + } + self.iter.nth(n) + } + + #[inline] + fn count(mut self) -> usize { + if self.n > 0 { + // nth(n) skips n+1 + if self.iter.nth(self.n - 1).is_none() { + return 0; + } + } + self.iter.count() + } + + #[inline] + fn last(mut self) -> Option<I::Item> { + if self.n > 0 { + // nth(n) skips n+1 + self.iter.nth(self.n - 1)?; + } + self.iter.last() + } + + #[inline] + fn size_hint(&self) -> (usize, Option<usize>) { + let (lower, upper) = self.iter.size_hint(); + + let lower = lower.saturating_sub(self.n); + let upper = match upper { + Some(x) => Some(x.saturating_sub(self.n)), + None => None, + }; + + (lower, upper) + } + + #[inline] + fn try_fold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> R, + R: Try<Output = Acc>, + { + let n = self.n; + self.n = 0; + if n > 0 { + // nth(n) skips n+1 + if self.iter.nth(n - 1).is_none() { + return try { init }; + } + } + self.iter.try_fold(init, fold) + } + + #[inline] + fn fold<Acc, Fold>(mut self, init: Acc, fold: Fold) -> Acc + where + Fold: FnMut(Acc, Self::Item) -> Acc, + { + if self.n > 0 { + // nth(n) skips n+1 + if self.iter.nth(self.n - 1).is_none() { + return init; + } + } + self.iter.fold(init, fold) + } + + #[inline] + #[rustc_inherit_overflow_checks] + fn advance_by(&mut self, n: usize) -> Result<(), usize> { + let mut rem = n; + let step_one = self.n.saturating_add(rem); + + match self.iter.advance_by(step_one) { + Ok(_) => { + rem -= step_one - self.n; + self.n = 0; + } + Err(advanced) => { + let advanced_without_skip = advanced.saturating_sub(self.n); + self.n = self.n.saturating_sub(advanced); + return if n == 0 { Ok(()) } else { Err(advanced_without_skip) }; + } + } + + // step_one calculation may have saturated + if unlikely(rem > 0) { + return match self.iter.advance_by(rem) { + ret @ Ok(_) => ret, + Err(advanced) => { + rem -= advanced; + Err(n - rem) + } + }; + } + + Ok(()) + } +} + +#[stable(feature = "rust1", since = "1.0.0")] +impl<I> ExactSizeIterator for Skip<I> where I: ExactSizeIterator {} + +#[stable(feature = "double_ended_skip_iterator", since = "1.9.0")] +impl<I> DoubleEndedIterator for Skip<I> +where + I: DoubleEndedIterator + ExactSizeIterator, +{ + fn next_back(&mut self) -> Option<Self::Item> { + if self.len() > 0 { self.iter.next_back() } else { None } + } + + #[inline] + fn nth_back(&mut self, n: usize) -> Option<I::Item> { + let len = self.len(); + if n < len { + self.iter.nth_back(n) + } else { + if len > 0 { + // consume the original iterator + self.iter.nth_back(len - 1); + } + None + } + } + + fn try_rfold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> R, + R: Try<Output = Acc>, + { + fn check<T, Acc, R: Try<Output = Acc>>( + mut n: usize, + mut fold: impl FnMut(Acc, T) -> R, + ) -> impl FnMut(Acc, T) -> ControlFlow<R, Acc> { + move |acc, x| { + n -= 1; + let r = fold(acc, x); + if n == 0 { ControlFlow::Break(r) } else { ControlFlow::from_try(r) } + } + } + + let n = self.len(); + if n == 0 { try { init } } else { self.iter.try_rfold(init, check(n, fold)).into_try() } + } + + fn rfold<Acc, Fold>(mut self, init: Acc, fold: Fold) -> Acc + where + Fold: FnMut(Acc, Self::Item) -> Acc, + { + #[inline] + fn ok<Acc, T>(mut f: impl FnMut(Acc, T) -> Acc) -> impl FnMut(Acc, T) -> Result<Acc, !> { + move |acc, x| Ok(f(acc, x)) + } + + self.try_rfold(init, ok(fold)).unwrap() + } + + #[inline] + fn advance_back_by(&mut self, n: usize) -> Result<(), usize> { + let min = crate::cmp::min(self.len(), n); + return match self.iter.advance_back_by(min) { + ret @ Ok(_) if n <= min => ret, + Ok(_) => Err(min), + _ => panic!("ExactSizeIterator contract violation"), + }; + } +} + +#[stable(feature = "fused", since = "1.26.0")] +impl<I> FusedIterator for Skip<I> where I: FusedIterator {} + +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl<I> SourceIter for Skip<I> +where + I: SourceIter, +{ + type Source = I::Source; + + #[inline] + unsafe fn as_inner(&mut self) -> &mut I::Source { + // SAFETY: unsafe function forwarding to unsafe function with the same requirements + unsafe { SourceIter::as_inner(&mut self.iter) } + } +} + +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl<I: InPlaceIterable> InPlaceIterable for Skip<I> {} diff --git a/library/core/src/iter/adapters/skip_while.rs b/library/core/src/iter/adapters/skip_while.rs new file mode 100644 index 000000000..f29661779 --- /dev/null +++ b/library/core/src/iter/adapters/skip_while.rs @@ -0,0 +1,125 @@ +use crate::fmt; +use crate::iter::{adapters::SourceIter, FusedIterator, InPlaceIterable}; +use crate::ops::Try; + +/// An iterator that rejects elements while `predicate` returns `true`. +/// +/// This `struct` is created by the [`skip_while`] method on [`Iterator`]. See its +/// documentation for more. +/// +/// [`skip_while`]: Iterator::skip_while +/// [`Iterator`]: trait.Iterator.html +#[must_use = "iterators are lazy and do nothing unless consumed"] +#[stable(feature = "rust1", since = "1.0.0")] +#[derive(Clone)] +pub struct SkipWhile<I, P> { + iter: I, + flag: bool, + predicate: P, +} + +impl<I, P> SkipWhile<I, P> { + pub(in crate::iter) fn new(iter: I, predicate: P) -> SkipWhile<I, P> { + SkipWhile { iter, flag: false, predicate } + } +} + +#[stable(feature = "core_impl_debug", since = "1.9.0")] +impl<I: fmt::Debug, P> fmt::Debug for SkipWhile<I, P> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("SkipWhile").field("iter", &self.iter).field("flag", &self.flag).finish() + } +} + +#[stable(feature = "rust1", since = "1.0.0")] +impl<I: Iterator, P> Iterator for SkipWhile<I, P> +where + P: FnMut(&I::Item) -> bool, +{ + type Item = I::Item; + + #[inline] + fn next(&mut self) -> Option<I::Item> { + fn check<'a, T>( + flag: &'a mut bool, + pred: &'a mut impl FnMut(&T) -> bool, + ) -> impl FnMut(&T) -> bool + 'a { + move |x| { + if *flag || !pred(x) { + *flag = true; + true + } else { + false + } + } + } + + let flag = &mut self.flag; + let pred = &mut self.predicate; + self.iter.find(check(flag, pred)) + } + + #[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 + } + + #[inline] + fn try_fold<Acc, Fold, R>(&mut self, mut init: Acc, mut fold: Fold) -> R + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> R, + R: Try<Output = Acc>, + { + if !self.flag { + match self.next() { + Some(v) => init = fold(init, v)?, + None => return try { init }, + } + } + self.iter.try_fold(init, fold) + } + + #[inline] + fn fold<Acc, Fold>(mut self, mut init: Acc, mut fold: Fold) -> Acc + where + Fold: FnMut(Acc, Self::Item) -> Acc, + { + if !self.flag { + match self.next() { + Some(v) => init = fold(init, v), + None => return init, + } + } + self.iter.fold(init, fold) + } +} + +#[stable(feature = "fused", since = "1.26.0")] +impl<I, P> FusedIterator for SkipWhile<I, P> +where + I: FusedIterator, + P: FnMut(&I::Item) -> bool, +{ +} + +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl<P, I> SourceIter for SkipWhile<I, P> +where + I: SourceIter, +{ + type Source = I::Source; + + #[inline] + unsafe fn as_inner(&mut self) -> &mut I::Source { + // SAFETY: unsafe function forwarding to unsafe function with the same requirements + unsafe { SourceIter::as_inner(&mut self.iter) } + } +} + +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl<I: InPlaceIterable, F> InPlaceIterable for SkipWhile<I, F> where + F: FnMut(&I::Item) -> bool +{ +} diff --git a/library/core/src/iter/adapters/step_by.rs b/library/core/src/iter/adapters/step_by.rs new file mode 100644 index 000000000..4252c34a0 --- /dev/null +++ b/library/core/src/iter/adapters/step_by.rs @@ -0,0 +1,235 @@ +use crate::{intrinsics, iter::from_fn, ops::Try}; + +/// An iterator for stepping iterators by a custom amount. +/// +/// This `struct` is created by the [`step_by`] method on [`Iterator`]. See +/// its documentation for more. +/// +/// [`step_by`]: Iterator::step_by +/// [`Iterator`]: trait.Iterator.html +#[must_use = "iterators are lazy and do nothing unless consumed"] +#[stable(feature = "iterator_step_by", since = "1.28.0")] +#[derive(Clone, Debug)] +pub struct StepBy<I> { + iter: I, + step: usize, + first_take: bool, +} + +impl<I> StepBy<I> { + pub(in crate::iter) fn new(iter: I, step: usize) -> StepBy<I> { + assert!(step != 0); + StepBy { iter, step: step - 1, first_take: true } + } +} + +#[stable(feature = "iterator_step_by", since = "1.28.0")] +impl<I> Iterator for StepBy<I> +where + I: Iterator, +{ + type Item = I::Item; + + #[inline] + fn next(&mut self) -> Option<Self::Item> { + if self.first_take { + self.first_take = false; + self.iter.next() + } else { + self.iter.nth(self.step) + } + } + + #[inline] + fn size_hint(&self) -> (usize, Option<usize>) { + #[inline] + fn first_size(step: usize) -> impl Fn(usize) -> usize { + move |n| if n == 0 { 0 } else { 1 + (n - 1) / (step + 1) } + } + + #[inline] + fn other_size(step: usize) -> impl Fn(usize) -> usize { + move |n| n / (step + 1) + } + + let (low, high) = self.iter.size_hint(); + + if self.first_take { + let f = first_size(self.step); + (f(low), high.map(f)) + } else { + let f = other_size(self.step); + (f(low), high.map(f)) + } + } + + #[inline] + fn nth(&mut self, mut n: usize) -> Option<Self::Item> { + if self.first_take { + self.first_take = false; + let first = self.iter.next(); + if n == 0 { + return first; + } + n -= 1; + } + // n and self.step are indices, we need to add 1 to get the amount of elements + // When calling `.nth`, we need to subtract 1 again to convert back to an index + // step + 1 can't overflow because `.step_by` sets `self.step` to `step - 1` + let mut step = self.step + 1; + // n + 1 could overflow + // thus, if n is usize::MAX, instead of adding one, we call .nth(step) + if n == usize::MAX { + self.iter.nth(step - 1); + } else { + n += 1; + } + + // overflow handling + loop { + let mul = n.checked_mul(step); + { + if intrinsics::likely(mul.is_some()) { + return self.iter.nth(mul.unwrap() - 1); + } + } + let div_n = usize::MAX / n; + let div_step = usize::MAX / step; + let nth_n = div_n * n; + let nth_step = div_step * step; + let nth = if nth_n > nth_step { + step -= div_n; + nth_n + } else { + n -= div_step; + nth_step + }; + self.iter.nth(nth - 1); + } + } + + fn try_fold<Acc, F, R>(&mut self, mut acc: Acc, mut f: F) -> R + where + F: FnMut(Acc, Self::Item) -> R, + R: Try<Output = Acc>, + { + #[inline] + fn nth<I: Iterator>(iter: &mut I, step: usize) -> impl FnMut() -> Option<I::Item> + '_ { + move || iter.nth(step) + } + + if self.first_take { + self.first_take = false; + match self.iter.next() { + None => return try { acc }, + Some(x) => acc = f(acc, x)?, + } + } + from_fn(nth(&mut self.iter, self.step)).try_fold(acc, f) + } + + fn fold<Acc, F>(mut self, mut acc: Acc, mut f: F) -> Acc + where + F: FnMut(Acc, Self::Item) -> Acc, + { + #[inline] + fn nth<I: Iterator>(iter: &mut I, step: usize) -> impl FnMut() -> Option<I::Item> + '_ { + move || iter.nth(step) + } + + if self.first_take { + self.first_take = false; + match self.iter.next() { + None => return acc, + Some(x) => acc = f(acc, x), + } + } + from_fn(nth(&mut self.iter, self.step)).fold(acc, f) + } +} + +impl<I> StepBy<I> +where + I: ExactSizeIterator, +{ + // The zero-based index starting from the end of the iterator of the + // last element. Used in the `DoubleEndedIterator` implementation. + fn next_back_index(&self) -> usize { + let rem = self.iter.len() % (self.step + 1); + if self.first_take { + if rem == 0 { self.step } else { rem - 1 } + } else { + rem + } + } +} + +#[stable(feature = "double_ended_step_by_iterator", since = "1.38.0")] +impl<I> DoubleEndedIterator for StepBy<I> +where + I: DoubleEndedIterator + ExactSizeIterator, +{ + #[inline] + fn next_back(&mut self) -> Option<Self::Item> { + self.iter.nth_back(self.next_back_index()) + } + + #[inline] + fn nth_back(&mut self, n: usize) -> Option<Self::Item> { + // `self.iter.nth_back(usize::MAX)` does the right thing here when `n` + // is out of bounds because the length of `self.iter` does not exceed + // `usize::MAX` (because `I: ExactSizeIterator`) and `nth_back` is + // zero-indexed + let n = n.saturating_mul(self.step + 1).saturating_add(self.next_back_index()); + self.iter.nth_back(n) + } + + fn try_rfold<Acc, F, R>(&mut self, init: Acc, mut f: F) -> R + where + F: FnMut(Acc, Self::Item) -> R, + R: Try<Output = Acc>, + { + #[inline] + fn nth_back<I: DoubleEndedIterator>( + iter: &mut I, + step: usize, + ) -> impl FnMut() -> Option<I::Item> + '_ { + move || iter.nth_back(step) + } + + match self.next_back() { + None => try { init }, + Some(x) => { + let acc = f(init, x)?; + from_fn(nth_back(&mut self.iter, self.step)).try_fold(acc, f) + } + } + } + + #[inline] + fn rfold<Acc, F>(mut self, init: Acc, mut f: F) -> Acc + where + Self: Sized, + F: FnMut(Acc, Self::Item) -> Acc, + { + #[inline] + fn nth_back<I: DoubleEndedIterator>( + iter: &mut I, + step: usize, + ) -> impl FnMut() -> Option<I::Item> + '_ { + move || iter.nth_back(step) + } + + match self.next_back() { + None => init, + Some(x) => { + let acc = f(init, x); + from_fn(nth_back(&mut self.iter, self.step)).fold(acc, f) + } + } + } +} + +// StepBy can only make the iterator shorter, so the len will still fit. +#[stable(feature = "iterator_step_by", since = "1.28.0")] +impl<I> ExactSizeIterator for StepBy<I> where I: ExactSizeIterator {} diff --git a/library/core/src/iter/adapters/take.rs b/library/core/src/iter/adapters/take.rs new file mode 100644 index 000000000..2962e0104 --- /dev/null +++ b/library/core/src/iter/adapters/take.rs @@ -0,0 +1,244 @@ +use crate::cmp; +use crate::iter::{adapters::SourceIter, FusedIterator, InPlaceIterable, TrustedLen}; +use crate::ops::{ControlFlow, Try}; + +/// An iterator that only iterates over the first `n` iterations of `iter`. +/// +/// This `struct` is created by the [`take`] method on [`Iterator`]. See its +/// documentation for more. +/// +/// [`take`]: Iterator::take +/// [`Iterator`]: trait.Iterator.html +#[derive(Clone, Debug)] +#[must_use = "iterators are lazy and do nothing unless consumed"] +#[stable(feature = "rust1", since = "1.0.0")] +pub struct Take<I> { + iter: I, + n: usize, +} + +impl<I> Take<I> { + pub(in crate::iter) fn new(iter: I, n: usize) -> Take<I> { + Take { iter, n } + } +} + +#[stable(feature = "rust1", since = "1.0.0")] +impl<I> Iterator for Take<I> +where + I: Iterator, +{ + type Item = <I as Iterator>::Item; + + #[inline] + fn next(&mut self) -> Option<<I as Iterator>::Item> { + if self.n != 0 { + self.n -= 1; + self.iter.next() + } else { + None + } + } + + #[inline] + fn nth(&mut self, n: usize) -> Option<I::Item> { + if self.n > n { + self.n -= n + 1; + self.iter.nth(n) + } else { + if self.n > 0 { + self.iter.nth(self.n - 1); + self.n = 0; + } + None + } + } + + #[inline] + fn size_hint(&self) -> (usize, Option<usize>) { + if self.n == 0 { + return (0, Some(0)); + } + + let (lower, upper) = self.iter.size_hint(); + + let lower = cmp::min(lower, self.n); + + let upper = match upper { + Some(x) if x < self.n => Some(x), + _ => Some(self.n), + }; + + (lower, upper) + } + + #[inline] + fn try_fold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> R, + R: Try<Output = Acc>, + { + fn check<'a, T, Acc, R: Try<Output = Acc>>( + n: &'a mut usize, + mut fold: impl FnMut(Acc, T) -> R + 'a, + ) -> impl FnMut(Acc, T) -> ControlFlow<R, Acc> + 'a { + move |acc, x| { + *n -= 1; + let r = fold(acc, x); + if *n == 0 { ControlFlow::Break(r) } else { ControlFlow::from_try(r) } + } + } + + if self.n == 0 { + try { init } + } else { + let n = &mut self.n; + self.iter.try_fold(init, check(n, fold)).into_try() + } + } + + #[inline] + fn fold<Acc, Fold>(mut self, init: Acc, fold: Fold) -> Acc + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> Acc, + { + #[inline] + fn ok<B, T>(mut f: impl FnMut(B, T) -> B) -> impl FnMut(B, T) -> Result<B, !> { + move |acc, x| Ok(f(acc, x)) + } + + self.try_fold(init, ok(fold)).unwrap() + } + + #[inline] + #[rustc_inherit_overflow_checks] + fn advance_by(&mut self, n: usize) -> Result<(), usize> { + let min = self.n.min(n); + match self.iter.advance_by(min) { + Ok(_) => { + self.n -= min; + if min < n { Err(min) } else { Ok(()) } + } + ret @ Err(advanced) => { + self.n -= advanced; + ret + } + } + } +} + +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl<I> SourceIter for Take<I> +where + I: SourceIter, +{ + type Source = I::Source; + + #[inline] + unsafe fn as_inner(&mut self) -> &mut I::Source { + // SAFETY: unsafe function forwarding to unsafe function with the same requirements + unsafe { SourceIter::as_inner(&mut self.iter) } + } +} + +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl<I: InPlaceIterable> InPlaceIterable for Take<I> {} + +#[stable(feature = "double_ended_take_iterator", since = "1.38.0")] +impl<I> DoubleEndedIterator for Take<I> +where + I: DoubleEndedIterator + ExactSizeIterator, +{ + #[inline] + fn next_back(&mut self) -> Option<Self::Item> { + if self.n == 0 { + None + } else { + let n = self.n; + self.n -= 1; + self.iter.nth_back(self.iter.len().saturating_sub(n)) + } + } + + #[inline] + fn nth_back(&mut self, n: usize) -> Option<Self::Item> { + let len = self.iter.len(); + if self.n > n { + let m = len.saturating_sub(self.n) + n; + self.n -= n + 1; + self.iter.nth_back(m) + } else { + if len > 0 { + self.iter.nth_back(len - 1); + } + None + } + } + + #[inline] + fn try_rfold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> R, + R: Try<Output = Acc>, + { + if self.n == 0 { + try { init } + } else { + let len = self.iter.len(); + if len > self.n && self.iter.nth_back(len - self.n - 1).is_none() { + try { init } + } else { + self.iter.try_rfold(init, fold) + } + } + } + + #[inline] + fn rfold<Acc, Fold>(mut self, init: Acc, fold: Fold) -> Acc + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> Acc, + { + if self.n == 0 { + init + } else { + let len = self.iter.len(); + if len > self.n && self.iter.nth_back(len - self.n - 1).is_none() { + init + } else { + self.iter.rfold(init, fold) + } + } + } + + #[inline] + #[rustc_inherit_overflow_checks] + fn advance_back_by(&mut self, n: usize) -> Result<(), usize> { + // The amount by which the inner iterator needs to be shortened for it to be + // at most as long as the take() amount. + let trim_inner = self.iter.len().saturating_sub(self.n); + // The amount we need to advance inner to fulfill the caller's request. + // take(), advance_by() and len() all can be at most usize, so we don't have to worry + // about having to advance more than usize::MAX here. + let advance_by = trim_inner.saturating_add(n); + + let advanced = match self.iter.advance_back_by(advance_by) { + Ok(_) => advance_by - trim_inner, + Err(advanced) => advanced - trim_inner, + }; + self.n -= advanced; + return if advanced < n { Err(advanced) } else { Ok(()) }; + } +} + +#[stable(feature = "rust1", since = "1.0.0")] +impl<I> ExactSizeIterator for Take<I> where I: ExactSizeIterator {} + +#[stable(feature = "fused", since = "1.26.0")] +impl<I> FusedIterator for Take<I> where I: FusedIterator {} + +#[unstable(feature = "trusted_len", issue = "37572")] +unsafe impl<I: TrustedLen> TrustedLen for Take<I> {} diff --git a/library/core/src/iter/adapters/take_while.rs b/library/core/src/iter/adapters/take_while.rs new file mode 100644 index 000000000..ded216da9 --- /dev/null +++ b/library/core/src/iter/adapters/take_while.rs @@ -0,0 +1,138 @@ +use crate::fmt; +use crate::iter::{adapters::SourceIter, FusedIterator, InPlaceIterable}; +use crate::ops::{ControlFlow, Try}; + +/// An iterator that only accepts elements while `predicate` returns `true`. +/// +/// This `struct` is created by the [`take_while`] method on [`Iterator`]. See its +/// documentation for more. +/// +/// [`take_while`]: Iterator::take_while +/// [`Iterator`]: trait.Iterator.html +#[must_use = "iterators are lazy and do nothing unless consumed"] +#[stable(feature = "rust1", since = "1.0.0")] +#[derive(Clone)] +pub struct TakeWhile<I, P> { + iter: I, + flag: bool, + predicate: P, +} + +impl<I, P> TakeWhile<I, P> { + pub(in crate::iter) fn new(iter: I, predicate: P) -> TakeWhile<I, P> { + TakeWhile { iter, flag: false, predicate } + } +} + +#[stable(feature = "core_impl_debug", since = "1.9.0")] +impl<I: fmt::Debug, P> fmt::Debug for TakeWhile<I, P> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("TakeWhile").field("iter", &self.iter).field("flag", &self.flag).finish() + } +} + +#[stable(feature = "rust1", since = "1.0.0")] +impl<I: Iterator, P> Iterator for TakeWhile<I, P> +where + P: FnMut(&I::Item) -> bool, +{ + type Item = I::Item; + + #[inline] + fn next(&mut self) -> Option<I::Item> { + if self.flag { + None + } else { + let x = self.iter.next()?; + if (self.predicate)(&x) { + Some(x) + } 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 + } + } + + #[inline] + fn try_fold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> R, + R: Try<Output = Acc>, + { + fn check<'a, T, Acc, R: Try<Output = Acc>>( + flag: &'a mut bool, + p: &'a mut impl FnMut(&T) -> bool, + mut fold: impl FnMut(Acc, T) -> R + 'a, + ) -> impl FnMut(Acc, T) -> ControlFlow<R, Acc> + 'a { + move |acc, x| { + if p(&x) { + ControlFlow::from_try(fold(acc, x)) + } else { + *flag = true; + ControlFlow::Break(try { acc }) + } + } + } + + if self.flag { + try { init } + } else { + let flag = &mut self.flag; + let p = &mut self.predicate; + self.iter.try_fold(init, check(flag, p, fold)).into_try() + } + } + + #[inline] + fn fold<Acc, Fold>(mut self, init: Acc, fold: Fold) -> Acc + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> Acc, + { + #[inline] + fn ok<B, T>(mut f: impl FnMut(B, T) -> B) -> impl FnMut(B, T) -> Result<B, !> { + move |acc, x| Ok(f(acc, x)) + } + + self.try_fold(init, ok(fold)).unwrap() + } +} + +#[stable(feature = "fused", since = "1.26.0")] +impl<I, P> FusedIterator for TakeWhile<I, P> +where + I: FusedIterator, + P: FnMut(&I::Item) -> bool, +{ +} + +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl<P, I> SourceIter for TakeWhile<I, P> +where + I: SourceIter, +{ + type Source = I::Source; + + #[inline] + unsafe fn as_inner(&mut self) -> &mut I::Source { + // SAFETY: unsafe function forwarding to unsafe function with the same requirements + unsafe { SourceIter::as_inner(&mut self.iter) } + } +} + +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl<I: InPlaceIterable, F> InPlaceIterable for TakeWhile<I, F> where + F: FnMut(&I::Item) -> bool +{ +} diff --git a/library/core/src/iter/adapters/zip.rs b/library/core/src/iter/adapters/zip.rs new file mode 100644 index 000000000..8153c8cfe --- /dev/null +++ b/library/core/src/iter/adapters/zip.rs @@ -0,0 +1,585 @@ +use crate::cmp; +use crate::fmt::{self, Debug}; +use crate::iter::{DoubleEndedIterator, ExactSizeIterator, FusedIterator, Iterator}; +use crate::iter::{InPlaceIterable, SourceIter, TrustedLen}; + +/// An iterator that iterates two other iterators simultaneously. +/// +/// This `struct` is created by [`zip`] or [`Iterator::zip`]. +/// See their documentation for more. +#[derive(Clone)] +#[must_use = "iterators are lazy and do nothing unless consumed"] +#[stable(feature = "rust1", since = "1.0.0")] +pub struct Zip<A, B> { + a: A, + b: B, + // index, len and a_len are only used by the specialized version of zip + index: usize, + len: usize, + a_len: usize, +} +impl<A: Iterator, B: Iterator> Zip<A, B> { + pub(in crate::iter) fn new(a: A, b: B) -> Zip<A, B> { + ZipImpl::new(a, b) + } + fn super_nth(&mut self, mut n: usize) -> Option<(A::Item, B::Item)> { + while let Some(x) = Iterator::next(self) { + if n == 0 { + return Some(x); + } + n -= 1; + } + None + } +} + +/// Converts the arguments to iterators and zips them. +/// +/// See the documentation of [`Iterator::zip`] for more. +/// +/// # Examples +/// +/// ``` +/// use std::iter::zip; +/// +/// let xs = [1, 2, 3]; +/// let ys = [4, 5, 6]; +/// +/// let mut iter = zip(xs, ys); +/// +/// assert_eq!(iter.next().unwrap(), (1, 4)); +/// assert_eq!(iter.next().unwrap(), (2, 5)); +/// assert_eq!(iter.next().unwrap(), (3, 6)); +/// assert!(iter.next().is_none()); +/// +/// // Nested zips are also possible: +/// let zs = [7, 8, 9]; +/// +/// let mut iter = zip(zip(xs, ys), zs); +/// +/// assert_eq!(iter.next().unwrap(), ((1, 4), 7)); +/// assert_eq!(iter.next().unwrap(), ((2, 5), 8)); +/// assert_eq!(iter.next().unwrap(), ((3, 6), 9)); +/// assert!(iter.next().is_none()); +/// ``` +#[stable(feature = "iter_zip", since = "1.59.0")] +pub fn zip<A, B>(a: A, b: B) -> Zip<A::IntoIter, B::IntoIter> +where + A: IntoIterator, + B: IntoIterator, +{ + ZipImpl::new(a.into_iter(), b.into_iter()) +} + +#[stable(feature = "rust1", since = "1.0.0")] +impl<A, B> Iterator for Zip<A, B> +where + A: Iterator, + B: Iterator, +{ + type Item = (A::Item, B::Item); + + #[inline] + fn next(&mut self) -> Option<Self::Item> { + ZipImpl::next(self) + } + + #[inline] + fn size_hint(&self) -> (usize, Option<usize>) { + ZipImpl::size_hint(self) + } + + #[inline] + fn nth(&mut self, n: usize) -> Option<Self::Item> { + ZipImpl::nth(self, n) + } + + #[inline] + unsafe fn __iterator_get_unchecked(&mut self, idx: usize) -> Self::Item + where + Self: TrustedRandomAccessNoCoerce, + { + // SAFETY: `ZipImpl::__iterator_get_unchecked` has same safety + // requirements as `Iterator::__iterator_get_unchecked`. + unsafe { ZipImpl::get_unchecked(self, idx) } + } +} + +#[stable(feature = "rust1", since = "1.0.0")] +impl<A, B> DoubleEndedIterator for Zip<A, B> +where + A: DoubleEndedIterator + ExactSizeIterator, + B: DoubleEndedIterator + ExactSizeIterator, +{ + #[inline] + fn next_back(&mut self) -> Option<(A::Item, B::Item)> { + ZipImpl::next_back(self) + } +} + +// Zip specialization trait +#[doc(hidden)] +trait ZipImpl<A, B> { + type Item; + fn new(a: A, b: B) -> Self; + fn next(&mut self) -> Option<Self::Item>; + fn size_hint(&self) -> (usize, Option<usize>); + fn nth(&mut self, n: usize) -> Option<Self::Item>; + fn next_back(&mut self) -> Option<Self::Item> + where + A: DoubleEndedIterator + ExactSizeIterator, + B: DoubleEndedIterator + ExactSizeIterator; + // This has the same safety requirements as `Iterator::__iterator_get_unchecked` + unsafe fn get_unchecked(&mut self, idx: usize) -> <Self as Iterator>::Item + where + Self: Iterator + TrustedRandomAccessNoCoerce; +} + +// Work around limitations of specialization, requiring `default` impls to be repeated +// in intermediary impls. +macro_rules! zip_impl_general_defaults { + () => { + default fn new(a: A, b: B) -> Self { + Zip { + a, + b, + index: 0, // unused + len: 0, // unused + a_len: 0, // unused + } + } + + #[inline] + default fn next(&mut self) -> Option<(A::Item, B::Item)> { + let x = self.a.next()?; + let y = self.b.next()?; + Some((x, y)) + } + + #[inline] + default fn nth(&mut self, n: usize) -> Option<Self::Item> { + self.super_nth(n) + } + + #[inline] + default fn next_back(&mut self) -> Option<(A::Item, B::Item)> + where + A: DoubleEndedIterator + ExactSizeIterator, + B: DoubleEndedIterator + ExactSizeIterator, + { + // The function body below only uses `self.a/b.len()` and `self.a/b.next_back()` + // and doesn’t call `next_back` too often, so this implementation is safe in + // the `TrustedRandomAccessNoCoerce` specialization + + let a_sz = self.a.len(); + let b_sz = self.b.len(); + if a_sz != b_sz { + // Adjust a, b to equal length + if a_sz > b_sz { + for _ in 0..a_sz - b_sz { + self.a.next_back(); + } + } else { + for _ in 0..b_sz - a_sz { + self.b.next_back(); + } + } + } + match (self.a.next_back(), self.b.next_back()) { + (Some(x), Some(y)) => Some((x, y)), + (None, None) => None, + _ => unreachable!(), + } + } + }; +} + +// General Zip impl +#[doc(hidden)] +impl<A, B> ZipImpl<A, B> for Zip<A, B> +where + A: Iterator, + B: Iterator, +{ + type Item = (A::Item, B::Item); + + zip_impl_general_defaults! {} + + #[inline] + default fn size_hint(&self) -> (usize, Option<usize>) { + let (a_lower, a_upper) = self.a.size_hint(); + let (b_lower, b_upper) = self.b.size_hint(); + + let lower = cmp::min(a_lower, b_lower); + + let upper = match (a_upper, b_upper) { + (Some(x), Some(y)) => Some(cmp::min(x, y)), + (Some(x), None) => Some(x), + (None, Some(y)) => Some(y), + (None, None) => None, + }; + + (lower, upper) + } + + default unsafe fn get_unchecked(&mut self, _idx: usize) -> <Self as Iterator>::Item + where + Self: TrustedRandomAccessNoCoerce, + { + unreachable!("Always specialized"); + } +} + +#[doc(hidden)] +impl<A, B> ZipImpl<A, B> for Zip<A, B> +where + A: TrustedRandomAccessNoCoerce + Iterator, + B: TrustedRandomAccessNoCoerce + Iterator, +{ + zip_impl_general_defaults! {} + + #[inline] + default fn size_hint(&self) -> (usize, Option<usize>) { + let size = cmp::min(self.a.size(), self.b.size()); + (size, Some(size)) + } + + #[inline] + unsafe fn get_unchecked(&mut self, idx: usize) -> <Self as Iterator>::Item { + let idx = self.index + idx; + // SAFETY: the caller must uphold the contract for + // `Iterator::__iterator_get_unchecked`. + unsafe { (self.a.__iterator_get_unchecked(idx), self.b.__iterator_get_unchecked(idx)) } + } +} + +#[doc(hidden)] +impl<A, B> ZipImpl<A, B> for Zip<A, B> +where + A: TrustedRandomAccess + Iterator, + B: TrustedRandomAccess + Iterator, +{ + fn new(a: A, b: B) -> Self { + let a_len = a.size(); + let len = cmp::min(a_len, b.size()); + Zip { a, b, index: 0, len, a_len } + } + + #[inline] + fn next(&mut self) -> Option<(A::Item, B::Item)> { + if self.index < self.len { + let i = self.index; + // since get_unchecked executes code which can panic we increment the counters beforehand + // so that the same index won't be accessed twice, as required by TrustedRandomAccess + self.index += 1; + // SAFETY: `i` is smaller than `self.len`, thus smaller than `self.a.len()` and `self.b.len()` + unsafe { + Some((self.a.__iterator_get_unchecked(i), self.b.__iterator_get_unchecked(i))) + } + } else if A::MAY_HAVE_SIDE_EFFECT && self.index < self.a_len { + let i = self.index; + // as above, increment before executing code that may panic + self.index += 1; + self.len += 1; + // match the base implementation's potential side effects + // SAFETY: we just checked that `i` < `self.a.len()` + unsafe { + self.a.__iterator_get_unchecked(i); + } + None + } else { + None + } + } + + #[inline] + fn size_hint(&self) -> (usize, Option<usize>) { + let len = self.len - self.index; + (len, Some(len)) + } + + #[inline] + fn nth(&mut self, n: usize) -> Option<Self::Item> { + let delta = cmp::min(n, self.len - self.index); + let end = self.index + delta; + while self.index < end { + let i = self.index; + // since get_unchecked executes code which can panic we increment the counters beforehand + // so that the same index won't be accessed twice, as required by TrustedRandomAccess + self.index += 1; + if A::MAY_HAVE_SIDE_EFFECT { + // SAFETY: the usage of `cmp::min` to calculate `delta` + // ensures that `end` is smaller than or equal to `self.len`, + // so `i` is also smaller than `self.len`. + unsafe { + self.a.__iterator_get_unchecked(i); + } + } + if B::MAY_HAVE_SIDE_EFFECT { + // SAFETY: same as above. + unsafe { + self.b.__iterator_get_unchecked(i); + } + } + } + + self.super_nth(n - delta) + } + + #[inline] + fn next_back(&mut self) -> Option<(A::Item, B::Item)> + where + A: DoubleEndedIterator + ExactSizeIterator, + B: DoubleEndedIterator + ExactSizeIterator, + { + if A::MAY_HAVE_SIDE_EFFECT || B::MAY_HAVE_SIDE_EFFECT { + let sz_a = self.a.size(); + let sz_b = self.b.size(); + // Adjust a, b to equal length, make sure that only the first call + // of `next_back` does this, otherwise we will break the restriction + // on calls to `self.next_back()` after calling `get_unchecked()`. + if sz_a != sz_b { + let sz_a = self.a.size(); + if A::MAY_HAVE_SIDE_EFFECT && sz_a > self.len { + for _ in 0..sz_a - self.len { + // since next_back() may panic we increment the counters beforehand + // to keep Zip's state in sync with the underlying iterator source + self.a_len -= 1; + self.a.next_back(); + } + debug_assert_eq!(self.a_len, self.len); + } + let sz_b = self.b.size(); + if B::MAY_HAVE_SIDE_EFFECT && sz_b > self.len { + for _ in 0..sz_b - self.len { + self.b.next_back(); + } + } + } + } + if self.index < self.len { + // since get_unchecked executes code which can panic we increment the counters beforehand + // so that the same index won't be accessed twice, as required by TrustedRandomAccess + self.len -= 1; + self.a_len -= 1; + let i = self.len; + // SAFETY: `i` is smaller than the previous value of `self.len`, + // which is also smaller than or equal to `self.a.len()` and `self.b.len()` + unsafe { + Some((self.a.__iterator_get_unchecked(i), self.b.__iterator_get_unchecked(i))) + } + } else { + None + } + } +} + +#[stable(feature = "rust1", since = "1.0.0")] +impl<A, B> ExactSizeIterator for Zip<A, B> +where + A: ExactSizeIterator, + B: ExactSizeIterator, +{ +} + +#[doc(hidden)] +#[unstable(feature = "trusted_random_access", issue = "none")] +unsafe impl<A, B> TrustedRandomAccess for Zip<A, B> +where + A: TrustedRandomAccess, + B: TrustedRandomAccess, +{ +} + +#[doc(hidden)] +#[unstable(feature = "trusted_random_access", issue = "none")] +unsafe impl<A, B> TrustedRandomAccessNoCoerce for Zip<A, B> +where + A: TrustedRandomAccessNoCoerce, + B: TrustedRandomAccessNoCoerce, +{ + const MAY_HAVE_SIDE_EFFECT: bool = A::MAY_HAVE_SIDE_EFFECT || B::MAY_HAVE_SIDE_EFFECT; +} + +#[stable(feature = "fused", since = "1.26.0")] +impl<A, B> FusedIterator for Zip<A, B> +where + A: FusedIterator, + B: FusedIterator, +{ +} + +#[unstable(feature = "trusted_len", issue = "37572")] +unsafe impl<A, B> TrustedLen for Zip<A, B> +where + A: TrustedLen, + B: TrustedLen, +{ +} + +// Arbitrarily selects the left side of the zip iteration as extractable "source" +// it would require negative trait bounds to be able to try both +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl<A, B> SourceIter for Zip<A, B> +where + A: SourceIter, +{ + type Source = A::Source; + + #[inline] + unsafe fn as_inner(&mut self) -> &mut A::Source { + // SAFETY: unsafe function forwarding to unsafe function with the same requirements + unsafe { SourceIter::as_inner(&mut self.a) } + } +} + +// Since SourceIter forwards the left hand side we do the same here +#[unstable(issue = "none", feature = "inplace_iteration")] +unsafe impl<A: InPlaceIterable, B: Iterator> InPlaceIterable for Zip<A, B> {} + +#[stable(feature = "rust1", since = "1.0.0")] +impl<A: Debug, B: Debug> Debug for Zip<A, B> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + ZipFmt::fmt(self, f) + } +} + +trait ZipFmt<A, B> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result; +} + +impl<A: Debug, B: Debug> ZipFmt<A, B> for Zip<A, B> { + default fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Zip").field("a", &self.a).field("b", &self.b).finish() + } +} + +impl<A: Debug + TrustedRandomAccessNoCoerce, B: Debug + TrustedRandomAccessNoCoerce> ZipFmt<A, B> + for Zip<A, B> +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + // It's *not safe* to call fmt on the contained iterators, since once + // we start iterating they're in strange, potentially unsafe, states. + f.debug_struct("Zip").finish() + } +} + +/// An iterator whose items are random-accessible efficiently +/// +/// # Safety +/// +/// The iterator's `size_hint` must be exact and cheap to call. +/// +/// `TrustedRandomAccessNoCoerce::size` may not be overridden. +/// +/// All subtypes and all supertypes of `Self` must also implement `TrustedRandomAccess`. +/// In particular, this means that types with non-invariant parameters usually can not have +/// an impl for `TrustedRandomAccess` that depends on any trait bounds on such parameters, except +/// for bounds that come from the respective struct/enum definition itself, or bounds involving +/// traits that themselves come with a guarantee similar to this one. +/// +/// If `Self: ExactSizeIterator` then `self.len()` must always produce results consistent +/// with `self.size()`. +/// +/// If `Self: Iterator`, then `<Self as Iterator>::__iterator_get_unchecked(&mut self, idx)` +/// must be safe to call provided the following conditions are met. +/// +/// 1. `0 <= idx` and `idx < self.size()`. +/// 2. If `Self: !Clone`, then `self.__iterator_get_unchecked(idx)` is never called with the same +/// index on `self` more than once. +/// 3. After `self.__iterator_get_unchecked(idx)` has been called, then `self.next_back()` will +/// only be called at most `self.size() - idx - 1` times. If `Self: Clone` and `self` is cloned, +/// then this number is calculated for `self` and its clone individually, +/// but `self.next_back()` calls that happened before the cloning count for both `self` and the clone. +/// 4. After `self.__iterator_get_unchecked(idx)` has been called, then only the following methods +/// will be called on `self` or on any new clones of `self`: +/// * `std::clone::Clone::clone` +/// * `std::iter::Iterator::size_hint` +/// * `std::iter::DoubleEndedIterator::next_back` +/// * `std::iter::ExactSizeIterator::len` +/// * `std::iter::Iterator::__iterator_get_unchecked` +/// * `std::iter::TrustedRandomAccessNoCoerce::size` +/// 5. If `T` is a subtype of `Self`, then `self` is allowed to be coerced +/// to `T`. If `self` is coerced to `T` after `self.__iterator_get_unchecked(idx)` has already +/// been called, then no methods except for the ones listed under 4. are allowed to be called +/// on the resulting value of type `T`, either. Multiple such coercion steps are allowed. +/// Regarding 2. and 3., the number of times `__iterator_get_unchecked(idx)` or `next_back()` is +/// called on `self` and the resulting value of type `T` (and on further coercion results with +/// sub-subtypes) are added together and their sums must not exceed the specified bounds. +/// +/// Further, given that these conditions are met, it must guarantee that: +/// +/// * It does not change the value returned from `size_hint` +/// * It must be safe to call the methods listed above on `self` after calling +/// `self.__iterator_get_unchecked(idx)`, assuming that the required traits are implemented. +/// * It must also be safe to drop `self` after calling `self.__iterator_get_unchecked(idx)`. +/// * If `T` is a subtype of `Self`, then it must be safe to coerce `self` to `T`. +// +// FIXME: Clarify interaction with SourceIter/InPlaceIterable. Calling `SourceIter::as_inner` +// after `__iterator_get_unchecked` is supposed to be allowed. +#[doc(hidden)] +#[unstable(feature = "trusted_random_access", issue = "none")] +#[rustc_specialization_trait] +pub unsafe trait TrustedRandomAccess: TrustedRandomAccessNoCoerce {} + +/// Like [`TrustedRandomAccess`] but without any of the requirements / guarantees around +/// coercions to subtypes after `__iterator_get_unchecked` (they aren’t allowed here!), and +/// without the requirement that subtypes / supertypes implement `TrustedRandomAccessNoCoerce`. +/// +/// This trait was created in PR #85874 to fix soundness issue #85873 without performance regressions. +/// It is subject to change as we might want to build a more generally useful (for performance +/// optimizations) and more sophisticated trait or trait hierarchy that replaces or extends +/// [`TrustedRandomAccess`] and `TrustedRandomAccessNoCoerce`. +#[doc(hidden)] +#[unstable(feature = "trusted_random_access", issue = "none")] +#[rustc_specialization_trait] +pub unsafe trait TrustedRandomAccessNoCoerce: Sized { + // Convenience method. + fn size(&self) -> usize + where + Self: Iterator, + { + self.size_hint().0 + } + /// `true` if getting an iterator element may have side effects. + /// Remember to take inner iterators into account. + const MAY_HAVE_SIDE_EFFECT: bool; +} + +/// Like `Iterator::__iterator_get_unchecked`, but doesn't require the compiler to +/// know that `U: TrustedRandomAccess`. +/// +/// ## Safety +/// +/// Same requirements calling `get_unchecked` directly. +#[doc(hidden)] +#[inline] +pub(in crate::iter::adapters) unsafe fn try_get_unchecked<I>(it: &mut I, idx: usize) -> I::Item +where + I: Iterator, +{ + // SAFETY: the caller must uphold the contract for + // `Iterator::__iterator_get_unchecked`. + unsafe { it.try_get_unchecked(idx) } +} + +unsafe trait SpecTrustedRandomAccess: Iterator { + /// If `Self: TrustedRandomAccess`, it must be safe to call + /// `Iterator::__iterator_get_unchecked(self, index)`. + unsafe fn try_get_unchecked(&mut self, index: usize) -> Self::Item; +} + +unsafe impl<I: Iterator> SpecTrustedRandomAccess for I { + default unsafe fn try_get_unchecked(&mut self, _: usize) -> Self::Item { + panic!("Should only be called on TrustedRandomAccess iterators"); + } +} + +unsafe impl<I: Iterator + TrustedRandomAccessNoCoerce> SpecTrustedRandomAccess for I { + #[inline] + unsafe fn try_get_unchecked(&mut self, index: usize) -> Self::Item { + // SAFETY: the caller must uphold the contract for + // `Iterator::__iterator_get_unchecked`. + unsafe { self.__iterator_get_unchecked(index) } + } +} |