Skip to content

Commit

Permalink
Refactor iter adapters with less macros
Browse files Browse the repository at this point in the history
  • Loading branch information
@camsteffen
camsteffen committed Jun 22, 2022
1 parent 10f4ce3 commit 6587dda
Show file tree
Hide file tree
Showing 3 changed files with 56 additions and 117 deletions.
72 changes: 19 additions & 53 deletions library/core/src/iter/adapters/chain.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,33 +37,6 @@ impl<A, 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")]
impl<A, B> Iterator for Chain<A, B>
where
Expand All @@ -74,10 +47,7 @@ where

#[inline]
fn next(&mut self) -> Option<A::Item> {
match fuse!(self.a.next()) {
None => maybe!(self.b.next()),
item => item,
}
and_then_or_clear(&mut self.a, Iterator::next).or_else(|| self.b.as_mut()?.next())
}

#[inline]
Expand Down Expand Up @@ -161,31 +131,23 @@ where
self.a = None;
}

maybe!(self.b.nth(n))
self.b.as_mut()?.nth(n)
}

#[inline]
fn find<P>(&mut self, mut predicate: P) -> Option<Self::Item>
where
P: FnMut(&Self::Item) -> bool,
{
match fuse!(self.a.find(&mut predicate)) {
None => maybe!(self.b.find(predicate)),
item => item,
}
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 = match self.a {
Some(a) => a.last(),
None => None,
};
let b_last = match self.b {
Some(b) => b.last(),
None => None,
};
let a_last = self.a.and_then(Iterator::last);
let b_last = self.b.and_then(Iterator::last);
b_last.or(a_last)
}

Expand Down Expand Up @@ -220,10 +182,7 @@ where
{
#[inline]
fn next_back(&mut self) -> Option<A::Item> {
match fuse!(self.b.next_back()) {
None => maybe!(self.a.next_back()),
item => item,
}
and_then_or_clear(&mut self.b, |b| b.next_back()).or_else(|| self.a.as_mut()?.next_back())
}

#[inline]
Expand Down Expand Up @@ -263,18 +222,16 @@ where
self.b = None;
}

maybe!(self.a.nth_back(n))
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,
{
match fuse!(self.b.rfind(&mut predicate)) {
None => maybe!(self.a.rfind(predicate)),
item => item,
}
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
Expand Down Expand Up @@ -324,3 +281,12 @@ where
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
}
41 changes: 16 additions & 25 deletions library/core/src/iter/adapters/flatten.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -290,20 +290,11 @@ where
#[inline]
fn next(&mut self) -> Option<U::Item> {
loop {
if let Some(ref mut inner) = self.frontiter {
match inner.next() {
None => self.frontiter = None,
elt @ Some(_) => return elt,
}
if let elt @ Some(_) = and_then_or_clear(&mut self.frontiter, Iterator::next) {
return elt;
}
match self.iter.next() {
None => match self.backiter.as_mut()?.next() {
None => {
self.backiter = None;
return None;
}
elt @ Some(_) => return elt,
},
None => return and_then_or_clear(&mut self.backiter, Iterator::next),
Some(inner) => self.frontiter = Some(inner.into_iter()),
}
}
Expand Down Expand Up @@ -436,21 +427,12 @@ where
#[inline]
fn next_back(&mut self) -> Option<U::Item> {
loop {
if let Some(ref mut inner) = self.backiter {
match inner.next_back() {
None => self.backiter = None,
elt @ Some(_) => return elt,
}
if let elt @ Some(_) = and_then_or_clear(&mut self.backiter, |b| b.next_back()) {
return elt;
}
match self.iter.next_back() {
None => match self.frontiter.as_mut()?.next_back() {
None => {
self.frontiter = None;
return None;
}
elt @ Some(_) => return elt,
},
next => self.backiter = next.map(IntoIterator::into_iter),
None => return and_then_or_clear(&mut self.frontiter, |f| f.next_back()),
Some(inner) => self.backiter = Some(inner.into_iter()),
}
}
}
Expand Down Expand Up @@ -606,3 +588,12 @@ 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")]
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
}
60 changes: 21 additions & 39 deletions library/core/src/iter/adapters/fuse.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,33 +29,6 @@ impl<I> Fuse<I> {
#[stable(feature = "fused", since = "1.26.0")]
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
// to avoid exposing default fns outside this trait.
#[stable(feature = "rust1", since = "1.0.0")]
Expand Down Expand Up @@ -281,12 +254,12 @@ where

#[inline]
default fn next(&mut self) -> Option<<I as Iterator>::Item> {
fuse!(self.iter.next())
and_then_or_clear(&mut self.iter, Iterator::next)
}

#[inline]
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]
Expand All @@ -308,23 +281,23 @@ where
where
P: FnMut(&Self::Item) -> bool,
{
fuse!(self.iter.find(predicate))
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,
{
fuse!(self.iter.next_back())
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,
{
fuse!(self.iter.nth_back(n))
and_then_or_clear(&mut self.iter, |iter| iter.nth_back(n))
}

#[inline]
Expand All @@ -348,7 +321,7 @@ where
P: FnMut(&Self::Item) -> bool,
I: DoubleEndedIterator,
{
fuse!(self.iter.rfind(predicate))
and_then_or_clear(&mut self.iter, |iter| iter.rfind(predicate))
}
}

Expand All @@ -361,12 +334,12 @@ where
{
#[inline]
fn next(&mut self) -> Option<<I as Iterator>::Item> {
spec!(self.iter.next())
self.iter.as_mut()?.next()
}

#[inline]
fn nth(&mut self, n: usize) -> Option<I::Item> {
spec!(self.iter.nth(n))
self.iter.as_mut()?.nth(n)
}

#[inline]
Expand All @@ -387,23 +360,23 @@ where
where
P: FnMut(&Self::Item) -> bool,
{
spec!(self.iter.find(predicate))
self.iter.as_mut()?.find(predicate)
}

#[inline]
fn next_back(&mut self) -> Option<<I as Iterator>::Item>
where
I: DoubleEndedIterator,
{
spec!(self.iter.next_back())
self.iter.as_mut()?.next_back()
}

#[inline]
fn nth_back(&mut self, n: usize) -> Option<<I as Iterator>::Item>
where
I: DoubleEndedIterator,
{
spec!(self.iter.nth_back(n))
self.iter.as_mut()?.nth_back(n)
}

#[inline]
Expand All @@ -426,6 +399,15 @@ where
P: FnMut(&Self::Item) -> bool,
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
}

0 comments on commit 6587dda

Please sign in to comment.