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`.
collect_into_vec<I, T>(pi: I, v: &mut Vec<T>) where I: IndexedParallelIterator<Item = T>, T: Send,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.
special_extend<I, T>(pi: I, len: usize, v: &mut Vec<T>) where I: ParallelIterator<Item = T>, T: Send,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`.
unzip_into_vecs<I, A, B>(pi: I, left: &mut Vec<A>, right: &mut Vec<B>) where I: IndexedParallelIterator<Item = (A, B)>, A: Send, B: Send,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.
collect_with_consumer<T, F>(vec: &mut Vec<T>, len: usize, scope_fn: F) where T: Send, F: FnOnce(CollectConsumer<'_, T>) -> CollectResult<'_, T>,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 }
117