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.

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

refactor: removing custom nth Zip fn #79173

Closed
wants to merge 5 commits into from
Closed
Changes from all commits
Commits
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
247 changes: 106 additions & 141 deletions library/core/src/iter/adapters/zip.rs
Original file line number Diff line number Diff line change
Expand Up @@ -18,18 +18,19 @@ pub struct Zip<A, B> {
len: usize,
a_len: usize,
}
impl<A: Iterator, B: Iterator> Zip<A, B> {

impl<A, B> Zip<A, B>
where
A: Iterator,
B: Iterator,
{
pub(in crate::iter) fn new(a: A, b: B) -> Zip<A, B> {
ZipImpl::new(a, b)
ZipNew::new(a, b)
}
fn super_nth(&mut self, mut n: usize) -> Option<(A::Item, B::Item)> {
while let Some(x) = Iterator::next(self) {
if n == 0 {
return Some(x);
}
n -= 1;
}
None

fn super_nth(&mut self, n: usize) -> Option<(A::Item, B::Item)> {
self.advance_by(n).ok()?;
self.next()
}
}

Expand Down Expand Up @@ -61,81 +62,20 @@ where
A: IntoIterator,
B: IntoIterator,
{
ZipImpl::new(a.into_iter(), b.into_iter())
ZipNew::new(a.into_iter(), b.into_iter())
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<A, B> Iterator for Zip<A, B>
where
A: Iterator,
B: Iterator,
{
type Item = (A::Item, B::Item);

#[inline]
fn next(&mut self) -> Option<Self::Item> {
ZipImpl::next(self)
}

#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
ZipImpl::size_hint(self)
}

#[inline]
fn nth(&mut self, n: usize) -> Option<Self::Item> {
ZipImpl::nth(self, n)
}

#[inline]
unsafe fn __iterator_get_unchecked(&mut self, idx: usize) -> Self::Item
where
Self: TrustedRandomAccess,
{
// SAFETY: `ZipImpl::__iterator_get_unchecked` has same safety
// requirements as `Iterator::__iterator_get_unchecked`.
unsafe { ZipImpl::get_unchecked(self, idx) }
}
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<A, B> DoubleEndedIterator for Zip<A, B>
where
A: DoubleEndedIterator + ExactSizeIterator,
B: DoubleEndedIterator + ExactSizeIterator,
{
#[inline]
fn next_back(&mut self) -> Option<(A::Item, B::Item)> {
ZipImpl::next_back(self)
}
}

// Zip specialization trait
#[doc(hidden)]
trait ZipImpl<A, B> {
type Item;
trait ZipNew<A, B> {
fn new(a: A, b: B) -> Self;
fn next(&mut self) -> Option<Self::Item>;
fn size_hint(&self) -> (usize, Option<usize>);
fn nth(&mut self, n: usize) -> Option<Self::Item>;
fn next_back(&mut self) -> Option<Self::Item>
where
A: DoubleEndedIterator + ExactSizeIterator,
B: DoubleEndedIterator + ExactSizeIterator;
// This has the same safety requirements as `Iterator::__iterator_get_unchecked`
unsafe fn get_unchecked(&mut self, idx: usize) -> <Self as Iterator>::Item
where
Self: Iterator + TrustedRandomAccess;
}

// General Zip impl
#[doc(hidden)]
impl<A, B> ZipImpl<A, B> for Zip<A, B>
impl<A, B> ZipNew<A, B> for Zip<A, B>
where
A: Iterator,
B: Iterator,
{
type Item = (A::Item, B::Item);
default fn new(a: A, b: B) -> Self {
Zip {
a,
Expand All @@ -145,25 +85,29 @@ where
a_len: 0, // unused
}
}
}

#[inline]
default fn next(&mut self) -> Option<(A::Item, B::Item)> {
let x = self.a.next()?;
let y = self.b.next()?;
Some((x, y))
}

#[inline]
default fn nth(&mut self, n: usize) -> Option<Self::Item> {
self.super_nth(n)
#[doc(hidden)]
impl<A, B> ZipNew<A, B> for Zip<A, B>
where
A: TrustedRandomAccess + Iterator,
B: TrustedRandomAccess + Iterator,
{
fn new(a: A, b: B) -> Self {
let a_len = a.size();
let len = cmp::min(a_len, b.size());
Zip { a, b, index: 0, len, a_len }
}
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<A, B> DoubleEndedIterator for Zip<A, B>
where
A: DoubleEndedIterator + ExactSizeIterator,
B: DoubleEndedIterator + ExactSizeIterator,
{
#[inline]
default fn next_back(&mut self) -> Option<(A::Item, B::Item)>
where
A: DoubleEndedIterator + ExactSizeIterator,
B: DoubleEndedIterator + ExactSizeIterator,
{
default fn next_back(&mut self) -> Option<(A::Item, B::Item)> {
let a_sz = self.a.len();
let b_sz = self.b.len();
if a_sz != b_sz {
Expand All @@ -184,6 +128,75 @@ where
_ => unreachable!(),
}
}
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<A, B> DoubleEndedIterator for Zip<A, B>
where
A: TrustedRandomAccess + DoubleEndedIterator + ExactSizeIterator,
B: TrustedRandomAccess + DoubleEndedIterator + ExactSizeIterator,
{
#[inline]
fn next_back(&mut self) -> Option<(A::Item, B::Item)>
where
A: DoubleEndedIterator + ExactSizeIterator,
B: DoubleEndedIterator + ExactSizeIterator,
{
if A::MAY_HAVE_SIDE_EFFECT || B::MAY_HAVE_SIDE_EFFECT {
let sz_a = self.a.size();
let sz_b = self.b.size();
// Adjust a, b to equal length, make sure that only the first call
// of `next_back` does this, otherwise we will break the restriction
// on calls to `self.next_back()` after calling `get_unchecked()`.
if sz_a != sz_b {
if A::MAY_HAVE_SIDE_EFFECT && sz_a > self.len {
for _ in 0..sz_a - self.len {
self.a.next_back();
}
self.a_len = self.len;
}

if B::MAY_HAVE_SIDE_EFFECT && sz_b > self.len {
for _ in 0..sz_b - self.len {
self.b.next_back();
}
}
}
}
if self.index < self.len {
self.len -= 1;
self.a_len -= 1;
let i = self.len;
// SAFETY: `i` is smaller than the previous value of `self.len`,
// which is also smaller than or equal to `self.a.len()` and `self.b.len()`
unsafe {
Some((self.a.__iterator_get_unchecked(i), self.b.__iterator_get_unchecked(i)))
}
} else {
None
}
}
}

#[stable(feature = "rust1", since = "1.0.0")]
impl<A, B> Iterator for Zip<A, B>
where
A: Iterator,
B: Iterator,
{
type Item = (A::Item, B::Item);

#[inline]
default fn next(&mut self) -> Option<(A::Item, B::Item)> {
let x = self.a.next()?;
let y = self.b.next()?;
Some((x, y))
}

#[inline]
default fn nth(&mut self, n: usize) -> Option<Self::Item> {
self.super_nth(n)
}

#[inline]
default fn size_hint(&self) -> (usize, Option<usize>) {
Expand All @@ -202,26 +215,20 @@ where
(lower, upper)
}

default unsafe fn get_unchecked(&mut self, _idx: usize) -> <Self as Iterator>::Item
default unsafe fn __iterator_get_unchecked(&mut self, _idx: usize) -> <Self as Iterator>::Item
where
Self: TrustedRandomAccess,
{
unreachable!("Always specialized");
}
}

#[doc(hidden)]
impl<A, B> ZipImpl<A, B> for Zip<A, B>
#[stable(feature = "rust1", since = "1.0.0")]
impl<A, B> Iterator for Zip<A, B>
where
A: TrustedRandomAccess + Iterator,
B: TrustedRandomAccess + Iterator,
{
fn new(a: A, b: B) -> Self {
let a_len = a.size();
let len = cmp::min(a_len, b.size());
Zip { a, b, index: 0, len, a_len }
}

#[inline]
fn next(&mut self) -> Option<(A::Item, B::Item)> {
if self.index < self.len {
Expand Down Expand Up @@ -279,49 +286,7 @@ where
}

#[inline]
fn next_back(&mut self) -> Option<(A::Item, B::Item)>
where
A: DoubleEndedIterator + ExactSizeIterator,
B: DoubleEndedIterator + ExactSizeIterator,
{
if A::MAY_HAVE_SIDE_EFFECT || B::MAY_HAVE_SIDE_EFFECT {
let sz_a = self.a.size();
let sz_b = self.b.size();
// Adjust a, b to equal length, make sure that only the first call
// of `next_back` does this, otherwise we will break the restriction
// on calls to `self.next_back()` after calling `get_unchecked()`.
if sz_a != sz_b {
let sz_a = self.a.size();
if A::MAY_HAVE_SIDE_EFFECT && sz_a > self.len {
for _ in 0..sz_a - self.len {
self.a.next_back();
}
self.a_len = self.len;
}
let sz_b = self.b.size();
if B::MAY_HAVE_SIDE_EFFECT && sz_b > self.len {
for _ in 0..sz_b - self.len {
self.b.next_back();
}
}
}
}
if self.index < self.len {
self.len -= 1;
self.a_len -= 1;
let i = self.len;
// SAFETY: `i` is smaller than the previous value of `self.len`,
// which is also smaller than or equal to `self.a.len()` and `self.b.len()`
unsafe {
Some((self.a.__iterator_get_unchecked(i), self.b.__iterator_get_unchecked(i)))
}
} else {
None
}
}

#[inline]
unsafe fn get_unchecked(&mut self, idx: usize) -> <Self as Iterator>::Item {
unsafe fn __iterator_get_unchecked(&mut self, idx: usize) -> <Self as Iterator>::Item {
let idx = self.index + idx;
// SAFETY: the caller must uphold the contract for
// `Iterator::__iterator_get_unchecked`.
Expand Down