jf_pcs/
toeplitz.rs

1// Copyright (c) 2023 Espresso Systems (espressosys.com)
2// This file is part of the Jellyfish library.
3
4// You should have received a copy of the MIT License
5// along with the Jellyfish library. If not, see <https://mit-license.org/>.
6
7//! Toeplitz matrices and Circulant matrices operations
8//! References: `https://eprint.iacr.org/2020/1516.pdf`
9
10use ark_ff::FftField;
11use ark_poly::{domain::DomainCoeff, EvaluationDomain, GeneralEvaluationDomain};
12use ark_std::{format, ops::Mul, string::ToString, vec::Vec};
13use jf_utils::hadamard_product;
14
15use crate::prelude::PCSError;
16
17/// An `NxN` [Circulant Matrix](https://en.wikipedia.org/wiki/Circulant_matrix)
18/// is unambiguously represented by its first column, and has the form:
19///
20/// [c_0   c_N-1 c_N-2 ... c_1]
21/// [c_1   c_0   c_N-1 ... c_2]
22/// [c_2   c_1   c_0   ... c_3]
23/// [..    ..    ..    ... .. ]
24/// [c_N-1 c_N-2 c_N-3 ... c_0]
25///
26/// It is a special case of [`ToeplitzMatrix`].
27#[derive(Debug, Clone, PartialEq, Eq)]
28pub struct CirculantMatrix<F: FftField> {
29    col: Vec<F>,
30}
31
32impl<F: FftField> CirculantMatrix<F> {
33    /// Construct a Circulant matrix by its first column vector
34    #[allow(dead_code)]
35    pub fn new(col: Vec<F>) -> Self {
36        Self { col }
37    }
38
39    /// Fast multiplication of a Circulant matrix by a vector via FFT
40    /// Details see Section 2.2.1 of [Tomescu20](https://eprint.iacr.org/2020/1516.pdf).
41    // TODO: (alex) think of ways to extend to arbitrary length vector (simply
42    // truncate doesn't work).
43    pub fn fast_vec_mul<T>(&self, m: &[T]) -> Result<Vec<T>, PCSError>
44    where
45        T: for<'a> Mul<&'a F, Output = T> + DomainCoeff<F>,
46    {
47        if !m.len().is_power_of_two() {
48            return Err(PCSError::InvalidParameters(
49                "Fast Circulant Matrix mul only supports vector size of power of two.".to_string(),
50            ));
51        }
52        if m.len() != self.col.len() {
53            return Err(PCSError::InvalidParameters(
54                "Wrong input dimension for matrix mul.".to_string(),
55            ));
56        }
57        let domain: GeneralEvaluationDomain<F> =
58            GeneralEvaluationDomain::new(m.len()).expect("Should init an evaluation domain");
59
60        let m_evals = domain.fft(m); // DFT(m)
61        let col_evals = domain.fft(&self.col); // DFT(c_N)
62        let eval_prod = // DFT(c_N) * DFT(m)
63            hadamard_product(col_evals, m_evals).expect("Hadamard product should succeed");
64        let res = domain.ifft(&eval_prod); // iDFT(DFT(c_N) * DFT(m))
65        Ok(res)
66    }
67}
68
69/// An `NxN` [Toeplitz Matrix](https://en.wikipedia.org/wiki/Toeplitz_matrix) is
70/// unambiguously represented by its first column and first row, and has the
71/// form:
72///
73/// [a_0   a_-1  a_-2  ...  a_-(N-1)]
74/// [a_1   a_0   a_-1  ...  a_-(N-2)]
75/// [a_2   a_1   a_0   ...  a_-(N-3)]
76/// [..    ..    ..    ...    ..    ]
77/// [a_N-1 a_N-2 ..    ...  a_0     ]
78#[derive(Debug, Clone, PartialEq, Eq)]
79pub struct ToeplitzMatrix<F: FftField> {
80    col: Vec<F>,
81    row: Vec<F>,
82}
83
84impl<F: FftField> ToeplitzMatrix<F> {
85    /// constructor for a new Toeplitz matrix.
86    pub fn new(col: Vec<F>, row: Vec<F>) -> Result<Self, PCSError> {
87        if col.is_empty() || col.len() != row.len() {
88            return Err(PCSError::InvalidParameters(format!(
89                "row: {}, col: {} should be both positive and equal",
90                row.len(),
91                col.len()
92            )));
93        }
94        if col[0] != row[0] {
95            return Err(PCSError::InvalidParameters(format!(
96                "1st value in 1st col: {:?} should be the same as that in 1st row: {:?}",
97                col[0], row[0]
98            )));
99        }
100
101        Ok(Self { col, row })
102    }
103
104    /// Embeds a Toeplitz matrix of size N to a Circulant matrix of size 2N.
105    ///
106    /// Details see Section 2.3.1 of [Tomescu20](https://eprint.iacr.org/2020/1516.pdf).
107    // TODO: (alex) turn this into concurrent code after: https://github.com/EspressoSystems/jellyfish/issues/111
108    pub fn circulant_embedding(&self) -> Result<CirculantMatrix<F>, PCSError> {
109        let mut extension_col = self.row.clone();
110        extension_col.rotate_left(1);
111        extension_col.reverse();
112
113        Ok(CirculantMatrix {
114            col: [self.col.clone(), extension_col].concat(),
115        })
116    }
117
118    /// Fast multiplication of a Toeplitz matrix by embedding it into a
119    /// circulant matrix and multiply there.
120    ///
121    /// Details see Section 2.3.1 of [Tomescu20](https://eprint.iacr.org/2020/1516.pdf).
122    pub fn fast_vec_mul<T>(&self, v: &[T]) -> Result<Vec<T>, PCSError>
123    where
124        T: for<'a> Mul<&'a F, Output = T> + DomainCoeff<F>,
125    {
126        if !v.len().is_power_of_two() {
127            return Err(PCSError::InvalidParameters(
128                "Fast Toeplitz Matrix mul only supports vector size of power of two.".to_string(),
129            ));
130        }
131        if v.len() != self.col.len() {
132            return Err(PCSError::InvalidParameters(
133                "Wrong input dimension for matrix mul.".to_string(),
134            ));
135        }
136
137        let cir_repr = self.circulant_embedding()?;
138        let mut padded_v = Vec::from(v);
139        padded_v.resize(2 * v.len(), T::zero());
140
141        let mut res = cir_repr.fast_vec_mul(&padded_v)?;
142        res.truncate(v.len());
143        Ok(res)
144    }
145}
146
147impl<F: FftField> From<CirculantMatrix<F>> for ToeplitzMatrix<F> {
148    fn from(c: CirculantMatrix<F>) -> Self {
149        let mut row = c.col.clone();
150        row.rotate_left(1);
151        row.reverse();
152
153        Self { col: c.col, row }
154    }
155}
156
157impl<F: FftField> TryFrom<ToeplitzMatrix<F>> for CirculantMatrix<F> {
158    type Error = PCSError;
159
160    fn try_from(t: ToeplitzMatrix<F>) -> Result<Self, Self::Error> {
161        let mut expected_col = t.row;
162        expected_col.reverse();
163        expected_col.rotate_right(1);
164        if expected_col != t.col {
165            return Err(PCSError::InvalidParameters(
166                "Not a Circulant Matrix".to_string(),
167            ));
168        }
169        Ok(Self { col: t.col })
170    }
171}
172
173#[cfg(test)]
174mod tests {
175    use super::*;
176    use ark_bls12_381::{Fr, G1Projective};
177    use ark_ff::Field;
178    use ark_std::{ops::AddAssign, vec, UniformRand};
179    use jf_utils::test_rng;
180
181    // a MxN matrix, M rows, N cols.
182    struct Matrix<T, const M: usize, const N: usize>([[T; N]; M]);
183    struct Vector<T, const N: usize>([T; N]);
184
185    impl<T: Copy, const M: usize, const N: usize> Matrix<T, M, N> {
186        fn transpose(self) -> Matrix<T, N, M> {
187            let mut transposed = [[self.0[0][0]; M]; N];
188
189            #[allow(clippy::needless_range_loop)]
190            for i in 0..M {
191                for j in 0..N {
192                    transposed[j][i] = self.0[i][j];
193                }
194            }
195            Matrix(transposed)
196        }
197    }
198
199    impl<T: Copy, const N: usize> From<Matrix<T, 1, N>> for Vector<T, N> {
200        fn from(m: Matrix<T, 1, N>) -> Self {
201            let mut v = [m.0[0][0]; N];
202            v.copy_from_slice(&m.0[0]);
203            Vector(v)
204        }
205    }
206
207    impl<T: Copy, const N: usize> From<Vector<T, N>> for Vec<T> {
208        fn from(v: Vector<T, N>) -> Self {
209            Vec::from(&v.0[..])
210        }
211    }
212
213    impl<F: FftField> CirculantMatrix<F> {
214        fn full_matrix<const N: usize>(self) -> Matrix<F, N, N> {
215            assert_eq!(self.col.len(), N);
216            let t: ToeplitzMatrix<F> = self.into();
217
218            let mut first_row = [F::default(); N];
219            first_row.copy_from_slice(&t.row);
220            let mut row_vecs = [first_row; N];
221            let mut cur_row = first_row;
222
223            #[allow(clippy::needless_range_loop)]
224            for i in 1..N {
225                cur_row.rotate_right(1);
226                row_vecs[i] = cur_row;
227            }
228            // some arbitrary sanity check
229            assert_eq!(row_vecs[N - 1][0], row_vecs[0][1]);
230            assert_eq!(row_vecs[1][0], row_vecs[0][N - 1]);
231            assert_eq!(row_vecs[N - 1][N - 1], row_vecs[0][0]);
232
233            Matrix(row_vecs)
234        }
235    }
236
237    impl<F: FftField> ToeplitzMatrix<F> {
238        fn full_matrix<const N: usize>(self) -> Matrix<F, N, N> {
239            assert_eq!(self.col.len(), N);
240            let mut matrix = [[F::zero(); N]; N];
241            for i in 0..N {
242                matrix[i][0] = self.col[i];
243                matrix[0][i] = self.row[i];
244            }
245            for i in 1..N {
246                for j in 1..N {
247                    matrix[i][j] = matrix[i - 1][j - 1];
248                }
249            }
250            Matrix(matrix)
251        }
252    }
253
254    fn naive_matrix_mul<F, T, const M: usize, const N: usize, const K: usize>(
255        a: Matrix<F, M, N>,
256        b: Matrix<T, N, K>,
257    ) -> Matrix<T, M, K>
258    where
259        F: Field,
260        T: for<'a> Mul<&'a F, Output = T> + AddAssign<T> + Copy + Default,
261    {
262        let mut c = [[T::default(); K]; M];
263
264        #[allow(clippy::needless_range_loop)]
265        for i in 0..M {
266            for j in 0..K {
267                for k in 0..N {
268                    c[i][j] += b.0[k][j] * &a.0[i][k];
269                }
270            }
271        }
272        Matrix(c)
273    }
274
275    #[test]
276    fn test_circulant_mul() -> Result<(), PCSError> {
277        let mut rng = test_rng();
278        // happy path
279        const N: usize = 16;
280
281        let cir_matrix = CirculantMatrix::new((0..N).map(|_| Fr::rand(&mut rng)).collect());
282
283        let msgs = [G1Projective::rand(&mut rng); N];
284        let msg_matrix = Matrix([msgs]);
285
286        let expected: Vector<G1Projective, N> =
287            naive_matrix_mul(cir_matrix.clone().full_matrix(), msg_matrix.transpose())
288                .transpose()
289                .into();
290        let got = cir_matrix.fast_vec_mul(&msgs)?;
291        assert_eq!(
292            <Vector<G1Projective, N> as Into<Vec<G1Projective>>>::into(expected),
293            got,
294            "Fast Circulant Matrix mul for EC group is incorrect."
295        );
296
297        let f_msgs = [Fr::rand(&mut rng); N];
298        let f_msg_matrix = Matrix([f_msgs]);
299
300        let expected: Vector<Fr, N> =
301            naive_matrix_mul(cir_matrix.clone().full_matrix(), f_msg_matrix.transpose())
302                .transpose()
303                .into();
304        let got = cir_matrix.fast_vec_mul(&f_msgs)?;
305        assert_eq!(
306            <Vector<Fr, N> as Into<Vec<Fr>>>::into(expected),
307            got,
308            "Fast Circulant Matrix mul for field is incorrect."
309        );
310
311        // bad path
312        // mismatched matrix.col.len() and msgs.len() should fail
313        let bad_msg = [msgs.to_vec(), vec![G1Projective::rand(&mut rng)]].concat();
314        assert!(cir_matrix.fast_vec_mul(&bad_msg).is_err());
315
316        // non power-of-two matrix fast mul should fail
317        let m = bad_msg.len(); // same dimension as the message, but not a power-of-two
318        let cir_matrix = CirculantMatrix::new((0..m).map(|_| Fr::rand(&mut rng)).collect());
319
320        assert!(
321            !m.is_power_of_two()
322                && m == cir_matrix.col.len()
323                && cir_matrix.fast_vec_mul(&bad_msg).is_err()
324        );
325
326        Ok(())
327    }
328
329    #[test]
330    fn test_toeplitz_mul() -> Result<(), PCSError> {
331        let mut rng = test_rng();
332        const N: usize = 16;
333
334        let rand_col: Vec<Fr> = (0..N).map(|_| Fr::rand(&mut rng)).collect();
335        let rand_row = (0..N)
336            .map(|i| {
337                if i == 0 {
338                    rand_col[0]
339                } else {
340                    Fr::rand(&mut rng)
341                }
342            })
343            .collect();
344        let toep_matrix = ToeplitzMatrix::new(rand_col, rand_row)?;
345
346        let mut msgs = [G1Projective::default(); N];
347        for m in msgs.iter_mut() {
348            *m = G1Projective::rand(&mut rng);
349        }
350        let msg_matrix = Matrix([msgs]);
351
352        let expected: Vector<G1Projective, N> =
353            naive_matrix_mul(toep_matrix.clone().full_matrix(), msg_matrix.transpose())
354                .transpose()
355                .into();
356        let got = toep_matrix.fast_vec_mul(&msgs)?;
357        assert_eq!(
358            <Vector<G1Projective, N> as Into<Vec<G1Projective>>>::into(expected),
359            got,
360            "Fast Toeplitz Matrix mul is incorrect."
361        );
362
363        Ok(())
364    }
365}