use core::marker::PhantomData;
use rayon::iter::{
plumbing::{
bridge,
Consumer,
Producer,
ProducerCallback,
UnindexedConsumer,
},
IndexedParallelIterator,
ParallelIterator,
};
pub struct RepeatNone<T> {
inner: PhantomData<T>,
count: usize,
}
impl<T> RepeatNone<T> {
pub(crate) fn new(count: usize) -> Self {
Self {
inner: PhantomData,
count,
}
}
}
impl<T> ParallelIterator for RepeatNone<T>
where
T: Send,
{
type Item = Option<T>;
fn drive_unindexed<C>(self, consumer: C) -> C::Result
where
C: UnindexedConsumer<Self::Item>,
{
bridge(self, consumer)
}
fn opt_len(&self) -> Option<usize> {
Some(self.count)
}
}
impl<T> IndexedParallelIterator for RepeatNone<T>
where
T: Send,
{
fn drive<C>(self, consumer: C) -> C::Result
where
C: Consumer<Self::Item>,
{
bridge(self, consumer)
}
fn with_producer<CB>(self, callback: CB) -> CB::Output
where
CB: ProducerCallback<Self::Item>,
{
callback.callback(RepeatNoneProducer {
inner: self.inner,
count: self.count,
})
}
fn len(&self) -> usize {
self.count
}
}
struct RepeatNoneProducer<T> {
inner: PhantomData<T>,
count: usize,
}
impl<T> Producer for RepeatNoneProducer<T>
where
T: Send,
{
type Item = Option<T>;
type IntoIter = RepeatNoneIter<T>;
fn into_iter(self) -> Self::IntoIter {
RepeatNoneIter {
inner: self.inner,
count: self.count,
}
}
fn split_at(self, index: usize) -> (Self, Self) {
(
Self {
inner: self.inner,
count: index,
},
Self {
inner: self.inner,
count: self.count - index,
},
)
}
}
struct RepeatNoneIter<T> {
inner: PhantomData<T>,
count: usize,
}
impl<T> Iterator for RepeatNoneIter<T> {
type Item = Option<T>;
#[inline]
fn next(&mut self) -> Option<Self::Item> {
if self.count > 0 {
self.count -= 1;
Some(None)
} else {
None
}
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
(self.count, Some(self.count))
}
}
impl<T> DoubleEndedIterator for RepeatNoneIter<T> {
#[inline]
fn next_back(&mut self) -> Option<Self::Item> {
self.next()
}
}
impl<T> ExactSizeIterator for RepeatNoneIter<T> {
#[inline]
fn len(&self) -> usize {
self.count
}
}