Rollup merge of #98430 - camsteffen:flatten-refactor, r=joshtriplett

Refactor iter adapters with less macros

Just some code cleanup. Introduced a util `and_then_or_clear` for each of chain, flatten and fuse iter adapter impls. This reduces code nicely for flatten, but admittedly the other modules are more of a lateral move replacing macros with a function. But I think consistency across the modules and avoiding macros when possible is good.

library/core/src/iter/adapters/chain.rs

@@ -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")]
 impl<A, B> Iterator for Chain<A, B>
 where
@@ -74,10 +47,7 @@ impl<A, B> Iterator for Chain<A, B>
 
     #[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]
@@ -161,7 +131,7 @@ fn nth(&mut self, mut n: usize) -> Option<Self::Item> {
             self.a = None;
         }
 
-        maybe!(self.b.nth(n))
+        self.b.as_mut()?.nth(n)
     }
 
     #[inline]
@@ -169,23 +139,15 @@ 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)
     }
 
@@ -220,10 +182,7 @@ impl<A, B> DoubleEndedIterator for Chain<A, B>
 {
     #[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]
@@ -263,7 +222,7 @@ fn nth_back(&mut self, mut n: usize) -> Option<Self::Item> {
             self.b = None;
         }
 
-        maybe!(self.a.nth_back(n))
+        self.a.as_mut()?.nth_back(n)
     }
 
     #[inline]
@@ -271,10 +230,8 @@ 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
@@ -324,3 +281,12 @@ unsafe impl<A, B> TrustedLen for Chain<A, B>
     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
+}

library/core/src/iter/adapters/flatten.rs

@@ -290,20 +290,11 @@ impl<I, U> Iterator for FlattenCompat<I, U>
     #[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()),
             }
         }
@@ -436,21 +427,12 @@ impl<I, U> DoubleEndedIterator for FlattenCompat<I, U>
     #[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()),
             }
         }
     }
@@ -606,3 +588,12 @@ pub unsafe trait TrustedConstSize: IntoIterator {}
 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
+}

library/core/src/iter/adapters/fuse.rs

@@ -29,33 +29,6 @@ pub(in crate::iter) fn new(iter: 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")]
@@ -281,12 +254,12 @@ impl<I> FuseImpl<I> for Fuse<I>
 
     #[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]
@@ -308,7 +281,7 @@ impl<I> FuseImpl<I> for Fuse<I>
     where
         P: FnMut(&Self::Item) -> bool,
     {
-        fuse!(self.iter.find(predicate))
+        and_then_or_clear(&mut self.iter, |iter| iter.find(predicate))
     }
 
     #[inline]
@@ -316,7 +289,7 @@ impl<I> FuseImpl<I> for Fuse<I>
     where
         I: DoubleEndedIterator,
     {
-        fuse!(self.iter.next_back())
+        and_then_or_clear(&mut self.iter, |iter| iter.next_back())
     }
 
     #[inline]
@@ -324,7 +297,7 @@ impl<I> FuseImpl<I> for Fuse<I>
     where
         I: DoubleEndedIterator,
     {
-        fuse!(self.iter.nth_back(n))
+        and_then_or_clear(&mut self.iter, |iter| iter.nth_back(n))
     }
 
     #[inline]
@@ -348,7 +321,7 @@ impl<I> FuseImpl<I> for Fuse<I>
         P: FnMut(&Self::Item) -> bool,
         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]
     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]
@@ -387,7 +360,7 @@ fn find<P>(&mut self, predicate: P) -> Option<Self::Item>
     where
         P: FnMut(&Self::Item) -> bool,
     {
-        spec!(self.iter.find(predicate))
+        self.iter.as_mut()?.find(predicate)
     }
 
     #[inline]
@@ -395,7 +368,7 @@ 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]
@@ -403,7 +376,7 @@ 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]
@@ -426,6 +399,15 @@ fn rfind<P>(&mut self, predicate: P) -> Option<Self::Item>
         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
 }