1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
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);