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

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