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

   1 use crate::cmp;
   2 use crate::fmt::{self, Debug};
   3 use crate::iter::{DoubleEndedIterator, ExactSizeIterator, FusedIterator, Iterator};
   4 use crate::iter::{InPlaceIterable, SourceIter, TrustedLen};
   5
   6 /// An iterator that iterates two other iterators simultaneously.
   7 ///
   8 /// This `struct` is created by [`zip`] or [`Iterator::zip`].
   9 /// See their documentation for more.
  10 #[derive(Clone)]
  11 #[must_use = "iterators are lazy and do nothing unless consumed"]
  12 #[stable(feature = "rust1", since = "1.0.0")]
  13 pub struct Zip<A, B> {
  14     a: A,
  15     b: B,
  16     // index, len and a_len are only used by the specialized version of zip
  17     index: usize,
  18     len: usize,
  19     a_len: usize,
  20 }
  21 impl<A: Iterator, B: Iterator> Zip<A, B> {
  22     pub(in crate::iter) fn new(a: A, b: B) -> Zip<A, B> {
  23         ZipImpl::new(a, b)
  24     }
  25     fn super_nth(&mut self, mut n: usize) -> Option<(A::Item, B::Item)> {
  26         while let Some(x) = Iterator::next(self) {
  27             if n == 0 {
  28                 return Some(x);
  29             }
  30             n -= 1;
  31         }
  32         None
  33     }
  34 }
  35
  36 /// Converts the arguments to iterators and zips them.
  37 ///
  38 /// See the documentation of [`Iterator::zip`] for more.
  39 ///
  40 /// # Examples
  41 ///
  42 /// ```
  43 /// #![feature(iter_zip)]
  44 /// use std::iter::zip;
  45 ///
  46 /// let xs = [1, 2, 3];
  47 /// let ys = [4, 5, 6];
  48 /// for (x, y) in zip(&xs, &ys) {
  49 ///     println!("x:{}, y:{}", x, y);
  50 /// }
  51 ///
  52 /// // Nested zips are also possible:
  53 /// let zs = [7, 8, 9];
  54 /// for ((x, y), z) in zip(zip(&xs, &ys), &zs) {
  55 ///     println!("x:{}, y:{}, z:{}", x, y, z);
  56 /// }
  57 /// ```
  58 #[unstable(feature = "iter_zip", issue = "83574")]
  59 pub fn zip<A, B>(a: A, b: B) -> Zip<A::IntoIter, B::IntoIter>
  60 where
  61     A: IntoIterator,
  62     B: IntoIterator,
  63 {
  64     ZipImpl::new(a.into_iter(), b.into_iter())
  65 }
  66
  67 #[stable(feature = "rust1", since = "1.0.0")]
  68 impl<A, B> Iterator for Zip<A, B>
  69 where
  70     A: Iterator,
  71     B: Iterator,
  72 {
  73     type Item = (A::Item, B::Item);
  74
  75     #[inline]
  76     fn next(&mut self) -> Option<Self::Item> {
  77         ZipImpl::next(self)
  78     }
  79
  80     #[inline]
  81     fn size_hint(&self) -> (usize, Option<usize>) {
  82         ZipImpl::size_hint(self)
  83     }
  84
  85     #[inline]
  86     fn nth(&mut self, n: usize) -> Option<Self::Item> {
  87         ZipImpl::nth(self, n)
  88     }
  89
  90     #[inline]
  91     #[doc(hidden)]
  92     unsafe fn __iterator_get_unchecked(&mut self, idx: usize) -> Self::Item
  93     where
  94         Self: TrustedRandomAccessNoCoerce,
  95     {
  96         // SAFETY: `ZipImpl::__iterator_get_unchecked` has same safety
  97         // requirements as `Iterator::__iterator_get_unchecked`.
  98         unsafe { ZipImpl::get_unchecked(self, idx) }
  99     }
 100 }
 101
 102 #[stable(feature = "rust1", since = "1.0.0")]
 103 impl<A, B> DoubleEndedIterator for Zip<A, B>
 104 where
 105     A: DoubleEndedIterator + ExactSizeIterator,
 106     B: DoubleEndedIterator + ExactSizeIterator,
 107 {
 108     #[inline]
 109     fn next_back(&mut self) -> Option<(A::Item, B::Item)> {
 110         ZipImpl::next_back(self)
 111     }
 112 }
 113
 114 // Zip specialization trait
 115 #[doc(hidden)]
 116 trait ZipImpl<A, B> {
 117     type Item;
 118     fn new(a: A, b: B) -> Self;
 119     fn next(&mut self) -> Option<Self::Item>;
 120     fn size_hint(&self) -> (usize, Option<usize>);
 121     fn nth(&mut self, n: usize) -> Option<Self::Item>;
 122     fn next_back(&mut self) -> Option<Self::Item>
 123     where
 124         A: DoubleEndedIterator + ExactSizeIterator,
 125         B: DoubleEndedIterator + ExactSizeIterator;
 126     // This has the same safety requirements as `Iterator::__iterator_get_unchecked`
 127     unsafe fn get_unchecked(&mut self, idx: usize) -> <Self as Iterator>::Item
 128     where
 129         Self: Iterator + TrustedRandomAccessNoCoerce;
 130 }
 131
 132 // Work around limitations of specialization, requiring `default` impls to be repeated
 133 // in intermediary impls.
 134 macro_rules! zip_impl_general_defaults {
 135     () => {
 136         default fn new(a: A, b: B) -> Self {
 137             Zip {
 138                 a,
 139                 b,
 140                 index: 0, // unused
 141                 len: 0,   // unused
 142                 a_len: 0, // unused
 143             }
 144         }
 145
 146         #[inline]
 147         default fn next(&mut self) -> Option<(A::Item, B::Item)> {
 148             let x = self.a.next()?;
 149             let y = self.b.next()?;
 150             Some((x, y))
 151         }
 152
 153         #[inline]
 154         default fn nth(&mut self, n: usize) -> Option<Self::Item> {
 155             self.super_nth(n)
 156         }
 157
 158         #[inline]
 159         default fn next_back(&mut self) -> Option<(A::Item, B::Item)>
 160         where
 161             A: DoubleEndedIterator + ExactSizeIterator,
 162             B: DoubleEndedIterator + ExactSizeIterator,
 163         {
 164             // The function body below only uses `self.a/b.len()` and `self.a/b.next_back()`
 165             // and doesn’t call `next_back` too often, so this implementation is safe in
 166             // the `TrustedRandomAccessNoCoerce` specialization
 167
 168             let a_sz = self.a.len();
 169             let b_sz = self.b.len();
 170             if a_sz != b_sz {
 171                 // Adjust a, b to equal length
 172                 if a_sz > b_sz {
 173                     for _ in 0..a_sz - b_sz {
 174                         self.a.next_back();
 175                     }
 176                 } else {
 177                     for _ in 0..b_sz - a_sz {
 178                         self.b.next_back();
 179                     }
 180                 }
 181             }
 182             match (self.a.next_back(), self.b.next_back()) {
 183                 (Some(x), Some(y)) => Some((x, y)),
 184                 (None, None) => None,
 185                 _ => unreachable!(),
 186             }
 187         }
 188     };
 189 }
 190
 191 // General Zip impl
 192 #[doc(hidden)]
 193 impl<A, B> ZipImpl<A, B> for Zip<A, B>
 194 where
 195     A: Iterator,
 196     B: Iterator,
 197 {
 198     type Item = (A::Item, B::Item);
 199
 200     zip_impl_general_defaults! {}
 201
 202     #[inline]
 203     default fn size_hint(&self) -> (usize, Option<usize>) {
 204         let (a_lower, a_upper) = self.a.size_hint();
 205         let (b_lower, b_upper) = self.b.size_hint();
 206
 207         let lower = cmp::min(a_lower, b_lower);
 208
 209         let upper = match (a_upper, b_upper) {
 210             (Some(x), Some(y)) => Some(cmp::min(x, y)),
 211             (Some(x), None) => Some(x),
 212             (None, Some(y)) => Some(y),
 213             (None, None) => None,
 214         };
 215
 216         (lower, upper)
 217     }
 218
 219     default unsafe fn get_unchecked(&mut self, _idx: usize) -> <Self as Iterator>::Item
 220     where
 221         Self: TrustedRandomAccessNoCoerce,
 222     {
 223         unreachable!("Always specialized");
 224     }
 225 }
 226
 227 #[doc(hidden)]
 228 impl<A, B> ZipImpl<A, B> for Zip<A, B>
 229 where
 230     A: TrustedRandomAccessNoCoerce + Iterator,
 231     B: TrustedRandomAccessNoCoerce + Iterator,
 232 {
 233     zip_impl_general_defaults! {}
 234
 235     #[inline]
 236     default fn size_hint(&self) -> (usize, Option<usize>) {
 237         let size = cmp::min(self.a.size(), self.b.size());
 238         (size, Some(size))
 239     }
 240
 241     #[inline]
 242     unsafe fn get_unchecked(&mut self, idx: usize) -> <Self as Iterator>::Item {
 243         let idx = self.index + idx;
 244         // SAFETY: the caller must uphold the contract for
 245         // `Iterator::__iterator_get_unchecked`.
 246         unsafe { (self.a.__iterator_get_unchecked(idx), self.b.__iterator_get_unchecked(idx)) }
 247     }
 248 }
 249
 250 #[doc(hidden)]
 251 impl<A, B> ZipImpl<A, B> for Zip<A, B>
 252 where
 253     A: TrustedRandomAccess + Iterator,
 254     B: TrustedRandomAccess + Iterator,
 255 {
 256     fn new(a: A, b: B) -> Self {
 257         let a_len = a.size();
 258         let len = cmp::min(a_len, b.size());
 259         Zip { a, b, index: 0, len, a_len }
 260     }
 261
 262     #[inline]
 263     fn next(&mut self) -> Option<(A::Item, B::Item)> {
 264         if self.index < self.len {
 265             let i = self.index;
 266             // since get_unchecked executes code which can panic we increment the counters beforehand
 267             // so that the same index won't be accessed twice, as required by TrustedRandomAccess
 268             self.index += 1;
 269             // SAFETY: `i` is smaller than `self.len`, thus smaller than `self.a.len()` and `self.b.len()`
 270             unsafe {
 271                 Some((self.a.__iterator_get_unchecked(i), self.b.__iterator_get_unchecked(i)))
 272             }
 273         } else if A::MAY_HAVE_SIDE_EFFECT && self.index < self.a_len {
 274             let i = self.index;
 275             // as above, increment before executing code that may panic
 276             self.index += 1;
 277             self.len += 1;
 278             // match the base implementation's potential side effects
 279             // SAFETY: we just checked that `i` < `self.a.len()`
 280             unsafe {
 281                 self.a.__iterator_get_unchecked(i);
 282             }
 283             None
 284         } else {
 285             None
 286         }
 287     }
 288
 289     #[inline]
 290     fn size_hint(&self) -> (usize, Option<usize>) {
 291         let len = self.len - self.index;
 292         (len, Some(len))
 293     }
 294
 295     #[inline]
 296     fn nth(&mut self, n: usize) -> Option<Self::Item> {
 297         let delta = cmp::min(n, self.len - self.index);
 298         let end = self.index + delta;
 299         while self.index < end {
 300             let i = self.index;
 301             // since get_unchecked executes code which can panic we increment the counters beforehand
 302             // so that the same index won't be accessed twice, as required by TrustedRandomAccess
 303             self.index += 1;
 304             if A::MAY_HAVE_SIDE_EFFECT {
 305                 // SAFETY: the usage of `cmp::min` to calculate `delta`
 306                 // ensures that `end` is smaller than or equal to `self.len`,
 307                 // so `i` is also smaller than `self.len`.
 308                 unsafe {
 309                     self.a.__iterator_get_unchecked(i);
 310                 }
 311             }
 312             if B::MAY_HAVE_SIDE_EFFECT {
 313                 // SAFETY: same as above.
 314                 unsafe {
 315                     self.b.__iterator_get_unchecked(i);
 316                 }
 317             }
 318         }
 319
 320         self.super_nth(n - delta)
 321     }
 322
 323     #[inline]
 324     fn next_back(&mut self) -> Option<(A::Item, B::Item)>
 325     where
 326         A: DoubleEndedIterator + ExactSizeIterator,
 327         B: DoubleEndedIterator + ExactSizeIterator,
 328     {
 329         if A::MAY_HAVE_SIDE_EFFECT || B::MAY_HAVE_SIDE_EFFECT {
 330             let sz_a = self.a.size();
 331             let sz_b = self.b.size();
 332             // Adjust a, b to equal length, make sure that only the first call
 333             // of `next_back` does this, otherwise we will break the restriction
 334             // on calls to `self.next_back()` after calling `get_unchecked()`.
 335             if sz_a != sz_b {
 336                 let sz_a = self.a.size();
 337                 if A::MAY_HAVE_SIDE_EFFECT && sz_a > self.len {
 338                     for _ in 0..sz_a - self.len {
 339                         // since next_back() may panic we increment the counters beforehand
 340                         // to keep Zip's state in sync with the underlying iterator source
 341                         self.a_len -= 1;
 342                         self.a.next_back();
 343                     }
 344                     debug_assert_eq!(self.a_len, self.len);
 345                 }
 346                 let sz_b = self.b.size();
 347                 if B::MAY_HAVE_SIDE_EFFECT && sz_b > self.len {
 348                     for _ in 0..sz_b - self.len {
 349                         self.b.next_back();
 350                     }
 351                 }
 352             }
 353         }
 354         if self.index < self.len {
 355             // since get_unchecked executes code which can panic we increment the counters beforehand
 356             // so that the same index won't be accessed twice, as required by TrustedRandomAccess
 357             self.len -= 1;
 358             self.a_len -= 1;
 359             let i = self.len;
 360             // SAFETY: `i` is smaller than the previous value of `self.len`,
 361             // which is also smaller than or equal to `self.a.len()` and `self.b.len()`
 362             unsafe {
 363                 Some((self.a.__iterator_get_unchecked(i), self.b.__iterator_get_unchecked(i)))
 364             }
 365         } else {
 366             None
 367         }
 368     }
 369 }
 370
 371 #[stable(feature = "rust1", since = "1.0.0")]
 372 impl<A, B> ExactSizeIterator for Zip<A, B>
 373 where
 374     A: ExactSizeIterator,
 375     B: ExactSizeIterator,
 376 {
 377 }
 378
 379 #[doc(hidden)]
 380 #[unstable(feature = "trusted_random_access", issue = "none")]
 381 unsafe impl<A, B> TrustedRandomAccess for Zip<A, B>
 382 where
 383     A: TrustedRandomAccess,
 384     B: TrustedRandomAccess,
 385 {
 386 }
 387
 388 #[doc(hidden)]
 389 #[unstable(feature = "trusted_random_access", issue = "none")]
 390 unsafe impl<A, B> TrustedRandomAccessNoCoerce for Zip<A, B>
 391 where
 392     A: TrustedRandomAccessNoCoerce,
 393     B: TrustedRandomAccessNoCoerce,
 394 {
 395     const MAY_HAVE_SIDE_EFFECT: bool = A::MAY_HAVE_SIDE_EFFECT || B::MAY_HAVE_SIDE_EFFECT;
 396 }
 397
 398 #[stable(feature = "fused", since = "1.26.0")]
 399 impl<A, B> FusedIterator for Zip<A, B>
 400 where
 401     A: FusedIterator,
 402     B: FusedIterator,
 403 {
 404 }
 405
 406 #[unstable(feature = "trusted_len", issue = "37572")]
 407 unsafe impl<A, B> TrustedLen for Zip<A, B>
 408 where
 409     A: TrustedLen,
 410     B: TrustedLen,
 411 {
 412 }
 413
 414 // Arbitrarily selects the left side of the zip iteration as extractable "source"
 415 // it would require negative trait bounds to be able to try both
 416 #[unstable(issue = "none", feature = "inplace_iteration")]
 417 unsafe impl<S, A, B> SourceIter for Zip<A, B>
 418 where
 419     A: SourceIter<Source = S>,
 420     B: Iterator,
 421     S: Iterator,
 422 {
 423     type Source = S;
 424
 425     #[inline]
 426     unsafe fn as_inner(&mut self) -> &mut S {
 427         // SAFETY: unsafe function forwarding to unsafe function with the same requirements
 428         unsafe { SourceIter::as_inner(&mut self.a) }
 429     }
 430 }
 431
 432 #[unstable(issue = "none", feature = "inplace_iteration")]
 433 // Limited to Item: Copy since interaction between Zip's use of TrustedRandomAccess
 434 // and Drop implementation of the source is unclear.
 435 //
 436 // An additional method returning the number of times the source has been logically advanced
 437 // (without calling next()) would be needed to properly drop the remainder of the source.
 438 unsafe impl<A: InPlaceIterable, B: Iterator> InPlaceIterable for Zip<A, B> where A::Item: Copy {}
 439
 440 #[stable(feature = "rust1", since = "1.0.0")]
 441 impl<A: Debug, B: Debug> Debug for Zip<A, B> {
 442     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 443         ZipFmt::fmt(self, f)
 444     }
 445 }
 446
 447 trait ZipFmt<A, B> {
 448     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result;
 449 }
 450
 451 impl<A: Debug, B: Debug> ZipFmt<A, B> for Zip<A, B> {
 452     default fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 453         f.debug_struct("Zip").field("a", &self.a).field("b", &self.b).finish()
 454     }
 455 }
 456
 457 impl<A: Debug + TrustedRandomAccessNoCoerce, B: Debug + TrustedRandomAccessNoCoerce> ZipFmt<A, B>
 458     for Zip<A, B>
 459 {
 460     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 461         // It's *not safe* to call fmt on the contained iterators, since once
 462         // we start iterating they're in strange, potentially unsafe, states.
 463         f.debug_struct("Zip").finish()
 464     }
 465 }
 466
 467 /// An iterator whose items are random-accessible efficiently
 468 ///
 469 /// # Safety
 470 ///
 471 /// The iterator's `size_hint` must be exact and cheap to call.
 472 ///
 473 /// `TrustedRandomAccessNoCoerce::size` may not be overridden.
 474 ///
 475 /// All subtypes and all supertypes of `Self` must also implement `TrustedRandomAccess`.
 476 /// In particular, this means that types with non-invariant parameters usually can not have
 477 /// an impl for `TrustedRandomAccess` that depends on any trait bounds on such parameters, except
 478 /// for bounds that come from the respective struct/enum definition itself, or bounds involving
 479 /// traits that themselves come with a guarantee similar to this one.
 480 ///
 481 /// If `Self: ExactSizeIterator` then `self.len()` must always produce results consistent
 482 /// with `self.size()`.
 483 ///
 484 /// If `Self: Iterator`, then `<Self as Iterator>::__iterator_get_unchecked(&mut self, idx)`
 485 /// must be safe to call provided the following conditions are met.
 486 ///
 487 /// 1. `0 <= idx` and `idx < self.size()`.
 488 /// 2. If `Self: !Clone`, then `self.__iterator_get_unchecked(idx)` is never called with the same
 489 ///    index on `self` more than once.
 490 /// 3. After `self.__iterator_get_unchecked(idx)` has been called, then `self.next_back()` will
 491 ///    only be called at most `self.size() - idx - 1` times. If `Self: Clone` and `self` is cloned,
 492 ///    then this number is calculated for `self` and its clone individually,
 493 ///    but `self.next_back()` calls that happened before the cloning count for both `self` and the clone.
 494 /// 4. After `self.__iterator_get_unchecked(idx)` has been called, then only the following methods
 495 ///    will be called on `self` or on any new clones of `self`:
 496 ///     * `std::clone::Clone::clone`
 497 ///     * `std::iter::Iterator::size_hint`
 498 ///     * `std::iter::DoubleEndedIterator::next_back`
 499 ///     * `std::iter::ExactSizeIterator::len`
 500 ///     * `std::iter::Iterator::__iterator_get_unchecked`
 501 ///     * `std::iter::TrustedRandomAccessNoCoerce::size`
 502 /// 5. If `T` is a subtype of `Self`, then `self` is allowed to be coerced
 503 ///    to `T`. If `self` is coerced to `T` after `self.__iterator_get_unchecked(idx)` has already
 504 ///    been called, then no methods except for the ones listed under 4. are allowed to be called
 505 ///    on the resulting value of type `T`, either. Multiple such coercion steps are allowed.
 506 ///    Regarding 2. and 3., the number of times `__iterator_get_unchecked(idx)` or `next_back()` is
 507 ///    called on `self` and the resulting value of type `T` (and on further coercion results with
 508 ///    sub-subtypes) are added together and their sums must not exceed the specified bounds.
 509 ///
 510 /// Further, given that these conditions are met, it must guarantee that:
 511 ///
 512 /// * It does not change the value returned from `size_hint`
 513 /// * It must be safe to call the methods listed above on `self` after calling
 514 ///   `self.__iterator_get_unchecked(idx)`, assuming that the required traits are implemented.
 515 /// * It must also be safe to drop `self` after calling `self.__iterator_get_unchecked(idx)`.
 516 /// * If `T` is a subtype of `Self`, then it must be safe to coerce `self` to `T`.
 517 //
 518 // FIXME: Clarify interaction with SourceIter/InPlaceIterable. Calling `SouceIter::as_inner`
 519 // after `__iterator_get_unchecked` is supposed to be allowed.
 520 #[doc(hidden)]
 521 #[unstable(feature = "trusted_random_access", issue = "none")]
 522 #[rustc_specialization_trait]
 523 pub unsafe trait TrustedRandomAccess: TrustedRandomAccessNoCoerce {}
 524
 525 /// Like [`TrustedRandomAccess`] but without any of the requirements / guarantees around
 526 /// coercions to subtypes after `__iterator_get_unchecked` (they aren’t allowed here!), and
 527 /// without the requirement that subtypes / supertypes implement `TrustedRandomAccessNoCoerce`.
 528 ///
 529 /// This trait was created in PR #85874 to fix soundness issue #85873 without performance regressions.
 530 /// It is subject to change as we might want to build a more generally useful (for performance
 531 /// optimizations) and more sophisticated trait or trait hierarchy that replaces or extends
 532 /// [`TrustedRandomAccess`] and `TrustedRandomAccessNoCoerce`.
 533 #[doc(hidden)]
 534 #[unstable(feature = "trusted_random_access", issue = "none")]
 535 #[rustc_specialization_trait]
 536 pub unsafe trait TrustedRandomAccessNoCoerce: Sized {
 537     // Convenience method.
 538     fn size(&self) -> usize
 539     where
 540         Self: Iterator,
 541     {
 542         self.size_hint().0
 543     }
 544     /// `true` if getting an iterator element may have side effects.
 545     /// Remember to take inner iterators into account.
 546     const MAY_HAVE_SIDE_EFFECT: bool;
 547 }
 548
 549 /// Like `Iterator::__iterator_get_unchecked`, but doesn't require the compiler to
 550 /// know that `U: TrustedRandomAccess`.
 551 ///
 552 /// ## Safety
 553 ///
 554 /// Same requirements calling `get_unchecked` directly.
 555 #[doc(hidden)]
 556 pub(in crate::iter::adapters) unsafe fn try_get_unchecked<I>(it: &mut I, idx: usize) -> I::Item
 557 where
 558     I: Iterator,
 559 {
 560     // SAFETY: the caller must uphold the contract for
 561     // `Iterator::__iterator_get_unchecked`.
 562     unsafe { it.try_get_unchecked(idx) }
 563 }
 564
 565 unsafe trait SpecTrustedRandomAccess: Iterator {
 566     /// If `Self: TrustedRandomAccess`, it must be safe to call a
 567     /// `Iterator::__iterator_get_unchecked(self, index)`.
 568     unsafe fn try_get_unchecked(&mut self, index: usize) -> Self::Item;
 569 }
 570
 571 unsafe impl<I: Iterator> SpecTrustedRandomAccess for I {
 572     default unsafe fn try_get_unchecked(&mut self, _: usize) -> Self::Item {
 573         panic!("Should only be called on TrustedRandomAccess iterators");
 574     }
 575 }
 576
 577 unsafe impl<I: Iterator + TrustedRandomAccessNoCoerce> SpecTrustedRandomAccess for I {
 578     unsafe fn try_get_unchecked(&mut self, index: usize) -> Self::Item {
 579         // SAFETY: the caller must uphold the contract for
 580         // `Iterator::__iterator_get_unchecked`.
 581         unsafe { self.__iterator_get_unchecked(index) }
 582     }
 583 }