Skip to content

Commit

Permalink
Specialize internal state of Fuse depending on underlying iterator
Browse files Browse the repository at this point in the history
  • Loading branch information
AnthonyMikh committed Mar 28, 2020
1 parent 7520894 commit 4ea75af
Showing 1 changed file with 79 additions and 182 deletions.
261 changes: 79 additions & 182 deletions src/libcore/iter/adapters/fuse.rs
Original file line number Diff line number Diff line change
@@ -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.
///
Expand All @@ -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);
}
}

Expand All @@ -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
}
Expand All @@ -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
}
Expand All @@ -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)
}
}

Expand All @@ -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(),
}
}

Expand Down

0 comments on commit 4ea75af

Please sign in to comment.