Refactor iter adapters with less macros

This commit is contained in:
Cameron Steffen 2022-06-22 10:56:26 -05:00
parent 10f4ce324b
commit 6587dda39e
3 changed files with 56 additions and 117 deletions

View File

@ -37,33 +37,6 @@ pub(in super::super) fn new(a: A, b: B) -> Chain<A, B> {
} }
} }
/// Fuse the iterator if the expression is `None`.
macro_rules! fuse {
($self:ident . $iter:ident . $($call:tt)+) => {
match $self.$iter {
Some(ref mut iter) => match iter.$($call)+ {
None => {
$self.$iter = None;
None
}
item => item,
},
None => None,
}
};
}
/// Try an iterator method without fusing,
/// like an inline `.as_mut().and_then(...)`
macro_rules! maybe {
($self:ident . $iter:ident . $($call:tt)+) => {
match $self.$iter {
Some(ref mut iter) => iter.$($call)+,
None => None,
}
};
}
#[stable(feature = "rust1", since = "1.0.0")] #[stable(feature = "rust1", since = "1.0.0")]
impl<A, B> Iterator for Chain<A, B> impl<A, B> Iterator for Chain<A, B>
where where
@ -74,10 +47,7 @@ impl<A, B> Iterator for Chain<A, B>
#[inline] #[inline]
fn next(&mut self) -> Option<A::Item> { fn next(&mut self) -> Option<A::Item> {
match fuse!(self.a.next()) { and_then_or_clear(&mut self.a, Iterator::next).or_else(|| self.b.as_mut()?.next())
None => maybe!(self.b.next()),
item => item,
}
} }
#[inline] #[inline]
@ -161,7 +131,7 @@ fn nth(&mut self, mut n: usize) -> Option<Self::Item> {
self.a = None; self.a = None;
} }
maybe!(self.b.nth(n)) self.b.as_mut()?.nth(n)
} }
#[inline] #[inline]
@ -169,23 +139,15 @@ fn find<P>(&mut self, mut predicate: P) -> Option<Self::Item>
where where
P: FnMut(&Self::Item) -> bool, P: FnMut(&Self::Item) -> bool,
{ {
match fuse!(self.a.find(&mut predicate)) { and_then_or_clear(&mut self.a, |a| a.find(&mut predicate))
None => maybe!(self.b.find(predicate)), .or_else(|| self.b.as_mut()?.find(predicate))
item => item,
}
} }
#[inline] #[inline]
fn last(self) -> Option<A::Item> { fn last(self) -> Option<A::Item> {
// Must exhaust a before b. // Must exhaust a before b.
let a_last = match self.a { let a_last = self.a.and_then(Iterator::last);
Some(a) => a.last(), let b_last = self.b.and_then(Iterator::last);
None => None,
};
let b_last = match self.b {
Some(b) => b.last(),
None => None,
};
b_last.or(a_last) b_last.or(a_last)
} }
@ -220,10 +182,7 @@ impl<A, B> DoubleEndedIterator for Chain<A, B>
{ {
#[inline] #[inline]
fn next_back(&mut self) -> Option<A::Item> { fn next_back(&mut self) -> Option<A::Item> {
match fuse!(self.b.next_back()) { and_then_or_clear(&mut self.b, |b| b.next_back()).or_else(|| self.a.as_mut()?.next_back())
None => maybe!(self.a.next_back()),
item => item,
}
} }
#[inline] #[inline]
@ -263,7 +222,7 @@ fn nth_back(&mut self, mut n: usize) -> Option<Self::Item> {
self.b = None; self.b = None;
} }
maybe!(self.a.nth_back(n)) self.a.as_mut()?.nth_back(n)
} }
#[inline] #[inline]
@ -271,10 +230,8 @@ fn rfind<P>(&mut self, mut predicate: P) -> Option<Self::Item>
where where
P: FnMut(&Self::Item) -> bool, P: FnMut(&Self::Item) -> bool,
{ {
match fuse!(self.b.rfind(&mut predicate)) { and_then_or_clear(&mut self.b, |b| b.rfind(&mut predicate))
None => maybe!(self.a.rfind(predicate)), .or_else(|| self.a.as_mut()?.rfind(predicate))
item => item,
}
} }
fn try_rfold<Acc, F, R>(&mut self, mut acc: Acc, mut f: F) -> R fn try_rfold<Acc, F, R>(&mut self, mut acc: Acc, mut f: F) -> R
@ -324,3 +281,12 @@ unsafe impl<A, B> TrustedLen for Chain<A, B>
B: TrustedLen<Item = A::Item>, 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
}

View File

@ -290,20 +290,11 @@ impl<I, U> Iterator for FlattenCompat<I, U>
#[inline] #[inline]
fn next(&mut self) -> Option<U::Item> { fn next(&mut self) -> Option<U::Item> {
loop { loop {
if let Some(ref mut inner) = self.frontiter { if let elt @ Some(_) = and_then_or_clear(&mut self.frontiter, Iterator::next) {
match inner.next() { return elt;
None => self.frontiter = None,
elt @ Some(_) => return elt,
}
} }
match self.iter.next() { match self.iter.next() {
None => match self.backiter.as_mut()?.next() { None => return and_then_or_clear(&mut self.backiter, Iterator::next),
None => {
self.backiter = None;
return None;
}
elt @ Some(_) => return elt,
},
Some(inner) => self.frontiter = Some(inner.into_iter()), Some(inner) => self.frontiter = Some(inner.into_iter()),
} }
} }
@ -436,21 +427,12 @@ impl<I, U> DoubleEndedIterator for FlattenCompat<I, U>
#[inline] #[inline]
fn next_back(&mut self) -> Option<U::Item> { fn next_back(&mut self) -> Option<U::Item> {
loop { loop {
if let Some(ref mut inner) = self.backiter { if let elt @ Some(_) = and_then_or_clear(&mut self.backiter, |b| b.next_back()) {
match inner.next_back() { return elt;
None => self.backiter = None,
elt @ Some(_) => return elt,
}
} }
match self.iter.next_back() { match self.iter.next_back() {
None => match self.frontiter.as_mut()?.next_back() { None => return and_then_or_clear(&mut self.frontiter, |f| f.next_back()),
None => { Some(inner) => self.backiter = Some(inner.into_iter()),
self.frontiter = None;
return None;
}
elt @ Some(_) => return elt,
},
next => self.backiter = next.map(IntoIterator::into_iter),
} }
} }
} }
@ -606,3 +588,12 @@ pub unsafe trait TrustedConstSize: IntoIterator {}
unsafe impl<T, const N: usize> TrustedConstSize for &'_ [T; N] {} unsafe impl<T, const N: usize> TrustedConstSize for &'_ [T; N] {}
#[unstable(feature = "std_internals", issue = "none")] #[unstable(feature = "std_internals", issue = "none")]
unsafe impl<T, const N: usize> TrustedConstSize for &'_ mut [T; N] {} 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
}

View File

@ -29,33 +29,6 @@ pub(in crate::iter) fn new(iter: I) -> Fuse<I> {
#[stable(feature = "fused", since = "1.26.0")] #[stable(feature = "fused", since = "1.26.0")]
impl<I> FusedIterator for Fuse<I> where I: Iterator {} impl<I> FusedIterator for Fuse<I> where I: Iterator {}
/// Fuse the iterator if the expression is `None`.
macro_rules! fuse {
($self:ident . iter . $($call:tt)+) => {
match $self.iter {
Some(ref mut iter) => match iter.$($call)+ {
None => {
$self.iter = None;
None
}
item => item,
},
None => None,
}
};
}
/// Specialized macro that doesn't check if the expression is `None`.
/// (We trust that a `FusedIterator` will fuse itself.)
macro_rules! spec {
($self:ident . iter . $($call:tt)+) => {
match $self.iter {
Some(ref mut iter) => iter.$($call)+,
None => None,
}
};
}
// Any specialized implementation here is made internal // Any specialized implementation here is made internal
// to avoid exposing default fns outside this trait. // to avoid exposing default fns outside this trait.
#[stable(feature = "rust1", since = "1.0.0")] #[stable(feature = "rust1", since = "1.0.0")]
@ -281,12 +254,12 @@ impl<I> FuseImpl<I> for Fuse<I>
#[inline] #[inline]
default fn next(&mut self) -> Option<<I as Iterator>::Item> { default fn next(&mut self) -> Option<<I as Iterator>::Item> {
fuse!(self.iter.next()) and_then_or_clear(&mut self.iter, Iterator::next)
} }
#[inline] #[inline]
default fn nth(&mut self, n: usize) -> Option<I::Item> { default fn nth(&mut self, n: usize) -> Option<I::Item> {
fuse!(self.iter.nth(n)) and_then_or_clear(&mut self.iter, |iter| iter.nth(n))
} }
#[inline] #[inline]
@ -308,7 +281,7 @@ impl<I> FuseImpl<I> for Fuse<I>
where where
P: FnMut(&Self::Item) -> bool, P: FnMut(&Self::Item) -> bool,
{ {
fuse!(self.iter.find(predicate)) and_then_or_clear(&mut self.iter, |iter| iter.find(predicate))
} }
#[inline] #[inline]
@ -316,7 +289,7 @@ impl<I> FuseImpl<I> for Fuse<I>
where where
I: DoubleEndedIterator, I: DoubleEndedIterator,
{ {
fuse!(self.iter.next_back()) and_then_or_clear(&mut self.iter, |iter| iter.next_back())
} }
#[inline] #[inline]
@ -324,7 +297,7 @@ impl<I> FuseImpl<I> for Fuse<I>
where where
I: DoubleEndedIterator, I: DoubleEndedIterator,
{ {
fuse!(self.iter.nth_back(n)) and_then_or_clear(&mut self.iter, |iter| iter.nth_back(n))
} }
#[inline] #[inline]
@ -348,7 +321,7 @@ impl<I> FuseImpl<I> for Fuse<I>
P: FnMut(&Self::Item) -> bool, P: FnMut(&Self::Item) -> bool,
I: DoubleEndedIterator, I: DoubleEndedIterator,
{ {
fuse!(self.iter.rfind(predicate)) and_then_or_clear(&mut self.iter, |iter| iter.rfind(predicate))
} }
} }
@ -361,12 +334,12 @@ impl<I> FuseImpl<I> for Fuse<I>
{ {
#[inline] #[inline]
fn next(&mut self) -> Option<<I as Iterator>::Item> { fn next(&mut self) -> Option<<I as Iterator>::Item> {
spec!(self.iter.next()) self.iter.as_mut()?.next()
} }
#[inline] #[inline]
fn nth(&mut self, n: usize) -> Option<I::Item> { fn nth(&mut self, n: usize) -> Option<I::Item> {
spec!(self.iter.nth(n)) self.iter.as_mut()?.nth(n)
} }
#[inline] #[inline]
@ -387,7 +360,7 @@ fn find<P>(&mut self, predicate: P) -> Option<Self::Item>
where where
P: FnMut(&Self::Item) -> bool, P: FnMut(&Self::Item) -> bool,
{ {
spec!(self.iter.find(predicate)) self.iter.as_mut()?.find(predicate)
} }
#[inline] #[inline]
@ -395,7 +368,7 @@ fn next_back(&mut self) -> Option<<I as Iterator>::Item>
where where
I: DoubleEndedIterator, I: DoubleEndedIterator,
{ {
spec!(self.iter.next_back()) self.iter.as_mut()?.next_back()
} }
#[inline] #[inline]
@ -403,7 +376,7 @@ fn nth_back(&mut self, n: usize) -> Option<<I as Iterator>::Item>
where where
I: DoubleEndedIterator, I: DoubleEndedIterator,
{ {
spec!(self.iter.nth_back(n)) self.iter.as_mut()?.nth_back(n)
} }
#[inline] #[inline]
@ -426,6 +399,15 @@ fn rfind<P>(&mut self, predicate: P) -> Option<Self::Item>
P: FnMut(&Self::Item) -> bool, P: FnMut(&Self::Item) -> bool,
I: DoubleEndedIterator, I: DoubleEndedIterator,
{ {
spec!(self.iter.rfind(predicate)) 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
}