use crate::{
matrix::csr_matrix::{CsrMatrixRowIter, CsrMatrixRowIterMut},
utils::{ParallelIteratorWrapper, ParallelProducerWrapper},
vec::css::{CssSlice, CssSliceMut},
};
use mop_common_deps::rayon::iter::{
plumbing::bridge,
plumbing::Producer,
plumbing::ProducerCallback,
plumbing::{Consumer, UnindexedConsumer},
IndexedParallelIterator, ParallelIterator,
};
macro_rules! impl_par_iter {
(
$csr_matrix_row_iter:ident,
$css_slice:ident
) => {
impl<'a, T> ParallelIterator for ParallelIteratorWrapper<$csr_matrix_row_iter<'a, T>>
where
T: Send + Sync,
{
type Item = $css_slice<'a, 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.0.len())
}
}
impl<'a, T> IndexedParallelIterator for ParallelIteratorWrapper<$csr_matrix_row_iter<'a, T>>
where
T: Send + Sync,
{
fn with_producer<Cb>(self, callback: Cb) -> Cb::Output
where
Cb: ProducerCallback<Self::Item>,
{
callback.callback(ParallelProducerWrapper(self.0))
}
fn len(&self) -> usize {
ExactSizeIterator::len(&self.0)
}
fn drive<C>(self, consumer: C) -> C::Result
where
C: Consumer<Self::Item>,
{
bridge(self, consumer)
}
}
impl<'a, T> IntoIterator for ParallelProducerWrapper<$csr_matrix_row_iter<'a, T>> {
type IntoIter = $csr_matrix_row_iter<'a, T>;
type Item = <Self::IntoIter as Iterator>::Item;
fn into_iter(self) -> Self::IntoIter {
self.0
}
}
impl<'a, T> Producer for ParallelProducerWrapper<$csr_matrix_row_iter<'a, T>> {
type IntoIter = $csr_matrix_row_iter<'a, T>;
type Item = <Self::IntoIter as Iterator>::Item;
fn into_iter(self) -> Self::IntoIter {
self.0
}
fn split_at(self, i: usize) -> (Self, Self) {
let (a, b) = self.0.split_at(i);
(ParallelProducerWrapper(a), ParallelProducerWrapper(b))
}
}
};
}
impl_par_iter!(CsrMatrixRowIter, CssSlice);
impl_par_iter!(CsrMatrixRowIterMut, CssSliceMut);