kimchi_msm/
witness.rs

1use ark_ff::Zero;
2use rayon::iter::{FromParallelIterator, IntoParallelIterator, ParallelIterator};
3use std::ops::Index;
4
5/// The witness columns used by a gate of the MSM circuits.
6/// It is generic over the number of columns, `N_WIT`, and the type of the
7/// witness, `T`.
8/// It is parametrized by a type `T` which can be either:
9/// - `Vec<G::ScalarField>` for the evaluations
10/// - `PolyComm<G>` for the commitments
11///
12/// It can be used to represent the different subcircuits used by the project.
13#[derive(Clone, Debug, PartialEq, Eq, Hash)]
14pub struct Witness<const N_WIT: usize, T> {
15    /// A witness row is represented by an array of N witness columns
16    /// When T is a vector, then the witness describes the rows of the circuit.
17    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            // Ideally the vector should be of domain size, but
62            // one-element vector should be a reasonable default too.
63            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
78// IMPLEMENTATION OF ITERATORS FOR THE WITNESS STRUCTURE
79
80impl<'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    /// Iterate over the columns in the circuit.
96    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    /// Iterate over the columns in the circuit, in parallel.
111    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}