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