Source code
Revision control
Copy as Markdown
Other Tools
use super::{IndexedParallelIterator, ParallelIterator};
mod consumer;
use self::consumer::CollectConsumer;
use self::consumer::CollectResult;
use super::unzip::unzip_indexed;
mod test;
/// Collects the results of the exact iterator into the specified vector.
///
/// This is called by `IndexedParallelIterator::collect_into_vec`.
pub(super) fn collect_into_vec<I, T>(pi: I, v: &mut Vec<T>)
where
I: IndexedParallelIterator<Item = T>,
T: Send,
{
v.truncate(0); // clear any old data
let len = pi.len();
collect_with_consumer(v, len, |consumer| pi.drive(consumer));
}
/// Collects the results of the iterator into the specified vector.
///
/// Technically, this only works for `IndexedParallelIterator`, but we're faking a
/// bit of specialization here until Rust can do that natively. Callers are
/// using `opt_len` to find the length before calling this, and only exact
/// iterators will return anything but `None` there.
///
/// Since the type system doesn't understand that contract, we have to allow
/// *any* `ParallelIterator` here, and `CollectConsumer` has to also implement
/// `UnindexedConsumer`. That implementation panics `unreachable!` in case
/// there's a bug where we actually do try to use this unindexed.
pub(super) fn special_extend<I, T>(pi: I, len: usize, v: &mut Vec<T>)
where
I: ParallelIterator<Item = T>,
T: Send,
{
collect_with_consumer(v, len, |consumer| pi.drive_unindexed(consumer));
}
/// Unzips the results of the exact iterator into the specified vectors.
///
/// This is called by `IndexedParallelIterator::unzip_into_vecs`.
pub(super) fn 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,
{
// clear any old data
left.truncate(0);
right.truncate(0);
let len = pi.len();
collect_with_consumer(right, len, |right_consumer| {
let mut right_result = None;
collect_with_consumer(left, len, |left_consumer| {
let (left_r, right_r) = unzip_indexed(pi, left_consumer, right_consumer);
right_result = Some(right_r);
left_r
});
right_result.unwrap()
});
}
/// Create a consumer on the slice of memory we are collecting into.
///
/// The consumer needs to be used inside the scope function, and the
/// complete collect result passed back.
///
/// This method will verify the collect result, and panic if the slice
/// was not fully written into. Otherwise, in the successful case,
/// the vector is complete with the collected result.
fn collect_with_consumer<T, F>(vec: &mut Vec<T>, len: usize, scope_fn: F)
where
T: Send,
F: FnOnce(CollectConsumer<'_, T>) -> CollectResult<'_, T>,
{
// Reserve space for `len` more elements in the vector,
vec.reserve(len);
// Create the consumer and run the callback for collection.
let result = scope_fn(CollectConsumer::appender(vec, len));
// The `CollectResult` represents a contiguous part of the slice, that has
// been written to. On unwind here, the `CollectResult` will be dropped. If
// some producers on the way did not produce enough elements, partial
// `CollectResult`s may have been dropped without being reduced to the final
// result, and we will see that as the length coming up short.
//
// Here, we assert that added length is fully initialized. This is checked
// by the following assert, which verifies if a complete `CollectResult`
// was produced; if the length is correct, it is necessarily covering the
// target slice. Since we know that the consumer cannot have escaped from
// `drive` (by parametricity, essentially), we know that any stores that
// will happen, have happened. Unless some code is buggy, that means we
// should have seen `len` total writes.
let actual_writes = result.len();
assert!(
actual_writes == len,
"expected {} total writes, but got {}",
len,
actual_writes
);
// Release the result's mutable borrow and "proxy ownership"
// of the elements, before the vector takes it over.
result.release_ownership();
let new_len = vec.len() + len;
unsafe {
vec.set_len(new_len);
}
}