zinc_poly/mle/dense/
try_collect_dense_mle.rs

1use itertools::Itertools;
2#[cfg(feature = "parallel")]
3use rayon::iter::ParallelIterator;
4use zinc_utils::log2;
5
6use crate::mle::DenseMultilinearExtension;
7
8pub trait TryCollectDenseMle: Iterator {
9    fn try_collect_dense_mle(self) -> Option<DenseMultilinearExtension<Self::Item>>;
10}
11
12impl<T> TryCollectDenseMle for T
13where
14    T: Iterator,
15{
16    fn try_collect_dense_mle(self) -> Option<DenseMultilinearExtension<Self::Item>> {
17        let evaluations = self.collect_vec();
18        let len = evaluations.len();
19
20        if len.is_power_of_two() {
21            Some(DenseMultilinearExtension {
22                evaluations,
23                num_vars: log2(len) as usize,
24            })
25        } else {
26            None
27        }
28    }
29}
30
31#[cfg(feature = "parallel")]
32pub trait ParallelTryCollectDenseMle: ParallelIterator {
33    fn try_collect_dense_mle(self) -> Option<DenseMultilinearExtension<Self::Item>>;
34}
35
36#[cfg(feature = "parallel")]
37impl<T> ParallelTryCollectDenseMle for T
38where
39    T: ParallelIterator,
40{
41    fn try_collect_dense_mle(self) -> Option<DenseMultilinearExtension<Self::Item>> {
42        let evaluations = self.collect::<Vec<_>>();
43        let len = evaluations.len();
44
45        if len.is_power_of_two() {
46            Some(DenseMultilinearExtension {
47                evaluations,
48                num_vars: log2(len) as usize,
49            })
50        } else {
51            None
52        }
53    }
54}
zinc_poly/mle/dense/try_collect_dense_mle.rs

zinc_poly/mle/dense/
try_collect_dense_mle.rs
