use super::plumbing::*;
use super::*;
struct BlocksCallback<S, C> {
sizes: S,
consumer: C,
len: usize,
}
impl<T, S, C> ProducerCallback<T> for BlocksCallback<S, C>
where
C: UnindexedConsumer<T>,
S: Iterator<Item = usize>,
{
type Output = C::Result;
fn callback<P: Producer<Item = T>>(mut self, mut producer: P) -> Self::Output {
let mut remaining_len = self.len;
let mut consumer = self.consumer;
let (left_consumer, right_consumer, _) = consumer.split_at(0);
let mut leftmost_res = left_consumer.into_folder().complete();
consumer = right_consumer;
while remaining_len > 0 && !consumer.full() {
let size = self.sizes.next().unwrap_or(std::usize::MAX);
let capped_size = remaining_len.min(size);
remaining_len -= capped_size;
let (left_producer, right_producer) = producer.split_at(capped_size);
producer = right_producer;
let (left_consumer, right_consumer, _) = consumer.split_at(capped_size);
consumer = right_consumer;
leftmost_res = consumer.to_reducer().reduce(
leftmost_res,
bridge_producer_consumer(capped_size, left_producer, left_consumer),
);
}
leftmost_res
}
}
#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
#[derive(Debug, Clone)]
pub struct ExponentialBlocks<I> {
base: I,
}
impl<I> ExponentialBlocks<I> {
pub(super) fn new(base: I) -> Self {
Self { base }
}
}
impl<I> ParallelIterator for ExponentialBlocks<I>
where
I: IndexedParallelIterator,
{
type Item = I::Item;
fn drive_unindexed<C>(self, consumer: C) -> C::Result
where
C: UnindexedConsumer<Self::Item>,
{
let first = crate::current_num_threads();
let callback = BlocksCallback {
consumer,
sizes: std::iter::successors(Some(first), exponential_size),
len: self.base.len(),
};
self.base.with_producer(callback)
}
}
fn exponential_size(size: &usize) -> Option<usize> {
Some(size.saturating_mul(2))
}
#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
#[derive(Debug, Clone)]
pub struct UniformBlocks<I> {
base: I,
block_size: usize,
}
impl<I> UniformBlocks<I> {
pub(super) fn new(base: I, block_size: usize) -> Self {
Self { base, block_size }
}
}
impl<I> ParallelIterator for UniformBlocks<I>
where
I: IndexedParallelIterator,
{
type Item = I::Item;
fn drive_unindexed<C>(self, consumer: C) -> C::Result
where
C: UnindexedConsumer<Self::Item>,
{
let callback = BlocksCallback {
consumer,
sizes: std::iter::repeat(self.block_size),
len: self.base.len(),
};
self.base.with_producer(callback)
}
}