diff --git a/src/libcore/iter/adapters/chain.rs b/src/libcore/iter/adapters/chain.rs index 3611a1aadaddb..0087efc2bab38 100644 --- a/src/libcore/iter/adapters/chain.rs +++ b/src/libcore/iter/adapters/chain.rs @@ -1,8 +1,7 @@ +use crate::iter::{DoubleEndedIterator, Fuse, FusedIterator, Iterator, TrustedLen}; use crate::ops::Try; use crate::usize; -use super::super::{DoubleEndedIterator, FusedIterator, Iterator, TrustedLen}; - /// An iterator that links two iterators together, in a chain. /// /// This `struct` is created by the [`chain`] method on [`Iterator`]. See its @@ -14,39 +13,19 @@ use super::super::{DoubleEndedIterator, FusedIterator, Iterator, TrustedLen}; #[must_use = "iterators are lazy and do nothing unless consumed"] #[stable(feature = "rust1", since = "1.0.0")] pub struct Chain { - a: A, - b: B, - state: ChainState, + // These are `Fuse`d so we don't have to manually track which part is already exhausted. + // However, the `Defuse` wrapper hides the `FusedIterator` implementation, so we don't + // use the specialized `Fuse` that unconditionally descends into the iterator, because + // that could be expensive to keep revisiting stuff like nested chains. + a: Fuse>, + b: Fuse>, } impl Chain { pub(in super::super) fn new(a: A, b: B) -> Chain { - Chain { a, b, state: ChainState::Both } + Chain { a: Fuse::new(Defuse(a)), b: Fuse::new(Defuse(b)) } } } -// The iterator protocol specifies that iteration ends with the return value -// `None` from `.next()` (or `.next_back()`) and it is unspecified what -// further calls return. The chain adaptor must account for this since it uses -// two subiterators. -// -// It uses three states: -// -// - Both: `a` and `b` are remaining -// - Front: `a` remaining -// - Back: `b` remaining -// -// The fourth state (neither iterator is remaining) only occurs after Chain has -// returned None once, so we don't need to store this state. -#[derive(Clone, Debug)] -enum ChainState { - // both front and back iterator are remaining - Both, - // only front is remaining - Front, - // only back is remaining - Back, -} - #[stable(feature = "rust1", since = "1.0.0")] impl Iterator for Chain where @@ -57,27 +36,16 @@ where #[inline] fn next(&mut self) -> Option { - match self.state { - ChainState::Both => match self.a.next() { - elt @ Some(..) => elt, - None => { - self.state = ChainState::Back; - self.b.next() - } - }, - ChainState::Front => self.a.next(), - ChainState::Back => self.b.next(), + match self.a.next() { + None => self.b.next(), + item => item, } } #[inline] #[rustc_inherit_overflow_checks] fn count(self) -> usize { - match self.state { - ChainState::Both => self.a.count() + self.b.count(), - ChainState::Front => self.a.count(), - ChainState::Back => self.b.count(), - } + self.a.count() + self.b.count() } fn try_fold(&mut self, init: Acc, mut f: F) -> R @@ -86,59 +54,27 @@ where F: FnMut(Acc, Self::Item) -> R, R: Try, { - let mut accum = init; - match self.state { - ChainState::Both | ChainState::Front => { - accum = self.a.try_fold(accum, &mut f)?; - if let ChainState::Both = self.state { - self.state = ChainState::Back; - } - } - _ => {} - } - if let ChainState::Back = self.state { - accum = self.b.try_fold(accum, &mut f)?; - } - Try::from_ok(accum) + let accum = self.a.try_fold(init, &mut f)?; + self.b.try_fold(accum, f) } fn fold(self, init: Acc, mut f: F) -> Acc where F: FnMut(Acc, Self::Item) -> Acc, { - let mut accum = init; - match self.state { - ChainState::Both | ChainState::Front => { - accum = self.a.fold(accum, &mut f); - } - _ => {} - } - match self.state { - ChainState::Both | ChainState::Back => { - accum = self.b.fold(accum, &mut f); - } - _ => {} - } - accum + let accum = self.a.fold(init, &mut f); + self.b.fold(accum, f) } #[inline] fn nth(&mut self, mut n: usize) -> Option { - match self.state { - ChainState::Both | ChainState::Front => { - for x in self.a.by_ref() { - if n == 0 { - return Some(x); - } - n -= 1; - } - if let ChainState::Both = self.state { - self.state = ChainState::Back; - } + for x in self.a.by_ref() { + if n == 0 { + return Some(x); } - ChainState::Back => {} + n -= 1; } - if let ChainState::Back = self.state { self.b.nth(n) } else { None } + self.b.nth(n) } #[inline] @@ -146,52 +82,33 @@ where where P: FnMut(&Self::Item) -> bool, { - match self.state { - ChainState::Both => match self.a.find(&mut predicate) { - None => { - self.state = ChainState::Back; - self.b.find(predicate) - } - v => v, - }, - ChainState::Front => self.a.find(predicate), - ChainState::Back => self.b.find(predicate), + match self.a.find(&mut predicate) { + None => self.b.find(predicate), + item => item, } } #[inline] fn last(self) -> Option { - match self.state { - ChainState::Both => { - // Must exhaust a before b. - let a_last = self.a.last(); - let b_last = self.b.last(); - b_last.or(a_last) - } - ChainState::Front => self.a.last(), - ChainState::Back => self.b.last(), - } + // Must exhaust a before b. + let a_last = self.a.last(); + let b_last = self.b.last(); + b_last.or(a_last) } #[inline] fn size_hint(&self) -> (usize, Option) { - match self.state { - ChainState::Both => { - let (a_lower, a_upper) = self.a.size_hint(); - let (b_lower, b_upper) = self.b.size_hint(); + let (a_lower, a_upper) = self.a.size_hint(); + let (b_lower, b_upper) = self.b.size_hint(); - let lower = a_lower.saturating_add(b_lower); + let lower = a_lower.saturating_add(b_lower); - let upper = match (a_upper, b_upper) { - (Some(x), Some(y)) => x.checked_add(y), - _ => None, - }; + let upper = match (a_upper, b_upper) { + (Some(x), Some(y)) => x.checked_add(y), + _ => None, + }; - (lower, upper) - } - ChainState::Front => self.a.size_hint(), - ChainState::Back => self.b.size_hint(), - } + (lower, upper) } } @@ -203,36 +120,32 @@ where { #[inline] fn next_back(&mut self) -> Option { - match self.state { - ChainState::Both => match self.b.next_back() { - elt @ Some(..) => elt, - None => { - self.state = ChainState::Front; - self.a.next_back() - } - }, - ChainState::Front => self.a.next_back(), - ChainState::Back => self.b.next_back(), + match self.b.next_back() { + None => self.a.next_back(), + item => item, } } #[inline] fn nth_back(&mut self, mut n: usize) -> Option { - match self.state { - ChainState::Both | ChainState::Back => { - for x in self.b.by_ref().rev() { - if n == 0 { - return Some(x); - } - n -= 1; - } - if let ChainState::Both = self.state { - self.state = ChainState::Front; - } + for x in self.b.by_ref().rev() { + if n == 0 { + return Some(x); } - ChainState::Front => {} + n -= 1; + } + self.a.nth_back(n) + } + + #[inline] + fn rfind

(&mut self, mut predicate: P) -> Option + where + P: FnMut(&Self::Item) -> bool, + { + match self.b.rfind(&mut predicate) { + None => self.a.rfind(predicate), + item => item, } - if let ChainState::Front = self.state { self.a.nth_back(n) } else { None } } fn try_rfold(&mut self, init: Acc, mut f: F) -> R @@ -241,44 +154,22 @@ where F: FnMut(Acc, Self::Item) -> R, R: Try, { - let mut accum = init; - match self.state { - ChainState::Both | ChainState::Back => { - accum = self.b.try_rfold(accum, &mut f)?; - if let ChainState::Both = self.state { - self.state = ChainState::Front; - } - } - _ => {} - } - if let ChainState::Front = self.state { - accum = self.a.try_rfold(accum, &mut f)?; - } - Try::from_ok(accum) + let accum = self.b.try_rfold(init, &mut f)?; + self.a.try_rfold(accum, f) } fn rfold(self, init: Acc, mut f: F) -> Acc where F: FnMut(Acc, Self::Item) -> Acc, { - let mut accum = init; - match self.state { - ChainState::Both | ChainState::Back => { - accum = self.b.rfold(accum, &mut f); - } - _ => {} - } - match self.state { - ChainState::Both | ChainState::Front => { - accum = self.a.rfold(accum, &mut f); - } - _ => {} - } - accum + let accum = self.b.rfold(init, &mut f); + self.a.rfold(accum, f) } } // Note: *both* must be fused to handle double-ended iterators. +// Now that we `Fuse` both sides, we *could* implement this unconditionally, +// but we should be cautious about committing to that in the public API. #[stable(feature = "fused", since = "1.26.0")] impl FusedIterator for Chain where @@ -294,3 +185,100 @@ where B: TrustedLen, { } + +/// Wrapper that forwards everything but `FusedIterator`. +#[derive(Clone, Debug)] +struct Defuse(I); + +impl Iterator for Defuse { + type Item = I::Item; + + #[inline] + fn next(&mut self) -> Option { + self.0.next() + } + + #[inline] + fn count(self) -> usize { + self.0.count() + } + + #[inline] + fn try_fold(&mut self, init: Acc, f: F) -> R + where + Self: Sized, + F: FnMut(Acc, I::Item) -> R, + R: Try, + { + self.0.try_fold(init, f) + } + + #[inline] + fn fold(self, init: Acc, f: F) -> Acc + where + F: FnMut(Acc, I::Item) -> Acc, + { + self.0.fold(init, f) + } + + #[inline] + fn nth(&mut self, n: usize) -> Option { + self.0.nth(n) + } + + #[inline] + fn find

(&mut self, predicate: P) -> Option + where + P: FnMut(&I::Item) -> bool, + { + self.0.find(predicate) + } + + #[inline] + fn last(self) -> Option { + self.0.last() + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + self.0.size_hint() + } +} + +impl DoubleEndedIterator for Defuse { + #[inline] + fn next_back(&mut self) -> Option { + self.0.next_back() + } + + #[inline] + fn nth_back(&mut self, n: usize) -> Option { + self.0.nth_back(n) + } + + #[inline] + fn rfind

(&mut self, predicate: P) -> Option + where + P: FnMut(&I::Item) -> bool, + { + self.0.rfind(predicate) + } + + #[inline] + fn try_rfold(&mut self, init: Acc, f: F) -> R + where + Self: Sized, + F: FnMut(Acc, I::Item) -> R, + R: Try, + { + self.0.try_rfold(init, f) + } + + #[inline] + fn rfold(self, init: Acc, f: F) -> Acc + where + F: FnMut(Acc, I::Item) -> Acc, + { + self.0.rfold(init, f) + } +}