Skip to content

Refactor iter adapters with less macros #98430

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jun 28, 2022
Merged

Refactor iter adapters with less macros #98430

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
72 changes: 19 additions & 53 deletions library/core/src/iter/adapters/chain.rs
View file
Open in desktop
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
Open in desktop
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
Open in desktop
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
}