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

   1 use crate::{
   2     cmp,
   3     fmt::{self, Debug},
   4     iter::{
   5         DoubleEndedIterator, ExactSizeIterator, FusedIterator, InPlaceIterable, Iterator,
   6         SourceIter, TrustedLen,
   7     },
   8 };
   9
  10 /// An iterator that iterates two other iterators simultaneously.
  11 ///
  12 /// This `struct` is created by [`Iterator::zip`]. See its documentation
  13 /// for more.
  14 #[derive(Clone)]
  15 #[must_use = "iterators are lazy and do nothing unless consumed"]
  16 #[stable(feature = "rust1", since = "1.0.0")]
  17 pub struct Zip<A, B> {
  18     a: A,
  19     b: B,
  20     // index and len are only used by the specialized version of zip
  21     index: usize,
  22     len: usize,
  23 }
  24 impl<A: Iterator, B: Iterator> Zip<A, B> {
  25     pub(in crate::iter) fn new(a: A, b: B) -> Zip<A, B> {
  26         ZipImpl::new(a, b)
  27     }
  28     fn super_nth(&mut self, mut n: usize) -> Option<(A::Item, B::Item)> {
  29         while let Some(x) = Iterator::next(self) {
  30             if n == 0 {
  31                 return Some(x);
  32             }
  33             n -= 1;
  34         }
  35         None
  36     }
  37 }
  38
  39 #[stable(feature = "rust1", since = "1.0.0")]
  40 impl<A, B> Iterator for Zip<A, B>
  41 where
  42     A: Iterator,
  43     B: Iterator,
  44 {
  45     type Item = (A::Item, B::Item);
  46
  47     #[inline]
  48     fn next(&mut self) -> Option<Self::Item> {
  49         ZipImpl::next(self)
  50     }
  51
  52     #[inline]
  53     fn size_hint(&self) -> (usize, Option<usize>) {
  54         ZipImpl::size_hint(self)
  55     }
  56
  57     #[inline]
  58     fn nth(&mut self, n: usize) -> Option<Self::Item> {
  59         ZipImpl::nth(self, n)
  60     }
  61
  62     #[inline]
  63     unsafe fn __iterator_get_unchecked(&mut self, idx: usize) -> Self::Item
  64     where
  65         Self: TrustedRandomAccess,
  66     {
  67         // SAFETY: `ZipImpl::__iterator_get_unchecked` has same safety
  68         // requirements as `Iterator::__iterator_get_unchecked`.
  69         unsafe { ZipImpl::get_unchecked(self, idx) }
  70     }
  71 }
  72
  73 #[stable(feature = "rust1", since = "1.0.0")]
  74 impl<A, B> DoubleEndedIterator for Zip<A, B>
  75 where
  76     A: DoubleEndedIterator + ExactSizeIterator,
  77     B: DoubleEndedIterator + ExactSizeIterator,
  78 {
  79     #[inline]
  80     fn next_back(&mut self) -> Option<(A::Item, B::Item)> {
  81         ZipImpl::next_back(self)
  82     }
  83 }
  84
  85 // Zip specialization trait
  86 #[doc(hidden)]
  87 trait ZipImpl<A, B> {
  88     type Item;
  89     fn new(a: A, b: B) -> Self;
  90     fn next(&mut self) -> Option<Self::Item>;
  91     fn size_hint(&self) -> (usize, Option<usize>);
  92     fn nth(&mut self, n: usize) -> Option<Self::Item>;
  93     fn next_back(&mut self) -> Option<Self::Item>
  94     where
  95         A: DoubleEndedIterator + ExactSizeIterator,
  96         B: DoubleEndedIterator + ExactSizeIterator;
  97     // This has the same safety requirements as `Iterator::__iterator_get_unchecked`
  98     unsafe fn get_unchecked(&mut self, idx: usize) -> <Self as Iterator>::Item
  99     where
 100         Self: Iterator + TrustedRandomAccess;
 101 }
 102
 103 // General Zip impl
 104 #[doc(hidden)]
 105 impl<A, B> ZipImpl<A, B> for Zip<A, B>
 106 where
 107     A: Iterator,
 108     B: Iterator,
 109 {
 110     type Item = (A::Item, B::Item);
 111     default fn new(a: A, b: B) -> Self {
 112         Zip {
 113             a,
 114             b,
 115             index: 0, // unused
 116             len: 0,   // unused
 117         }
 118     }
 119
 120     #[inline]
 121     default fn next(&mut self) -> Option<(A::Item, B::Item)> {
 122         let x = self.a.next()?;
 123         let y = self.b.next()?;
 124         Some((x, y))
 125     }
 126
 127     #[inline]
 128     default fn nth(&mut self, n: usize) -> Option<Self::Item> {
 129         self.super_nth(n)
 130     }
 131
 132     #[inline]
 133     default fn next_back(&mut self) -> Option<(A::Item, B::Item)>
 134     where
 135         A: DoubleEndedIterator + ExactSizeIterator,
 136         B: DoubleEndedIterator + ExactSizeIterator,
 137     {
 138         let a_sz = self.a.len();
 139         let b_sz = self.b.len();
 140         if a_sz != b_sz {
 141             // Adjust a, b to equal length
 142             if a_sz > b_sz {
 143                 for _ in 0..a_sz - b_sz {
 144                     self.a.next_back();
 145                 }
 146             } else {
 147                 for _ in 0..b_sz - a_sz {
 148                     self.b.next_back();
 149                 }
 150             }
 151         }
 152         match (self.a.next_back(), self.b.next_back()) {
 153             (Some(x), Some(y)) => Some((x, y)),
 154             (None, None) => None,
 155             _ => unreachable!(),
 156         }
 157     }
 158
 159     #[inline]
 160     default fn size_hint(&self) -> (usize, Option<usize>) {
 161         let (a_lower, a_upper) = self.a.size_hint();
 162         let (b_lower, b_upper) = self.b.size_hint();
 163
 164         let lower = cmp::min(a_lower, b_lower);
 165
 166         let upper = match (a_upper, b_upper) {
 167             (Some(x), Some(y)) => Some(cmp::min(x, y)),
 168             (Some(x), None) => Some(x),
 169             (None, Some(y)) => Some(y),
 170             (None, None) => None,
 171         };
 172
 173         (lower, upper)
 174     }
 175
 176     default unsafe fn get_unchecked(&mut self, _idx: usize) -> <Self as Iterator>::Item
 177     where
 178         Self: TrustedRandomAccess,
 179     {
 180         unreachable!("Always specialized");
 181     }
 182 }
 183
 184 #[doc(hidden)]
 185 impl<A, B> ZipImpl<A, B> for Zip<A, B>
 186 where
 187     A: TrustedRandomAccess + Iterator,
 188     B: TrustedRandomAccess + Iterator,
 189 {
 190     fn new(a: A, b: B) -> Self {
 191         let len = cmp::min(a.size(), b.size());
 192         Zip { a, b, index: 0, len }
 193     }
 194
 195     #[inline]
 196     fn next(&mut self) -> Option<(A::Item, B::Item)> {
 197         if self.index < self.len {
 198             let i = self.index;
 199             self.index += 1;
 200             // SAFETY: `i` is smaller than `self.len`, thus smaller than `self.a.len()` and `self.b.len()`
 201             unsafe {
 202                 Some((self.a.__iterator_get_unchecked(i), self.b.__iterator_get_unchecked(i)))
 203             }
 204         } else if A::may_have_side_effect() && self.index < self.a.size() {
 205             // match the base implementation's potential side effects
 206             // SAFETY: we just checked that `self.index` < `self.a.len()`
 207             unsafe {
 208                 self.a.__iterator_get_unchecked(self.index);
 209             }
 210             self.index += 1;
 211             None
 212         } else {
 213             None
 214         }
 215     }
 216
 217     #[inline]
 218     fn size_hint(&self) -> (usize, Option<usize>) {
 219         let len = self.len - self.index;
 220         (len, Some(len))
 221     }
 222
 223     #[inline]
 224     fn nth(&mut self, n: usize) -> Option<Self::Item> {
 225         let delta = cmp::min(n, self.len - self.index);
 226         let end = self.index + delta;
 227         while self.index < end {
 228             let i = self.index;
 229             self.index += 1;
 230             if A::may_have_side_effect() {
 231                 // SAFETY: the usage of `cmp::min` to calculate `delta`
 232                 // ensures that `end` is smaller than or equal to `self.len`,
 233                 // so `i` is also smaller than `self.len`.
 234                 unsafe {
 235                     self.a.__iterator_get_unchecked(i);
 236                 }
 237             }
 238             if B::may_have_side_effect() {
 239                 // SAFETY: same as above.
 240                 unsafe {
 241                     self.b.__iterator_get_unchecked(i);
 242                 }
 243             }
 244         }
 245
 246         self.super_nth(n - delta)
 247     }
 248
 249     #[inline]
 250     fn next_back(&mut self) -> Option<(A::Item, B::Item)>
 251     where
 252         A: DoubleEndedIterator + ExactSizeIterator,
 253         B: DoubleEndedIterator + ExactSizeIterator,
 254     {
 255         let a_side_effect = A::may_have_side_effect();
 256         let b_side_effect = B::may_have_side_effect();
 257         if a_side_effect || b_side_effect {
 258             let sz_a = self.a.size();
 259             let sz_b = self.b.size();
 260             // Adjust a, b to equal length, make sure that only the first call
 261             // of `next_back` does this, otherwise we will break the restriction
 262             // on calls to `self.next_back()` after calling `get_unchecked()`.
 263             if sz_a != sz_b {
 264                 let sz_a = self.a.size();
 265                 if a_side_effect && sz_a > self.len {
 266                     for _ in 0..sz_a - cmp::max(self.len, self.index) {
 267                         self.a.next_back();
 268                     }
 269                 }
 270                 let sz_b = self.b.size();
 271                 if b_side_effect && sz_b > self.len {
 272                     for _ in 0..sz_b - self.len {
 273                         self.b.next_back();
 274                     }
 275                 }
 276             }
 277         }
 278         if self.index < self.len {
 279             self.len -= 1;
 280             let i = self.len;
 281             // SAFETY: `i` is smaller than the previous value of `self.len`,
 282             // which is also smaller than or equal to `self.a.len()` and `self.b.len()`
 283             unsafe {
 284                 Some((self.a.__iterator_get_unchecked(i), self.b.__iterator_get_unchecked(i)))
 285             }
 286         } else {
 287             None
 288         }
 289     }
 290
 291     #[inline]
 292     unsafe fn get_unchecked(&mut self, idx: usize) -> <Self as Iterator>::Item {
 293         // SAFETY: the caller must uphold the contract for
 294         // `Iterator::__iterator_get_unchecked`.
 295         unsafe { (self.a.__iterator_get_unchecked(idx), self.b.__iterator_get_unchecked(idx)) }
 296     }
 297 }
 298
 299 #[stable(feature = "rust1", since = "1.0.0")]
 300 impl<A, B> ExactSizeIterator for Zip<A, B>
 301 where
 302     A: ExactSizeIterator,
 303     B: ExactSizeIterator,
 304 {
 305 }
 306
 307 #[doc(hidden)]
 308 #[unstable(feature = "trusted_random_access", issue = "none")]
 309 unsafe impl<A, B> TrustedRandomAccess for Zip<A, B>
 310 where
 311     A: TrustedRandomAccess,
 312     B: TrustedRandomAccess,
 313 {
 314     fn may_have_side_effect() -> bool {
 315         A::may_have_side_effect() || B::may_have_side_effect()
 316     }
 317 }
 318
 319 #[stable(feature = "fused", since = "1.26.0")]
 320 impl<A, B> FusedIterator for Zip<A, B>
 321 where
 322     A: FusedIterator,
 323     B: FusedIterator,
 324 {
 325 }
 326
 327 #[unstable(feature = "trusted_len", issue = "37572")]
 328 unsafe impl<A, B> TrustedLen for Zip<A, B>
 329 where
 330     A: TrustedLen,
 331     B: TrustedLen,
 332 {
 333 }
 334
 335 // Arbitrarily selects the left side of the zip iteration as extractable "source"
 336 // it would require negative trait bounds to be able to try both
 337 #[unstable(issue = "none", feature = "inplace_iteration")]
 338 unsafe impl<S, A, B> SourceIter for Zip<A, B>
 339 where
 340     A: SourceIter<Source = S>,
 341     B: Iterator,
 342     S: Iterator,
 343 {
 344     type Source = S;
 345
 346     #[inline]
 347     unsafe fn as_inner(&mut self) -> &mut S {
 348         // SAFETY: unsafe function forwarding to unsafe function with the same requirements
 349         unsafe { SourceIter::as_inner(&mut self.a) }
 350     }
 351 }
 352
 353 #[unstable(issue = "none", feature = "inplace_iteration")]
 354 // Limited to Item: Copy since interaction between Zip's use of TrustedRandomAccess
 355 // and Drop implementation of the source is unclear.
 356 //
 357 // An additional method returning the number of times the source has been logically advanced
 358 // (without calling next()) would be needed to properly drop the remainder of the source.
 359 unsafe impl<A: InPlaceIterable, B: Iterator> InPlaceIterable for Zip<A, B> where A::Item: Copy {}
 360
 361 #[stable(feature = "rust1", since = "1.0.0")]
 362 impl<A: Debug, B: Debug> Debug for Zip<A, B> {
 363     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 364         ZipFmt::fmt(self, f)
 365     }
 366 }
 367
 368 trait ZipFmt<A, B> {
 369     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result;
 370 }
 371
 372 impl<A: Debug, B: Debug> ZipFmt<A, B> for Zip<A, B> {
 373     default fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 374         f.debug_struct("Zip").field("a", &self.a).field("b", &self.b).finish()
 375     }
 376 }
 377
 378 impl<A: Debug + TrustedRandomAccess, B: Debug + TrustedRandomAccess> ZipFmt<A, B> for Zip<A, B> {
 379     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 380         // It's *not safe* to call fmt on the contained iterators, since once
 381         // we start iterating they're in strange, potentially unsafe, states.
 382         f.debug_struct("Zip").finish()
 383     }
 384 }
 385
 386 /// An iterator whose items are random-accessible efficiently
 387 ///
 388 /// # Safety
 389 ///
 390 /// The iterator's `size_hint` must be exact and cheap to call.
 391 ///
 392 /// `size` may not be overridden.
 393 ///
 394 /// `<Self as Iterator>::__iterator_get_unchecked` must be safe to call
 395 /// provided the following conditions are met.
 396 ///
 397 /// 1. `0 <= idx` and `idx < self.size()`.
 398 /// 2. If `self: !Clone`, then `get_unchecked` is never called with the same
 399 ///    index on `self` more than once.
 400 /// 3. After `self.get_unchecked(idx)` has been called then `next_back` will
 401 ///    only be called at most `self.size() - idx - 1` times.
 402 /// 4. After `get_unchecked` is called, then only the following methods will be
 403 ///    called on `self`:
 404 ///     * `std::clone::Clone::clone`
 405 ///     * `std::iter::Iterator::size_hint()`
 406 ///     * `std::iter::Iterator::next_back()`
 407 ///     * `std::iter::Iterator::__iterator_get_unchecked()`
 408 ///     * `std::iter::TrustedRandomAccess::size()`
 409 ///
 410 /// Further, given that these conditions are met, it must guarantee that:
 411 ///
 412 /// * It does not change the value returned from `size_hint`
 413 /// * It must be safe to call the methods listed above on `self` after calling
 414 ///   `get_unchecked`, assuming that the required traits are implemented.
 415 /// * It must also be safe to drop `self` after calling `get_unchecked`.
 416 #[doc(hidden)]
 417 #[unstable(feature = "trusted_random_access", issue = "none")]
 418 #[rustc_specialization_trait]
 419 pub unsafe trait TrustedRandomAccess: Sized {
 420     // Convenience method.
 421     fn size(&self) -> usize
 422     where
 423         Self: Iterator,
 424     {
 425         self.size_hint().0
 426     }
 427     /// Returns `true` if getting an iterator element may have
 428     /// side effects. Remember to take inner iterators into account.
 429     fn may_have_side_effect() -> bool;
 430 }
 431
 432 /// Like `Iterator::__iterator_get_unchecked`, but doesn't require the compiler to
 433 /// know that `U: TrustedRandomAccess`.
 434 ///
 435 /// ## Safety
 436 ///
 437 /// Same requirements calling `get_unchecked` directly.
 438 #[doc(hidden)]
 439 pub(in crate::iter::adapters) unsafe fn try_get_unchecked<I>(it: &mut I, idx: usize) -> I::Item
 440 where
 441     I: Iterator,
 442 {
 443     // SAFETY: the caller must uphold the contract for
 444     // `Iterator::__iterator_get_unchecked`.
 445     unsafe { it.try_get_unchecked(idx) }
 446 }
 447
 448 unsafe trait SpecTrustedRandomAccess: Iterator {
 449     /// If `Self: TrustedRandomAccess`, it must be safe to call a
 450     /// `Iterator::__iterator_get_unchecked(self, index)`.
 451     unsafe fn try_get_unchecked(&mut self, index: usize) -> Self::Item;
 452 }
 453
 454 unsafe impl<I: Iterator> SpecTrustedRandomAccess for I {
 455     default unsafe fn try_get_unchecked(&mut self, _: usize) -> Self::Item {
 456         panic!("Should only be called on TrustedRandomAccess iterators");
 457     }
 458 }
 459
 460 unsafe impl<I: Iterator + TrustedRandomAccess> SpecTrustedRandomAccess for I {
 461     unsafe fn try_get_unchecked(&mut self, index: usize) -> Self::Item {
 462         // SAFETY: the caller must uphold the contract for
 463         // `Iterator::__iterator_get_unchecked`.
 464         unsafe { self.__iterator_get_unchecked(index) }
 465     }
 466 }