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: TrustedRandomAccess,
  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 + TrustedRandomAccess;
 130 }
 131
 132 // General Zip impl
 133 #[doc(hidden)]
 134 impl<A, B> ZipImpl<A, B> for Zip<A, B>
 135 where
 136     A: Iterator,
 137     B: Iterator,
 138 {
 139     type Item = (A::Item, B::Item);
 140     default fn new(a: A, b: B) -> Self {
 141         Zip {
 142             a,
 143             b,
 144             index: 0, // unused
 145             len: 0,   // unused
 146             a_len: 0, // unused
 147         }
 148     }
 149
 150     #[inline]
 151     default fn next(&mut self) -> Option<(A::Item, B::Item)> {
 152         let x = self.a.next()?;
 153         let y = self.b.next()?;
 154         Some((x, y))
 155     }
 156
 157     #[inline]
 158     default fn nth(&mut self, n: usize) -> Option<Self::Item> {
 159         self.super_nth(n)
 160     }
 161
 162     #[inline]
 163     default fn next_back(&mut self) -> Option<(A::Item, B::Item)>
 164     where
 165         A: DoubleEndedIterator + ExactSizeIterator,
 166         B: DoubleEndedIterator + ExactSizeIterator,
 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     #[inline]
 190     default fn size_hint(&self) -> (usize, Option<usize>) {
 191         let (a_lower, a_upper) = self.a.size_hint();
 192         let (b_lower, b_upper) = self.b.size_hint();
 193
 194         let lower = cmp::min(a_lower, b_lower);
 195
 196         let upper = match (a_upper, b_upper) {
 197             (Some(x), Some(y)) => Some(cmp::min(x, y)),
 198             (Some(x), None) => Some(x),
 199             (None, Some(y)) => Some(y),
 200             (None, None) => None,
 201         };
 202
 203         (lower, upper)
 204     }
 205
 206     default unsafe fn get_unchecked(&mut self, _idx: usize) -> <Self as Iterator>::Item
 207     where
 208         Self: TrustedRandomAccess,
 209     {
 210         unreachable!("Always specialized");
 211     }
 212 }
 213
 214 #[doc(hidden)]
 215 impl<A, B> ZipImpl<A, B> for Zip<A, B>
 216 where
 217     A: TrustedRandomAccess + Iterator,
 218     B: TrustedRandomAccess + Iterator,
 219 {
 220     fn new(a: A, b: B) -> Self {
 221         let a_len = a.size();
 222         let len = cmp::min(a_len, b.size());
 223         Zip { a, b, index: 0, len, a_len }
 224     }
 225
 226     #[inline]
 227     fn next(&mut self) -> Option<(A::Item, B::Item)> {
 228         if self.index < self.len {
 229             let i = self.index;
 230             // since get_unchecked executes code which can panic we increment the counters beforehand
 231             // so that the same index won't be accessed twice, as required by TrustedRandomAccess
 232             self.index += 1;
 233             // SAFETY: `i` is smaller than `self.len`, thus smaller than `self.a.len()` and `self.b.len()`
 234             unsafe {
 235                 Some((self.a.__iterator_get_unchecked(i), self.b.__iterator_get_unchecked(i)))
 236             }
 237         } else if A::MAY_HAVE_SIDE_EFFECT && self.index < self.a_len {
 238             let i = self.index;
 239             // as above, increment before executing code that may panic
 240             self.index += 1;
 241             self.len += 1;
 242             // match the base implementation's potential side effects
 243             // SAFETY: we just checked that `i` < `self.a.len()`
 244             unsafe {
 245                 self.a.__iterator_get_unchecked(i);
 246             }
 247             None
 248         } else {
 249             None
 250         }
 251     }
 252
 253     #[inline]
 254     fn size_hint(&self) -> (usize, Option<usize>) {
 255         let len = self.len - self.index;
 256         (len, Some(len))
 257     }
 258
 259     #[inline]
 260     fn nth(&mut self, n: usize) -> Option<Self::Item> {
 261         let delta = cmp::min(n, self.len - self.index);
 262         let end = self.index + delta;
 263         while self.index < end {
 264             let i = self.index;
 265             // since get_unchecked executes code which can panic we increment the counters beforehand
 266             // so that the same index won't be accessed twice, as required by TrustedRandomAccess
 267             self.index += 1;
 268             if A::MAY_HAVE_SIDE_EFFECT {
 269                 // SAFETY: the usage of `cmp::min` to calculate `delta`
 270                 // ensures that `end` is smaller than or equal to `self.len`,
 271                 // so `i` is also smaller than `self.len`.
 272                 unsafe {
 273                     self.a.__iterator_get_unchecked(i);
 274                 }
 275             }
 276             if B::MAY_HAVE_SIDE_EFFECT {
 277                 // SAFETY: same as above.
 278                 unsafe {
 279                     self.b.__iterator_get_unchecked(i);
 280                 }
 281             }
 282         }
 283
 284         self.super_nth(n - delta)
 285     }
 286
 287     #[inline]
 288     fn next_back(&mut self) -> Option<(A::Item, B::Item)>
 289     where
 290         A: DoubleEndedIterator + ExactSizeIterator,
 291         B: DoubleEndedIterator + ExactSizeIterator,
 292     {
 293         if A::MAY_HAVE_SIDE_EFFECT || B::MAY_HAVE_SIDE_EFFECT {
 294             let sz_a = self.a.size();
 295             let sz_b = self.b.size();
 296             // Adjust a, b to equal length, make sure that only the first call
 297             // of `next_back` does this, otherwise we will break the restriction
 298             // on calls to `self.next_back()` after calling `get_unchecked()`.
 299             if sz_a != sz_b {
 300                 let sz_a = self.a.size();
 301                 if A::MAY_HAVE_SIDE_EFFECT && sz_a > self.len {
 302                     for _ in 0..sz_a - self.len {
 303                         // since next_back() may panic we increment the counters beforehand
 304                         // to keep Zip's state in sync with the underlying iterator source
 305                         self.a_len -= 1;
 306                         self.a.next_back();
 307                     }
 308                     debug_assert_eq!(self.a_len, self.len);
 309                 }
 310                 let sz_b = self.b.size();
 311                 if B::MAY_HAVE_SIDE_EFFECT && sz_b > self.len {
 312                     for _ in 0..sz_b - self.len {
 313                         self.b.next_back();
 314                     }
 315                 }
 316             }
 317         }
 318         if self.index < self.len {
 319             // since get_unchecked executes code which can panic we increment the counters beforehand
 320             // so that the same index won't be accessed twice, as required by TrustedRandomAccess
 321             self.len -= 1;
 322             self.a_len -= 1;
 323             let i = self.len;
 324             // SAFETY: `i` is smaller than the previous value of `self.len`,
 325             // which is also smaller than or equal to `self.a.len()` and `self.b.len()`
 326             unsafe {
 327                 Some((self.a.__iterator_get_unchecked(i), self.b.__iterator_get_unchecked(i)))
 328             }
 329         } else {
 330             None
 331         }
 332     }
 333
 334     #[inline]
 335     unsafe fn get_unchecked(&mut self, idx: usize) -> <Self as Iterator>::Item {
 336         let idx = self.index + idx;
 337         // SAFETY: the caller must uphold the contract for
 338         // `Iterator::__iterator_get_unchecked`.
 339         unsafe { (self.a.__iterator_get_unchecked(idx), self.b.__iterator_get_unchecked(idx)) }
 340     }
 341 }
 342
 343 #[stable(feature = "rust1", since = "1.0.0")]
 344 impl<A, B> ExactSizeIterator for Zip<A, B>
 345 where
 346     A: ExactSizeIterator,
 347     B: ExactSizeIterator,
 348 {
 349 }
 350
 351 #[doc(hidden)]
 352 #[unstable(feature = "trusted_random_access", issue = "none")]
 353 unsafe impl<A, B> TrustedRandomAccess for Zip<A, B>
 354 where
 355     A: TrustedRandomAccess,
 356     B: TrustedRandomAccess,
 357 {
 358     const MAY_HAVE_SIDE_EFFECT: bool = A::MAY_HAVE_SIDE_EFFECT || B::MAY_HAVE_SIDE_EFFECT;
 359 }
 360
 361 #[stable(feature = "fused", since = "1.26.0")]
 362 impl<A, B> FusedIterator for Zip<A, B>
 363 where
 364     A: FusedIterator,
 365     B: FusedIterator,
 366 {
 367 }
 368
 369 #[unstable(feature = "trusted_len", issue = "37572")]
 370 unsafe impl<A, B> TrustedLen for Zip<A, B>
 371 where
 372     A: TrustedLen,
 373     B: TrustedLen,
 374 {
 375 }
 376
 377 // Arbitrarily selects the left side of the zip iteration as extractable "source"
 378 // it would require negative trait bounds to be able to try both
 379 #[unstable(issue = "none", feature = "inplace_iteration")]
 380 unsafe impl<S, A, B> SourceIter for Zip<A, B>
 381 where
 382     A: SourceIter<Source = S>,
 383     B: Iterator,
 384     S: Iterator,
 385 {
 386     type Source = S;
 387
 388     #[inline]
 389     unsafe fn as_inner(&mut self) -> &mut S {
 390         // SAFETY: unsafe function forwarding to unsafe function with the same requirements
 391         unsafe { SourceIter::as_inner(&mut self.a) }
 392     }
 393 }
 394
 395 #[unstable(issue = "none", feature = "inplace_iteration")]
 396 // Limited to Item: Copy since interaction between Zip's use of TrustedRandomAccess
 397 // and Drop implementation of the source is unclear.
 398 //
 399 // An additional method returning the number of times the source has been logically advanced
 400 // (without calling next()) would be needed to properly drop the remainder of the source.
 401 unsafe impl<A: InPlaceIterable, B: Iterator> InPlaceIterable for Zip<A, B> where A::Item: Copy {}
 402
 403 #[stable(feature = "rust1", since = "1.0.0")]
 404 impl<A: Debug, B: Debug> Debug for Zip<A, B> {
 405     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 406         ZipFmt::fmt(self, f)
 407     }
 408 }
 409
 410 trait ZipFmt<A, B> {
 411     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result;
 412 }
 413
 414 impl<A: Debug, B: Debug> ZipFmt<A, B> for Zip<A, B> {
 415     default fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 416         f.debug_struct("Zip").field("a", &self.a).field("b", &self.b).finish()
 417     }
 418 }
 419
 420 impl<A: Debug + TrustedRandomAccess, B: Debug + TrustedRandomAccess> ZipFmt<A, B> for Zip<A, B> {
 421     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 422         // It's *not safe* to call fmt on the contained iterators, since once
 423         // we start iterating they're in strange, potentially unsafe, states.
 424         f.debug_struct("Zip").finish()
 425     }
 426 }
 427
 428 /// An iterator whose items are random-accessible efficiently
 429 ///
 430 /// # Safety
 431 ///
 432 /// The iterator's `size_hint` must be exact and cheap to call.
 433 ///
 434 /// `size` may not be overridden.
 435 ///
 436 /// `<Self as Iterator>::__iterator_get_unchecked` must be safe to call
 437 /// provided the following conditions are met.
 438 ///
 439 /// 1. `0 <= idx` and `idx < self.size()`.
 440 /// 2. If `self: !Clone`, then `get_unchecked` is never called with the same
 441 ///    index on `self` more than once.
 442 /// 3. After `self.get_unchecked(idx)` has been called then `next_back` will
 443 ///    only be called at most `self.size() - idx - 1` times.
 444 /// 4. After `get_unchecked` is called, then only the following methods will be
 445 ///    called on `self`:
 446 ///     * `std::clone::Clone::clone()`
 447 ///     * `std::iter::Iterator::size_hint()`
 448 ///     * `std::iter::DoubleEndedIterator::next_back()`
 449 ///     * `std::iter::Iterator::__iterator_get_unchecked()`
 450 ///     * `std::iter::TrustedRandomAccess::size()`
 451 ///
 452 /// Further, given that these conditions are met, it must guarantee that:
 453 ///
 454 /// * It does not change the value returned from `size_hint`
 455 /// * It must be safe to call the methods listed above on `self` after calling
 456 ///   `get_unchecked`, assuming that the required traits are implemented.
 457 /// * It must also be safe to drop `self` after calling `get_unchecked`.
 458 #[doc(hidden)]
 459 #[unstable(feature = "trusted_random_access", issue = "none")]
 460 #[rustc_specialization_trait]
 461 pub unsafe trait TrustedRandomAccess: Sized {
 462     // Convenience method.
 463     fn size(&self) -> usize
 464     where
 465         Self: Iterator,
 466     {
 467         self.size_hint().0
 468     }
 469     /// `true` if getting an iterator element may have side effects.
 470     /// Remember to take inner iterators into account.
 471     const MAY_HAVE_SIDE_EFFECT: bool;
 472 }
 473
 474 /// Like `Iterator::__iterator_get_unchecked`, but doesn't require the compiler to
 475 /// know that `U: TrustedRandomAccess`.
 476 ///
 477 /// ## Safety
 478 ///
 479 /// Same requirements calling `get_unchecked` directly.
 480 #[doc(hidden)]
 481 pub(in crate::iter::adapters) unsafe fn try_get_unchecked<I>(it: &mut I, idx: usize) -> I::Item
 482 where
 483     I: Iterator,
 484 {
 485     // SAFETY: the caller must uphold the contract for
 486     // `Iterator::__iterator_get_unchecked`.
 487     unsafe { it.try_get_unchecked(idx) }
 488 }
 489
 490 unsafe trait SpecTrustedRandomAccess: Iterator {
 491     /// If `Self: TrustedRandomAccess`, it must be safe to call a
 492     /// `Iterator::__iterator_get_unchecked(self, index)`.
 493     unsafe fn try_get_unchecked(&mut self, index: usize) -> Self::Item;
 494 }
 495
 496 unsafe impl<I: Iterator> SpecTrustedRandomAccess for I {
 497     default unsafe fn try_get_unchecked(&mut self, _: usize) -> Self::Item {
 498         panic!("Should only be called on TrustedRandomAccess iterators");
 499     }
 500 }
 501
 502 unsafe impl<I: Iterator + TrustedRandomAccess> SpecTrustedRandomAccess for I {
 503     unsafe fn try_get_unchecked(&mut self, index: usize) -> Self::Item {
 504         // SAFETY: the caller must uphold the contract for
 505         // `Iterator::__iterator_get_unchecked`.
 506         unsafe { self.__iterator_get_unchecked(index) }
 507     }
 508 }