Specialize internal state of Fuse depending on underlying iterator

Author: AnthonyMikh

src/libcore/iter/adapters/fuse.rs

@@ -1,9 +1,44 @@
+use crate::convert::Infallible;
 use crate::intrinsics;
 use crate::iter::{
     DoubleEndedIterator, ExactSizeIterator, FusedIterator, Iterator, TrustedRandomAccess,
 };
 use crate::ops::Try;
 
+trait Stop: Sized {
+    /// Switch internal state of `Fuse` to indicate that underlying iterator returned `None`
+    fn stop<I>(state: &mut Result<I, Self>);
+}
+
+impl Stop for () {
+    fn stop<I>(state: &mut Result<I, Self>) {
+        *state = Err(())
+    }
+}
+
+impl Stop for Infallible {
+    #[inline(always)]
+    fn stop<I>(_state: &mut Result<I, Self>) {
+        // Intentionally does nothing: fused iterator returns `None`s after returning `None`,
+        // so there is no need to alter state
+    }
+}
+
+trait StopState {
+    /// Type of value used to indicate that iterator returned `None`
+    type State: Stop + crate::fmt::Debug + Clone;
+}
+
+impl<I: ?Sized> StopState for I {
+    default type State = ();
+}
+
+impl<I: FusedIterator + ?Sized> StopState for I {
+    type State = Infallible;
+}
+
+type StopStateOf<I> = <I as StopState>::State;
+
 /// An iterator that yields `None` forever after the underlying iterator
 /// yields `None` once.
 ///
@@ -16,12 +51,17 @@ use crate::ops::Try;
 #[must_use = "iterators are lazy and do nothing unless consumed"]
 #[stable(feature = "rust1", since = "1.0.0")]
 pub struct Fuse<I> {
-    // NOTE: for `I: FusedIterator`, this is always assumed `Some`!
-    iter: Option<I>,
+    iter: Result<I, StopStateOf<I>>,
 }
+
 impl<I> Fuse<I> {
     pub(in crate::iter) fn new(iter: I) -> Fuse<I> {
-        Fuse { iter: Some(iter) }
+        Fuse { iter: Ok(iter) }
+    }
+
+    #[inline(always)]
+    fn stop(&mut self) {
+        StopStateOf::<I>::stop(&mut self.iter);
     }
 }
 
@@ -36,71 +76,71 @@ where
     type Item = <I as Iterator>::Item;
 
     #[inline]
-    default fn next(&mut self) -> Option<<I as Iterator>::Item> {
-        let next = self.iter.as_mut()?.next();
+    fn next(&mut self) -> Option<<I as Iterator>::Item> {
+        let next = self.iter.as_mut().ok()?.next();
         if next.is_none() {
-            self.iter = None;
+            self.stop();
         }
         next
     }
 
     #[inline]
-    default fn nth(&mut self, n: usize) -> Option<I::Item> {
-        let nth = self.iter.as_mut()?.nth(n);
+    fn nth(&mut self, n: usize) -> Option<I::Item> {
+        let nth = self.iter.as_mut().ok()?.nth(n);
         if nth.is_none() {
-            self.iter = None;
+            self.stop();
         }
         nth
     }
 
     #[inline]
-    default fn last(self) -> Option<I::Item> {
-        self.iter?.last()
+    fn last(self) -> Option<I::Item> {
+        self.iter.ok()?.last()
     }
 
     #[inline]
-    default fn count(self) -> usize {
+    fn count(self) -> usize {
         self.iter.map_or(0, I::count)
     }
 
     #[inline]
-    default fn size_hint(&self) -> (usize, Option<usize>) {
+    fn size_hint(&self) -> (usize, Option<usize>) {
         self.iter.as_ref().map_or((0, Some(0)), I::size_hint)
     }
 
     #[inline]
-    default fn try_fold<Acc, Fold, R>(&mut self, mut acc: Acc, fold: Fold) -> R
+    fn try_fold<Acc, Fold, R>(&mut self, mut acc: Acc, fold: Fold) -> R
     where
         Self: Sized,
         Fold: FnMut(Acc, Self::Item) -> R,
         R: Try<Ok = Acc>,
     {
-        if let Some(ref mut iter) = self.iter {
+        if let Ok(ref mut iter) = self.iter {
             acc = iter.try_fold(acc, fold)?;
-            self.iter = None;
+            self.stop();
         }
         Try::from_ok(acc)
     }
 
     #[inline]
-    default fn fold<Acc, Fold>(self, mut acc: Acc, fold: Fold) -> Acc
+    fn fold<Acc, Fold>(self, mut acc: Acc, fold: Fold) -> Acc
     where
         Fold: FnMut(Acc, Self::Item) -> Acc,
     {
-        if let Some(iter) = self.iter {
+        if let Ok(iter) = self.iter {
             acc = iter.fold(acc, fold);
         }
         acc
     }
 
     #[inline]
-    default fn find<P>(&mut self, predicate: P) -> Option<Self::Item>
+    fn find<P>(&mut self, predicate: P) -> Option<Self::Item>
     where
         P: FnMut(&Self::Item) -> bool,
     {
-        let found = self.iter.as_mut()?.find(predicate);
+        let found = self.iter.as_mut().ok()?.find(predicate);
         if found.is_none() {
-            self.iter = None;
+            self.stop();
         }
         found
     }
@@ -112,56 +152,56 @@ where
     I: DoubleEndedIterator,
 {
     #[inline]
-    default fn next_back(&mut self) -> Option<<I as Iterator>::Item> {
-        let next = self.iter.as_mut()?.next_back();
+    fn next_back(&mut self) -> Option<<I as Iterator>::Item> {
+        let next = self.iter.as_mut().ok()?.next_back();
         if next.is_none() {
-            self.iter = None;
+            self.stop();
         }
         next
     }
 
     #[inline]
-    default fn nth_back(&mut self, n: usize) -> Option<<I as Iterator>::Item> {
-        let nth = self.iter.as_mut()?.nth_back(n);
+    fn nth_back(&mut self, n: usize) -> Option<<I as Iterator>::Item> {
+        let nth = self.iter.as_mut().ok()?.nth_back(n);
         if nth.is_none() {
-            self.iter = None;
+            self.stop();
         }
         nth
     }
 
     #[inline]
-    default fn try_rfold<Acc, Fold, R>(&mut self, mut acc: Acc, fold: Fold) -> R
+    fn try_rfold<Acc, Fold, R>(&mut self, mut acc: Acc, fold: Fold) -> R
     where
         Self: Sized,
         Fold: FnMut(Acc, Self::Item) -> R,
         R: Try<Ok = Acc>,
     {
-        if let Some(ref mut iter) = self.iter {
+        if let Ok(ref mut iter) = self.iter {
             acc = iter.try_rfold(acc, fold)?;
-            self.iter = None;
+            self.stop();
         }
         Try::from_ok(acc)
     }
 
     #[inline]
-    default fn rfold<Acc, Fold>(self, mut acc: Acc, fold: Fold) -> Acc
+    fn rfold<Acc, Fold>(self, mut acc: Acc, fold: Fold) -> Acc
     where
         Fold: FnMut(Acc, Self::Item) -> Acc,
     {
-        if let Some(iter) = self.iter {
+        if let Ok(iter) = self.iter {
             acc = iter.rfold(acc, fold);
         }
         acc
     }
 
     #[inline]
-    default fn rfind<P>(&mut self, predicate: P) -> Option<Self::Item>
+    fn rfind<P>(&mut self, predicate: P) -> Option<Self::Item>
     where
         P: FnMut(&Self::Item) -> bool,
     {
-        let found = self.iter.as_mut()?.rfind(predicate);
+        let found = self.iter.as_mut().ok()?.rfind(predicate);
         if found.is_none() {
-            self.iter = None;
+            self.stop();
         }
         found
     }
@@ -171,156 +211,13 @@ where
 impl<I> ExactSizeIterator for Fuse<I>
 where
     I: ExactSizeIterator,
-{
-    default fn len(&self) -> usize {
-        self.iter.as_ref().map_or(0, I::len)
-    }
-
-    default fn is_empty(&self) -> bool {
-        self.iter.as_ref().map_or(true, I::is_empty)
-    }
-}
-
-// NOTE: for `I: FusedIterator`, we assume that the iterator is always `Some`
-impl<I: FusedIterator> Fuse<I> {
-    #[inline(always)]
-    fn as_inner(&self) -> &I {
-        match self.iter {
-            Some(ref iter) => iter,
-            // SAFETY: the specialized iterator never sets `None`
-            None => unsafe { intrinsics::unreachable() },
-        }
-    }
-
-    #[inline(always)]
-    fn as_inner_mut(&mut self) -> &mut I {
-        match self.iter {
-            Some(ref mut iter) => iter,
-            // SAFETY: the specialized iterator never sets `None`
-            None => unsafe { intrinsics::unreachable() },
-        }
-    }
-
-    #[inline(always)]
-    fn into_inner(self) -> I {
-        match self.iter {
-            Some(iter) => iter,
-            // SAFETY: the specialized iterator never sets `None`
-            None => unsafe { intrinsics::unreachable() },
-        }
-    }
-}
-
-#[stable(feature = "fused", since = "1.26.0")]
-impl<I> Iterator for Fuse<I>
-where
-    I: FusedIterator,
-{
-    #[inline]
-    fn next(&mut self) -> Option<<I as Iterator>::Item> {
-        self.as_inner_mut().next()
-    }
-
-    #[inline]
-    fn nth(&mut self, n: usize) -> Option<I::Item> {
-        self.as_inner_mut().nth(n)
-    }
-
-    #[inline]
-    fn last(self) -> Option<I::Item> {
-        self.into_inner().last()
-    }
-
-    #[inline]
-    fn count(self) -> usize {
-        self.into_inner().count()
-    }
-
-    #[inline]
-    fn size_hint(&self) -> (usize, Option<usize>) {
-        self.as_inner().size_hint()
-    }
-
-    #[inline]
-    fn try_fold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R
-    where
-        Self: Sized,
-        Fold: FnMut(Acc, Self::Item) -> R,
-        R: Try<Ok = Acc>,
-    {
-        self.as_inner_mut().try_fold(init, fold)
-    }
-
-    #[inline]
-    fn fold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc
-    where
-        Fold: FnMut(Acc, Self::Item) -> Acc,
-    {
-        self.into_inner().fold(init, fold)
-    }
-
-    #[inline]
-    fn find<P>(&mut self, predicate: P) -> Option<Self::Item>
-    where
-        P: FnMut(&Self::Item) -> bool,
-    {
-        self.as_inner_mut().find(predicate)
-    }
-}
-
-#[stable(feature = "fused", since = "1.26.0")]
-impl<I> DoubleEndedIterator for Fuse<I>
-where
-    I: DoubleEndedIterator + FusedIterator,
-{
-    #[inline]
-    fn next_back(&mut self) -> Option<<I as Iterator>::Item> {
-        self.as_inner_mut().next_back()
-    }
-
-    #[inline]
-    fn nth_back(&mut self, n: usize) -> Option<<I as Iterator>::Item> {
-        self.as_inner_mut().nth_back(n)
-    }
-
-    #[inline]
-    fn try_rfold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R
-    where
-        Self: Sized,
-        Fold: FnMut(Acc, Self::Item) -> R,
-        R: Try<Ok = Acc>,
-    {
-        self.as_inner_mut().try_rfold(init, fold)
-    }
-
-    #[inline]
-    fn rfold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc
-    where
-        Fold: FnMut(Acc, Self::Item) -> Acc,
-    {
-        self.into_inner().rfold(init, fold)
-    }
-
-    #[inline]
-    fn rfind<P>(&mut self, predicate: P) -> Option<Self::Item>
-    where
-        P: FnMut(&Self::Item) -> bool,
-    {
-        self.as_inner_mut().rfind(predicate)
-    }
-}
-
-#[stable(feature = "rust1", since = "1.0.0")]
-impl<I> ExactSizeIterator for Fuse<I>
-where
-    I: ExactSizeIterator + FusedIterator,
 {
     fn len(&self) -> usize {
-        self.as_inner().len()
+        self.iter.as_ref().map_or(0, I::len)
     }
 
     fn is_empty(&self) -> bool {
-        self.as_inner().is_empty()
+        self.iter.as_ref().map_or(true, I::is_empty)
     }
 }
 
@@ -330,9 +227,9 @@ where
 {
     unsafe fn get_unchecked(&mut self, i: usize) -> I::Item {
         match self.iter {
-            Some(ref mut iter) => iter.get_unchecked(i),
+            Ok(ref mut iter) => iter.get_unchecked(i),
             // SAFETY: the caller asserts there is an item at `i`, so we're not exhausted.
-            None => intrinsics::unreachable(),
+            Err(_) => intrinsics::unreachable(),
         }
     }