1use ark_ff::Zero;
2use rayon::iter::{FromParallelIterator, IntoParallelIterator, ParallelIterator};
3use std::ops::Index;
4
5#[derive(Clone, Debug, PartialEq, Eq, Hash)]
14pub struct Witness<const N_WIT: usize, T> {
15 pub cols: Box<[T; N_WIT]>,
18}
19
20impl<const N_WIT: usize, T: Zero + Clone> Default for Witness<N_WIT, T> {
21 fn default() -> Self {
22 Witness {
23 cols: Box::new(std::array::from_fn(|_| T::zero())),
24 }
25 }
26}
27
28impl<const N_WIT: usize, T> TryFrom<Vec<T>> for Witness<N_WIT, T> {
29 type Error = String;
30
31 fn try_from(value: Vec<T>) -> Result<Self, Self::Error> {
32 let len = value.len();
33 let cols: Box<[T; N_WIT]> = value
34 .try_into()
35 .map_err(|_| format!("Size mismatch: Expected {N_WIT:?} got {len:?}"))?;
36 Ok(Witness { cols })
37 }
38}
39
40impl<const N_WIT: usize, T> Index<usize> for Witness<N_WIT, T> {
41 type Output = T;
42
43 fn index(&self, index: usize) -> &Self::Output {
44 &self.cols[index]
45 }
46}
47
48impl<const N_WIT: usize, T> Witness<N_WIT, T> {
49 pub fn len(&self) -> usize {
50 self.cols.len()
51 }
52
53 pub fn is_empty(&self) -> bool {
54 self.cols.is_empty()
55 }
56}
57
58impl<const N_WIT: usize, T: Zero + Clone> Witness<N_WIT, Vec<T>> {
59 pub fn zero_vec(domain_size: usize) -> Self {
60 Witness {
61 cols: Box::new(std::array::from_fn(|_| vec![T::zero(); domain_size])),
64 }
65 }
66
67 pub fn to_pub_columns<const NPUB: usize>(&self) -> Witness<NPUB, Vec<T>> {
68 let mut newcols: [Vec<T>; NPUB] = std::array::from_fn(|_| vec![]);
69 for (i, vec) in self.cols[0..NPUB].iter().enumerate() {
70 newcols[i].clone_from(vec);
71 }
72 Witness {
73 cols: Box::new(newcols),
74 }
75 }
76}
77
78impl<'lt, const N_WIT: usize, G> IntoIterator for &'lt Witness<N_WIT, G> {
81 type Item = &'lt G;
82 type IntoIter = std::vec::IntoIter<&'lt G>;
83
84 fn into_iter(self) -> Self::IntoIter {
85 let mut iter_contents = Vec::with_capacity(N_WIT);
86 iter_contents.extend(&*self.cols);
87 iter_contents.into_iter()
88 }
89}
90
91impl<const N_WIT: usize, F: Clone> IntoIterator for Witness<N_WIT, F> {
92 type Item = F;
93 type IntoIter = std::vec::IntoIter<F>;
94
95 fn into_iter(self) -> Self::IntoIter {
97 let mut iter_contents = Vec::with_capacity(N_WIT);
98 iter_contents.extend(*self.cols);
99 iter_contents.into_iter()
100 }
101}
102
103impl<const N_WIT: usize, G> IntoParallelIterator for Witness<N_WIT, G>
104where
105 Vec<G>: IntoParallelIterator,
106{
107 type Iter = <Vec<G> as IntoParallelIterator>::Iter;
108 type Item = <Vec<G> as IntoParallelIterator>::Item;
109
110 fn into_par_iter(self) -> Self::Iter {
112 let mut iter_contents = Vec::with_capacity(N_WIT);
113 iter_contents.extend(*self.cols);
114 iter_contents.into_par_iter()
115 }
116}
117
118impl<const N_WIT: usize, G: Send + std::fmt::Debug> FromParallelIterator<G> for Witness<N_WIT, G> {
119 fn from_par_iter<I>(par_iter: I) -> Self
120 where
121 I: IntoParallelIterator<Item = G>,
122 {
123 let mut iter_contents = par_iter.into_par_iter().collect::<Vec<_>>();
124 let cols = iter_contents
125 .drain(..N_WIT)
126 .collect::<Vec<G>>()
127 .try_into()
128 .unwrap();
129 Witness { cols }
130 }
131}
132
133impl<'data, const N_WIT: usize, G> IntoParallelIterator for &'data Witness<N_WIT, G>
134where
135 Vec<&'data G>: IntoParallelIterator,
136{
137 type Iter = <Vec<&'data G> as IntoParallelIterator>::Iter;
138 type Item = <Vec<&'data G> as IntoParallelIterator>::Item;
139
140 fn into_par_iter(self) -> Self::Iter {
141 let mut iter_contents = Vec::with_capacity(N_WIT);
142 iter_contents.extend(&*self.cols);
143 iter_contents.into_par_iter()
144 }
145}
146
147impl<'data, const N_WIT: usize, G> IntoParallelIterator for &'data mut Witness<N_WIT, G>
148where
149 Vec<&'data mut G>: IntoParallelIterator,
150{
151 type Iter = <Vec<&'data mut G> as IntoParallelIterator>::Iter;
152 type Item = <Vec<&'data mut G> as IntoParallelIterator>::Item;
153
154 fn into_par_iter(self) -> Self::Iter {
155 let mut iter_contents = Vec::with_capacity(N_WIT);
156 iter_contents.extend(&mut *self.cols);
157 iter_contents.into_par_iter()
158 }
159}