Skip to content
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

Jump to bottom

Implement Chain with Option fuses #70896

Merged
merged 4 commits into from
Apr 9, 2020
Merged

Implement Chain with Option fuses #70896

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
859b8da
Implement Chain with Option fuses
cuviper Apr 7, 2020
8aac107
Reduce callsites in Chain::count()
cuviper Apr 7, 2020
2c4cffd
Reduce callsites in Chain::last()
cuviper Apr 7, 2020
ce8abc6
Avoid extra &mut in Chain::fold and try_fold
cuviper Apr 7, 2020
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
267 changes: 113 additions & 154 deletions src/libcore/iter/adapters/chain.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,7 @@
use crate::iter::{DoubleEndedIterator, FusedIterator, Iterator, TrustedLen};
use crate::ops::Try;
use crate::usize;

use super::super::{DoubleEndedIterator, FusedIterator, Iterator, TrustedLen};

/// An iterator that links two iterators together, in a chain.
///
/// This `struct` is created by the [`chain`] method on [`Iterator`]. See its
Expand All @@ -14,37 +13,34 @@ use super::super::{DoubleEndedIterator, FusedIterator, Iterator, TrustedLen};
#[must_use = "iterators are lazy and do nothing unless consumed"]
#[stable(feature = "rust1", since = "1.0.0")]
pub struct Chain<A, B> {
a: A,
b: B,
state: ChainState,
// These are "fused" with `Option` so we don't need separate state to track which part is
// already exhausted, and we may also get niche layout for `None`. We don't use the real `Fuse`
// adapter because its specialization for `FusedIterator` unconditionally descends into the
// iterator, and that could be expensive to keep revisiting stuff like nested chains. It also
// hurts compiler performance to add more iterator layers to `Chain`.
a: Option<A>,
b: Option<B>,
}
impl<A, B> Chain<A, B> {
pub(in super::super) fn new(a: A, b: B) -> Chain<A, B> {
Chain { a, b, state: ChainState::Both }
Chain { a: Some(a), b: Some(b) }
}
}

// The iterator protocol specifies that iteration ends with the return value
// `None` from `.next()` (or `.next_back()`) and it is unspecified what
// further calls return. The chain adaptor must account for this since it uses
// two subiterators.
//
// It uses three states:
//
// - Both: `a` and `b` are remaining
// - Front: `a` remaining
// - Back: `b` remaining
//
// The fourth state (neither iterator is remaining) only occurs after Chain has
// returned None once, so we don't need to store this state.
#[derive(Clone, Debug)]
enum ChainState {
// both front and back iterator are remaining
Both,
// only front is remaining
Front,
// only back is remaining
Back,
/// 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,
}
};
}

#[stable(feature = "rust1", since = "1.0.0")]
Expand All @@ -57,128 +53,101 @@ where

#[inline]
fn next(&mut self) -> Option<A::Item> {
match self.state {
ChainState::Both => match self.a.next() {
elt @ Some(..) => elt,
None => {
self.state = ChainState::Back;
self.b.next()
}
},
ChainState::Front => self.a.next(),
ChainState::Back => self.b.next(),
match fuse!(self.a.next()) {
None => fuse!(self.b.next()),
item => item,
}
}

#[inline]
#[rustc_inherit_overflow_checks]
fn count(self) -> usize {
match self.state {
ChainState::Both => self.a.count() + self.b.count(),
ChainState::Front => self.a.count(),
ChainState::Back => self.b.count(),
}
let a_count = match self.a {
Some(a) => a.count(),
None => 0,
};
let b_count = match self.b {
Some(b) => b.count(),
None => 0,
};
a_count + b_count
}

fn try_fold<Acc, F, R>(&mut self, init: Acc, mut f: F) -> R
fn try_fold<Acc, F, R>(&mut self, mut acc: Acc, mut f: F) -> R
where
Self: Sized,
F: FnMut(Acc, Self::Item) -> R,
R: Try<Ok = Acc>,
{
let mut accum = init;
match self.state {
ChainState::Both | ChainState::Front => {
accum = self.a.try_fold(accum, &mut f)?;
if let ChainState::Both = self.state {
self.state = ChainState::Back;
}
}
_ => {}
if let Some(ref mut a) = self.a {
acc = a.try_fold(acc, &mut f)?;
self.a = None;
}
if let ChainState::Back = self.state {
accum = self.b.try_fold(accum, &mut f)?;
if let Some(ref mut b) = self.b {
acc = b.try_fold(acc, f)?;
self.b = None;
}
Try::from_ok(accum)
Try::from_ok(acc)
}

fn fold<Acc, F>(self, init: Acc, mut f: F) -> Acc
fn fold<Acc, F>(self, mut acc: Acc, mut f: F) -> Acc
where
F: FnMut(Acc, Self::Item) -> Acc,
{
let mut accum = init;
match self.state {
ChainState::Both | ChainState::Front => {
accum = self.a.fold(accum, &mut f);
}
_ => {}
if let Some(a) = self.a {
acc = a.fold(acc, &mut f);
}
match self.state {
ChainState::Both | ChainState::Back => {
accum = self.b.fold(accum, &mut f);
}
_ => {}
if let Some(b) = self.b {
acc = b.fold(acc, f);
}
accum
acc
}

#[inline]
fn nth(&mut self, mut n: usize) -> Option<A::Item> {
match self.state {
ChainState::Both | ChainState::Front => {
for x in self.a.by_ref() {
if n == 0 {
return Some(x);
}
n -= 1;
}
if let ChainState::Both = self.state {
self.state = ChainState::Back;
if let Some(ref mut a) = self.a {
while let Some(x) = a.next() {
if n == 0 {
return Some(x);
}
n -= 1;
}
ChainState::Back => {}
self.a = None;
}
if let ChainState::Back = self.state { self.b.nth(n) } else { None }
fuse!(self.b.nth(n))
}

#[inline]
fn find<P>(&mut self, mut predicate: P) -> Option<Self::Item>
where
P: FnMut(&Self::Item) -> bool,
{
match self.state {
ChainState::Both => match self.a.find(&mut predicate) {
None => {
self.state = ChainState::Back;
self.b.find(predicate)
}
v => v,
},
ChainState::Front => self.a.find(predicate),
ChainState::Back => self.b.find(predicate),
match fuse!(self.a.find(&mut predicate)) {
None => fuse!(self.b.find(predicate)),
item => item,
}
}

#[inline]
fn last(self) -> Option<A::Item> {
match self.state {
ChainState::Both => {
// Must exhaust a before b.
let a_last = self.a.last();
let b_last = self.b.last();
b_last.or(a_last)
}
ChainState::Front => self.a.last(),
ChainState::Back => self.b.last(),
}
// 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,
};
b_last.or(a_last)
}

#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
match self.state {
ChainState::Both => {
let (a_lower, a_upper) = self.a.size_hint();
let (b_lower, b_upper) = self.b.size_hint();
match self {
Chain { a: Some(a), b: Some(b) } => {
let (a_lower, a_upper) = a.size_hint();
let (b_lower, b_upper) = b.size_hint();

let lower = a_lower.saturating_add(b_lower);

Expand All @@ -189,8 +158,9 @@ where

(lower, upper)
}
ChainState::Front => self.a.size_hint(),
ChainState::Back => self.b.size_hint(),
Chain { a: Some(a), b: None } => a.size_hint(),
Chain { a: None, b: Some(b) } => b.size_hint(),
Chain { a: None, b: None } => (0, Some(0)),
}
}
}
Expand All @@ -203,82 +173,71 @@ where
{
#[inline]
fn next_back(&mut self) -> Option<A::Item> {
match self.state {
ChainState::Both => match self.b.next_back() {
elt @ Some(..) => elt,
None => {
self.state = ChainState::Front;
self.a.next_back()
}
},
ChainState::Front => self.a.next_back(),
ChainState::Back => self.b.next_back(),
match fuse!(self.b.next_back()) {
None => fuse!(self.a.next_back()),
item => item,
}
}

#[inline]
fn nth_back(&mut self, mut n: usize) -> Option<A::Item> {
match self.state {
ChainState::Both | ChainState::Back => {
for x in self.b.by_ref().rev() {
if n == 0 {
return Some(x);
}
n -= 1;
}
if let ChainState::Both = self.state {
self.state = ChainState::Front;
if let Some(ref mut b) = self.b {
while let Some(x) = b.next_back() {
if n == 0 {
return Some(x);
}
n -= 1;
}
ChainState::Front => {}
self.b = None;
}
if let ChainState::Front = self.state { self.a.nth_back(n) } else { None }
fuse!(self.a.nth_back(n))
}

fn try_rfold<Acc, F, R>(&mut self, init: Acc, mut f: F) -> R
#[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 => fuse!(self.a.rfind(predicate)),
item => item,
}
}

fn try_rfold<Acc, F, R>(&mut self, mut acc: Acc, mut f: F) -> R
where
Self: Sized,
F: FnMut(Acc, Self::Item) -> R,
R: Try<Ok = Acc>,
{
let mut accum = init;
match self.state {
ChainState::Both | ChainState::Back => {
accum = self.b.try_rfold(accum, &mut f)?;
if let ChainState::Both = self.state {
self.state = ChainState::Front;
}
}
_ => {}
if let Some(ref mut b) = self.b {
acc = b.try_rfold(acc, &mut f)?;
self.b = None;
}
if let ChainState::Front = self.state {
accum = self.a.try_rfold(accum, &mut f)?;
if let Some(ref mut a) = self.a {
acc = a.try_rfold(acc, f)?;
self.a = None;
}
Try::from_ok(accum)
Try::from_ok(acc)
}

fn rfold<Acc, F>(self, init: Acc, mut f: F) -> Acc
fn rfold<Acc, F>(self, mut acc: Acc, mut f: F) -> Acc
where
F: FnMut(Acc, Self::Item) -> Acc,
{
let mut accum = init;
match self.state {
ChainState::Both | ChainState::Back => {
accum = self.b.rfold(accum, &mut f);
}
_ => {}
if let Some(b) = self.b {
acc = b.rfold(acc, &mut f);
}
match self.state {
ChainState::Both | ChainState::Front => {
accum = self.a.rfold(accum, &mut f);
}
_ => {}
if let Some(a) = self.a {
acc = a.rfold(acc, f);
}
accum
acc
}
}

// Note: *both* must be fused to handle double-ended iterators.
// Now that we "fuse" both sides, we *could* implement this unconditionally,
// but we should be cautious about committing to that in the public API.
#[stable(feature = "fused", since = "1.26.0")]
impl<A, B> FusedIterator for Chain<A, B>
where
Expand Down