vendor/rayon/src/iter/collect/mod.rs

   1 use super::{IndexedParallelIterator, ParallelIterator};
   2
   3 mod consumer;
   4 use self::consumer::CollectConsumer;
   5 use self::consumer::CollectResult;
   6 use super::unzip::unzip_indexed;
   7
   8 mod test;
   9
  10 /// Collects the results of the exact iterator into the specified vector.
  11 ///
  12 /// This is called by `IndexedParallelIterator::collect_into_vec`.
  13 pub(super) fn collect_into_vec<I, T>(pi: I, v: &mut Vec<T>)
  14 where
  15     I: IndexedParallelIterator<Item = T>,
  16     T: Send,
  17 {
  18     v.truncate(0); // clear any old data
  19     let len = pi.len();
  20     collect_with_consumer(v, len, |consumer| pi.drive(consumer));
  21 }
  22
  23 /// Collects the results of the iterator into the specified vector.
  24 ///
  25 /// Technically, this only works for `IndexedParallelIterator`, but we're faking a
  26 /// bit of specialization here until Rust can do that natively.  Callers are
  27 /// using `opt_len` to find the length before calling this, and only exact
  28 /// iterators will return anything but `None` there.
  29 ///
  30 /// Since the type system doesn't understand that contract, we have to allow
  31 /// *any* `ParallelIterator` here, and `CollectConsumer` has to also implement
  32 /// `UnindexedConsumer`.  That implementation panics `unreachable!` in case
  33 /// there's a bug where we actually do try to use this unindexed.
  34 pub(super) fn special_extend<I, T>(pi: I, len: usize, v: &mut Vec<T>)
  35 where
  36     I: ParallelIterator<Item = T>,
  37     T: Send,
  38 {
  39     collect_with_consumer(v, len, |consumer| pi.drive_unindexed(consumer));
  40 }
  41
  42 /// Unzips the results of the exact iterator into the specified vectors.
  43 ///
  44 /// This is called by `IndexedParallelIterator::unzip_into_vecs`.
  45 pub(super) fn unzip_into_vecs<I, A, B>(pi: I, left: &mut Vec<A>, right: &mut Vec<B>)
  46 where
  47     I: IndexedParallelIterator<Item = (A, B)>,
  48     A: Send,
  49     B: Send,
  50 {
  51     // clear any old data
  52     left.truncate(0);
  53     right.truncate(0);
  54
  55     let len = pi.len();
  56     collect_with_consumer(right, len, |right_consumer| {
  57         let mut right_result = None;
  58         collect_with_consumer(left, len, |left_consumer| {
  59             let (left_r, right_r) = unzip_indexed(pi, left_consumer, right_consumer);
  60             right_result = Some(right_r);
  61             left_r
  62         });
  63         right_result.unwrap()
  64     });
  65 }
  66
  67 /// Create a consumer on the slice of memory we are collecting into.
  68 ///
  69 /// The consumer needs to be used inside the scope function, and the
  70 /// complete collect result passed back.
  71 ///
  72 /// This method will verify the collect result, and panic if the slice
  73 /// was not fully written into. Otherwise, in the successful case,
  74 /// the vector is complete with the collected result.
  75 fn collect_with_consumer<T, F>(vec: &mut Vec<T>, len: usize, scope_fn: F)
  76 where
  77     T: Send,
  78     F: FnOnce(CollectConsumer<'_, T>) -> CollectResult<'_, T>,
  79 {
  80     // Reserve space for `len` more elements in the vector,
  81     vec.reserve(len);
  82
  83     // Create the consumer and run the callback for collection.
  84     let result = scope_fn(CollectConsumer::appender(vec, len));
  85
  86     // The `CollectResult` represents a contiguous part of the slice, that has
  87     // been written to. On unwind here, the `CollectResult` will be dropped. If
  88     // some producers on the way did not produce enough elements, partial
  89     // `CollectResult`s may have been dropped without being reduced to the final
  90     // result, and we will see that as the length coming up short.
  91     //
  92     // Here, we assert that added length is fully initialized. This is checked
  93     // by the following assert, which verifies if a complete `CollectResult`
  94     // was produced; if the length is correct, it is necessarily covering the
  95     // target slice. Since we know that the consumer cannot have escaped from
  96     // `drive` (by parametricity, essentially), we know that any stores that
  97     // will happen, have happened. Unless some code is buggy, that means we
  98     // should have seen `len` total writes.
  99     let actual_writes = result.len();
 100     assert!(
 101         actual_writes == len,
 102         "expected {} total writes, but got {}",
 103         len,
 104         actual_writes
 105     );
 106
 107     // Release the result's mutable borrow and "proxy ownership"
 108     // of the elements, before the vector takes it over.
 109     result.release_ownership();
 110
 111     let new_len = vec.len() + len;
 112
 113     unsafe {
 114         vec.set_len(new_len);
 115     }
 116 }