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 /// use std::iter::zip;
  44 ///
  45 /// let xs = [1, 2, 3];
  46 /// let ys = [4, 5, 6];
  47 ///
  48 /// let mut iter = zip(xs, ys);
  49 ///
  50 /// assert_eq!(iter.next().unwrap(), (1, 4));
  51 /// assert_eq!(iter.next().unwrap(), (2, 5));
  52 /// assert_eq!(iter.next().unwrap(), (3, 6));
  53 /// assert!(iter.next().is_none());
  54 ///
  55 /// // Nested zips are also possible:
  56 /// let zs = [7, 8, 9];
  57 ///
  58 /// let mut iter = zip(zip(xs, ys), zs);
  59 ///
  60 /// assert_eq!(iter.next().unwrap(), ((1, 4), 7));
  61 /// assert_eq!(iter.next().unwrap(), ((2, 5), 8));
  62 /// assert_eq!(iter.next().unwrap(), ((3, 6), 9));
  63 /// assert!(iter.next().is_none());
  64 /// ```
  65 #[stable(feature = "iter_zip", since = "1.59.0")]
  66 pub fn zip<A, B>(a: A, b: B) -> Zip<A::IntoIter, B::IntoIter>
  67 where
  68     A: IntoIterator,
  69     B: IntoIterator,
  70 {
  71     ZipImpl::new(a.into_iter(), b.into_iter())
  72 }
  73
  74 #[stable(feature = "rust1", since = "1.0.0")]
  75 impl<A, B> Iterator for Zip<A, B>
  76 where
  77     A: Iterator,
  78     B: Iterator,
  79 {
  80     type Item = (A::Item, B::Item);
  81
  82     #[inline]
  83     fn next(&mut self) -> Option<Self::Item> {
  84         ZipImpl::next(self)
  85     }
  86
  87     #[inline]
  88     fn size_hint(&self) -> (usize, Option<usize>) {
  89         ZipImpl::size_hint(self)
  90     }
  91
  92     #[inline]
  93     fn nth(&mut self, n: usize) -> Option<Self::Item> {
  94         ZipImpl::nth(self, n)
  95     }
  96
  97     #[inline]
  98     unsafe fn __iterator_get_unchecked(&mut self, idx: usize) -> Self::Item
  99     where
 100         Self: TrustedRandomAccessNoCoerce,
 101     {
 102         // SAFETY: `ZipImpl::__iterator_get_unchecked` has same safety
 103         // requirements as `Iterator::__iterator_get_unchecked`.
 104         unsafe { ZipImpl::get_unchecked(self, idx) }
 105     }
 106 }
 107
 108 #[stable(feature = "rust1", since = "1.0.0")]
 109 impl<A, B> DoubleEndedIterator for Zip<A, B>
 110 where
 111     A: DoubleEndedIterator + ExactSizeIterator,
 112     B: DoubleEndedIterator + ExactSizeIterator,
 113 {
 114     #[inline]
 115     fn next_back(&mut self) -> Option<(A::Item, B::Item)> {
 116         ZipImpl::next_back(self)
 117     }
 118 }
 119
 120 // Zip specialization trait
 121 #[doc(hidden)]
 122 trait ZipImpl<A, B> {
 123     type Item;
 124     fn new(a: A, b: B) -> Self;
 125     fn next(&mut self) -> Option<Self::Item>;
 126     fn size_hint(&self) -> (usize, Option<usize>);
 127     fn nth(&mut self, n: usize) -> Option<Self::Item>;
 128     fn next_back(&mut self) -> Option<Self::Item>
 129     where
 130         A: DoubleEndedIterator + ExactSizeIterator,
 131         B: DoubleEndedIterator + ExactSizeIterator;
 132     // This has the same safety requirements as `Iterator::__iterator_get_unchecked`
 133     unsafe fn get_unchecked(&mut self, idx: usize) -> <Self as Iterator>::Item
 134     where
 135         Self: Iterator + TrustedRandomAccessNoCoerce;
 136 }
 137
 138 // Work around limitations of specialization, requiring `default` impls to be repeated
 139 // in intermediary impls.
 140 macro_rules! zip_impl_general_defaults {
 141     () => {
 142         default fn new(a: A, b: B) -> Self {
 143             Zip {
 144                 a,
 145                 b,
 146                 index: 0, // unused
 147                 len: 0,   // unused
 148                 a_len: 0, // unused
 149             }
 150         }
 151
 152         #[inline]
 153         default fn next(&mut self) -> Option<(A::Item, B::Item)> {
 154             let x = self.a.next()?;
 155             let y = self.b.next()?;
 156             Some((x, y))
 157         }
 158
 159         #[inline]
 160         default fn nth(&mut self, n: usize) -> Option<Self::Item> {
 161             self.super_nth(n)
 162         }
 163
 164         #[inline]
 165         default fn next_back(&mut self) -> Option<(A::Item, B::Item)>
 166         where
 167             A: DoubleEndedIterator + ExactSizeIterator,
 168             B: DoubleEndedIterator + ExactSizeIterator,
 169         {
 170             // The function body below only uses `self.a/b.len()` and `self.a/b.next_back()`
 171             // and doesn’t call `next_back` too often, so this implementation is safe in
 172             // the `TrustedRandomAccessNoCoerce` specialization
 173
 174             let a_sz = self.a.len();
 175             let b_sz = self.b.len();
 176             if a_sz != b_sz {
 177                 // Adjust a, b to equal length
 178                 if a_sz > b_sz {
 179                     for _ in 0..a_sz - b_sz {
 180                         self.a.next_back();
 181                     }
 182                 } else {
 183                     for _ in 0..b_sz - a_sz {
 184                         self.b.next_back();
 185                     }
 186                 }
 187             }
 188             match (self.a.next_back(), self.b.next_back()) {
 189                 (Some(x), Some(y)) => Some((x, y)),
 190                 (None, None) => None,
 191                 _ => unreachable!(),
 192             }
 193         }
 194     };
 195 }
 196
 197 // General Zip impl
 198 #[doc(hidden)]
 199 impl<A, B> ZipImpl<A, B> for Zip<A, B>
 200 where
 201     A: Iterator,
 202     B: Iterator,
 203 {
 204     type Item = (A::Item, B::Item);
 205
 206     zip_impl_general_defaults! {}
 207
 208     #[inline]
 209     default fn size_hint(&self) -> (usize, Option<usize>) {
 210         let (a_lower, a_upper) = self.a.size_hint();
 211         let (b_lower, b_upper) = self.b.size_hint();
 212
 213         let lower = cmp::min(a_lower, b_lower);
 214
 215         let upper = match (a_upper, b_upper) {
 216             (Some(x), Some(y)) => Some(cmp::min(x, y)),
 217             (Some(x), None) => Some(x),
 218             (None, Some(y)) => Some(y),
 219             (None, None) => None,
 220         };
 221
 222         (lower, upper)
 223     }
 224
 225     default unsafe fn get_unchecked(&mut self, _idx: usize) -> <Self as Iterator>::Item
 226     where
 227         Self: TrustedRandomAccessNoCoerce,
 228     {
 229         unreachable!("Always specialized");
 230     }
 231 }
 232
 233 #[doc(hidden)]
 234 impl<A, B> ZipImpl<A, B> for Zip<A, B>
 235 where
 236     A: TrustedRandomAccessNoCoerce + Iterator,
 237     B: TrustedRandomAccessNoCoerce + Iterator,
 238 {
 239     zip_impl_general_defaults! {}
 240
 241     #[inline]
 242     default fn size_hint(&self) -> (usize, Option<usize>) {
 243         let size = cmp::min(self.a.size(), self.b.size());
 244         (size, Some(size))
 245     }
 246
 247     #[inline]
 248     unsafe fn get_unchecked(&mut self, idx: usize) -> <Self as Iterator>::Item {
 249         let idx = self.index + idx;
 250         // SAFETY: the caller must uphold the contract for
 251         // `Iterator::__iterator_get_unchecked`.
 252         unsafe { (self.a.__iterator_get_unchecked(idx), self.b.__iterator_get_unchecked(idx)) }
 253     }
 254 }
 255
 256 #[doc(hidden)]
 257 impl<A, B> ZipImpl<A, B> for Zip<A, B>
 258 where
 259     A: TrustedRandomAccess + Iterator,
 260     B: TrustedRandomAccess + Iterator,
 261 {
 262     fn new(a: A, b: B) -> Self {
 263         let a_len = a.size();
 264         let len = cmp::min(a_len, b.size());
 265         Zip { a, b, index: 0, len, a_len }
 266     }
 267
 268     #[inline]
 269     fn next(&mut self) -> Option<(A::Item, B::Item)> {
 270         if self.index < self.len {
 271             let i = self.index;
 272             // since get_unchecked executes code which can panic we increment the counters beforehand
 273             // so that the same index won't be accessed twice, as required by TrustedRandomAccess
 274             self.index += 1;
 275             // SAFETY: `i` is smaller than `self.len`, thus smaller than `self.a.len()` and `self.b.len()`
 276             unsafe {
 277                 Some((self.a.__iterator_get_unchecked(i), self.b.__iterator_get_unchecked(i)))
 278             }
 279         } else if A::MAY_HAVE_SIDE_EFFECT && self.index < self.a_len {
 280             let i = self.index;
 281             // as above, increment before executing code that may panic
 282             self.index += 1;
 283             self.len += 1;
 284             // match the base implementation's potential side effects
 285             // SAFETY: we just checked that `i` < `self.a.len()`
 286             unsafe {
 287                 self.a.__iterator_get_unchecked(i);
 288             }
 289             None
 290         } else {
 291             None
 292         }
 293     }
 294
 295     #[inline]
 296     fn size_hint(&self) -> (usize, Option<usize>) {
 297         let len = self.len - self.index;
 298         (len, Some(len))
 299     }
 300
 301     #[inline]
 302     fn nth(&mut self, n: usize) -> Option<Self::Item> {
 303         let delta = cmp::min(n, self.len - self.index);
 304         let end = self.index + delta;
 305         while self.index < end {
 306             let i = self.index;
 307             // since get_unchecked executes code which can panic we increment the counters beforehand
 308             // so that the same index won't be accessed twice, as required by TrustedRandomAccess
 309             self.index += 1;
 310             if A::MAY_HAVE_SIDE_EFFECT {
 311                 // SAFETY: the usage of `cmp::min` to calculate `delta`
 312                 // ensures that `end` is smaller than or equal to `self.len`,
 313                 // so `i` is also smaller than `self.len`.
 314                 unsafe {
 315                     self.a.__iterator_get_unchecked(i);
 316                 }
 317             }
 318             if B::MAY_HAVE_SIDE_EFFECT {
 319                 // SAFETY: same as above.
 320                 unsafe {
 321                     self.b.__iterator_get_unchecked(i);
 322                 }
 323             }
 324         }
 325
 326         self.super_nth(n - delta)
 327     }
 328
 329     #[inline]
 330     fn next_back(&mut self) -> Option<(A::Item, B::Item)>
 331     where
 332         A: DoubleEndedIterator + ExactSizeIterator,
 333         B: DoubleEndedIterator + ExactSizeIterator,
 334     {
 335         if A::MAY_HAVE_SIDE_EFFECT || B::MAY_HAVE_SIDE_EFFECT {
 336             let sz_a = self.a.size();
 337             let sz_b = self.b.size();
 338             // Adjust a, b to equal length, make sure that only the first call
 339             // of `next_back` does this, otherwise we will break the restriction
 340             // on calls to `self.next_back()` after calling `get_unchecked()`.
 341             if sz_a != sz_b {
 342                 let sz_a = self.a.size();
 343                 if A::MAY_HAVE_SIDE_EFFECT && sz_a > self.len {
 344                     for _ in 0..sz_a - self.len {
 345                         // since next_back() may panic we increment the counters beforehand
 346                         // to keep Zip's state in sync with the underlying iterator source
 347                         self.a_len -= 1;
 348                         self.a.next_back();
 349                     }
 350                     debug_assert_eq!(self.a_len, self.len);
 351                 }
 352                 let sz_b = self.b.size();
 353                 if B::MAY_HAVE_SIDE_EFFECT && sz_b > self.len {
 354                     for _ in 0..sz_b - self.len {
 355                         self.b.next_back();
 356                     }
 357                 }
 358             }
 359         }
 360         if self.index < self.len {
 361             // since get_unchecked executes code which can panic we increment the counters beforehand
 362             // so that the same index won't be accessed twice, as required by TrustedRandomAccess
 363             self.len -= 1;
 364             self.a_len -= 1;
 365             let i = self.len;
 366             // SAFETY: `i` is smaller than the previous value of `self.len`,
 367             // which is also smaller than or equal to `self.a.len()` and `self.b.len()`
 368             unsafe {
 369                 Some((self.a.__iterator_get_unchecked(i), self.b.__iterator_get_unchecked(i)))
 370             }
 371         } else {
 372             None
 373         }
 374     }
 375 }
 376
 377 #[stable(feature = "rust1", since = "1.0.0")]
 378 impl<A, B> ExactSizeIterator for Zip<A, B>
 379 where
 380     A: ExactSizeIterator,
 381     B: ExactSizeIterator,
 382 {
 383 }
 384
 385 #[doc(hidden)]
 386 #[unstable(feature = "trusted_random_access", issue = "none")]
 387 unsafe impl<A, B> TrustedRandomAccess for Zip<A, B>
 388 where
 389     A: TrustedRandomAccess,
 390     B: TrustedRandomAccess,
 391 {
 392 }
 393
 394 #[doc(hidden)]
 395 #[unstable(feature = "trusted_random_access", issue = "none")]
 396 unsafe impl<A, B> TrustedRandomAccessNoCoerce for Zip<A, B>
 397 where
 398     A: TrustedRandomAccessNoCoerce,
 399     B: TrustedRandomAccessNoCoerce,
 400 {
 401     const MAY_HAVE_SIDE_EFFECT: bool = A::MAY_HAVE_SIDE_EFFECT || B::MAY_HAVE_SIDE_EFFECT;
 402 }
 403
 404 #[stable(feature = "fused", since = "1.26.0")]
 405 impl<A, B> FusedIterator for Zip<A, B>
 406 where
 407     A: FusedIterator,
 408     B: FusedIterator,
 409 {
 410 }
 411
 412 #[unstable(feature = "trusted_len", issue = "37572")]
 413 unsafe impl<A, B> TrustedLen for Zip<A, B>
 414 where
 415     A: TrustedLen,
 416     B: TrustedLen,
 417 {
 418 }
 419
 420 // Arbitrarily selects the left side of the zip iteration as extractable "source"
 421 // it would require negative trait bounds to be able to try both
 422 #[unstable(issue = "none", feature = "inplace_iteration")]
 423 unsafe impl<A, B> SourceIter for Zip<A, B>
 424 where
 425     A: SourceIter,
 426 {
 427     type Source = A::Source;
 428
 429     #[inline]
 430     unsafe fn as_inner(&mut self) -> &mut A::Source {
 431         // SAFETY: unsafe function forwarding to unsafe function with the same requirements
 432         unsafe { SourceIter::as_inner(&mut self.a) }
 433     }
 434 }
 435
 436 // Since SourceIter forwards the left hand side we do the same here
 437 #[unstable(issue = "none", feature = "inplace_iteration")]
 438 unsafe impl<A: InPlaceIterable, B: Iterator> InPlaceIterable for Zip<A, B> {}
 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 `SourceIter::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 #[inline]
 557 pub(in crate::iter::adapters) unsafe fn try_get_unchecked<I>(it: &mut I, idx: usize) -> I::Item
 558 where
 559     I: Iterator,
 560 {
 561     // SAFETY: the caller must uphold the contract for
 562     // `Iterator::__iterator_get_unchecked`.
 563     unsafe { it.try_get_unchecked(idx) }
 564 }
 565
 566 unsafe trait SpecTrustedRandomAccess: Iterator {
 567     /// If `Self: TrustedRandomAccess`, it must be safe to call
 568     /// `Iterator::__iterator_get_unchecked(self, index)`.
 569     unsafe fn try_get_unchecked(&mut self, index: usize) -> Self::Item;
 570 }
 571
 572 unsafe impl<I: Iterator> SpecTrustedRandomAccess for I {
 573     default unsafe fn try_get_unchecked(&mut self, _: usize) -> Self::Item {
 574         panic!("Should only be called on TrustedRandomAccess iterators");
 575     }
 576 }
 577
 578 unsafe impl<I: Iterator + TrustedRandomAccessNoCoerce> SpecTrustedRandomAccess for I {
 579     #[inline]
 580     unsafe fn try_get_unchecked(&mut self, index: usize) -> Self::Item {
 581         // SAFETY: the caller must uphold the contract for
 582         // `Iterator::__iterator_get_unchecked`.
 583         unsafe { self.__iterator_get_unchecked(index) }
 584     }
 585 }