Source code
Revision control
Copy as Markdown
Other Tools
use crate::raw::Bucket;
use crate::raw::{Allocator, Global, RawIter, RawIterRange, RawTable};
use crate::scopeguard::guard;
use core::marker::PhantomData;
use core::mem;
use core::ptr::NonNull;
use rayon::iter::{
plumbing::{self, Folder, UnindexedConsumer, UnindexedProducer},
ParallelIterator,
};
/// Parallel iterator which returns a raw pointer to every full bucket in the table.
pub struct RawParIter<T> {
iter: RawIterRange<T>,
}
impl<T> RawParIter<T> {
#[cfg_attr(feature = "inline-more", inline)]
pub(super) unsafe fn iter(&self) -> RawIterRange<T> {
self.iter.clone()
}
}
impl<T> Clone for RawParIter<T> {
#[cfg_attr(feature = "inline-more", inline)]
fn clone(&self) -> Self {
Self {
iter: self.iter.clone(),
}
}
}
impl<T> From<RawIter<T>> for RawParIter<T> {
fn from(it: RawIter<T>) -> Self {
RawParIter { iter: it.iter }
}
}
impl<T> ParallelIterator for RawParIter<T> {
type Item = Bucket<T>;
#[cfg_attr(feature = "inline-more", inline)]
fn drive_unindexed<C>(self, consumer: C) -> C::Result
where
C: UnindexedConsumer<Self::Item>,
{
let producer = ParIterProducer { iter: self.iter };
plumbing::bridge_unindexed(producer, consumer)
}
}
/// Producer which returns a `Bucket<T>` for every element.
struct ParIterProducer<T> {
iter: RawIterRange<T>,
}
impl<T> UnindexedProducer for ParIterProducer<T> {
type Item = Bucket<T>;
#[cfg_attr(feature = "inline-more", inline)]
fn split(self) -> (Self, Option<Self>) {
let (left, right) = self.iter.split();
let left = ParIterProducer { iter: left };
let right = right.map(|right| ParIterProducer { iter: right });
(left, right)
}
#[cfg_attr(feature = "inline-more", inline)]
fn fold_with<F>(self, folder: F) -> F
where
F: Folder<Self::Item>,
{
folder.consume_iter(self.iter)
}
}
/// Parallel iterator which consumes a table and returns elements.
pub struct RawIntoParIter<T, A: Allocator = Global> {
table: RawTable<T, A>,
}
impl<T, A: Allocator> RawIntoParIter<T, A> {
#[cfg_attr(feature = "inline-more", inline)]
pub(super) unsafe fn par_iter(&self) -> RawParIter<T> {
self.table.par_iter()
}
}
impl<T: Send, A: Allocator + Send> ParallelIterator for RawIntoParIter<T, A> {
type Item = T;
#[cfg_attr(feature = "inline-more", inline)]
fn drive_unindexed<C>(self, consumer: C) -> C::Result
where
C: UnindexedConsumer<Self::Item>,
{
let iter = unsafe { self.table.iter().iter };
let _guard = guard(self.table.into_allocation(), |alloc| {
if let Some((ptr, layout, ref alloc)) = *alloc {
unsafe {
alloc.deallocate(ptr, layout);
}
}
});
let producer = ParDrainProducer { iter };
plumbing::bridge_unindexed(producer, consumer)
}
}
/// Parallel iterator which consumes elements without freeing the table storage.
pub struct RawParDrain<'a, T, A: Allocator = Global> {
// We don't use a &'a mut RawTable<T> because we want RawParDrain to be
// covariant over T.
table: NonNull<RawTable<T, A>>,
marker: PhantomData<&'a RawTable<T, A>>,
}
unsafe impl<T: Send, A: Allocator> Send for RawParDrain<'_, T, A> {}
impl<T, A: Allocator> RawParDrain<'_, T, A> {
#[cfg_attr(feature = "inline-more", inline)]
pub(super) unsafe fn par_iter(&self) -> RawParIter<T> {
self.table.as_ref().par_iter()
}
}
impl<T: Send, A: Allocator> ParallelIterator for RawParDrain<'_, T, A> {
type Item = T;
#[cfg_attr(feature = "inline-more", inline)]
fn drive_unindexed<C>(self, consumer: C) -> C::Result
where
C: UnindexedConsumer<Self::Item>,
{
let _guard = guard(self.table, |table| unsafe {
table.as_mut().clear_no_drop();
});
let iter = unsafe { self.table.as_ref().iter().iter };
mem::forget(self);
let producer = ParDrainProducer { iter };
plumbing::bridge_unindexed(producer, consumer)
}
}
impl<T, A: Allocator> Drop for RawParDrain<'_, T, A> {
fn drop(&mut self) {
// If drive_unindexed is not called then simply clear the table.
unsafe {
self.table.as_mut().clear();
}
}
}
/// Producer which will consume all elements in the range, even if it is dropped
/// halfway through.
struct ParDrainProducer<T> {
iter: RawIterRange<T>,
}
impl<T: Send> UnindexedProducer for ParDrainProducer<T> {
type Item = T;
#[cfg_attr(feature = "inline-more", inline)]
fn split(self) -> (Self, Option<Self>) {
let (left, right) = self.iter.clone().split();
mem::forget(self);
let left = ParDrainProducer { iter: left };
let right = right.map(|right| ParDrainProducer { iter: right });
(left, right)
}
#[cfg_attr(feature = "inline-more", inline)]
fn fold_with<F>(mut self, mut folder: F) -> F
where
F: Folder<Self::Item>,
{
// Make sure to modify the iterator in-place so that any remaining
// elements are processed in our Drop impl.
for item in &mut self.iter {
folder = folder.consume(unsafe { item.read() });
if folder.full() {
return folder;
}
}
// If we processed all elements then we don't need to run the drop.
mem::forget(self);
folder
}
}
impl<T> Drop for ParDrainProducer<T> {
#[cfg_attr(feature = "inline-more", inline)]
fn drop(&mut self) {
// Drop all remaining elements
if mem::needs_drop::<T>() {
for item in &mut self.iter {
unsafe {
item.drop();
}
}
}
}
}
impl<T, A: Allocator> RawTable<T, A> {
/// Returns a parallel iterator over the elements in a `RawTable`.
#[cfg_attr(feature = "inline-more", inline)]
pub unsafe fn par_iter(&self) -> RawParIter<T> {
RawParIter {
iter: self.iter().iter,
}
}
/// Returns a parallel iterator over the elements in a `RawTable`.
#[cfg_attr(feature = "inline-more", inline)]
pub fn into_par_iter(self) -> RawIntoParIter<T, A> {
RawIntoParIter { table: self }
}
/// Returns a parallel iterator which consumes all elements of a `RawTable`
/// without freeing its memory allocation.
#[cfg_attr(feature = "inline-more", inline)]
pub fn par_drain(&mut self) -> RawParDrain<'_, T, A> {
RawParDrain {
table: NonNull::from(self),
marker: PhantomData,
}
}
}