shockwave-plus/tensor_pcs/src/tensor_pcs.rs
tensor_pcs/src/tensor_pcs.rsRaw
1use crate::rs_config::ecfft::ECFFTConfig;
2use crate::tree::BaseOpening;
3use crate::FieldExt;
4use ecfft::extend;
5use serde::{Deserialize, Serialize};
6
7use crate::fft::fft;
8use crate::polynomial::eq_poly::EqPoly;
9use crate::polynomial::sparse_ml_poly::SparseMLPoly;
10use crate::tensor_code::TensorCode;
11use crate::transcript::Transcript;
12use crate::utils::{dot_prod, hash_all, rlc_rows, sample_indices};
13
14use super::tensor_code::CommittedTensorCode;
15
16#[derive(Clone)]
17pub struct TensorRSMultilinearPCSConfig<F: FieldExt> {
18 pub expansion_factor: usize,
19 pub domain_powers: Option<Vec<Vec<F>>>,
20 pub fft_domain: Option<Vec<F>>,
21 pub ecfft_config: Option<ECFFTConfig<F>>,
22 pub l: usize,
23 pub num_entries: usize,
24 pub num_rows: usize,
25}
26
27impl<F: FieldExt> TensorRSMultilinearPCSConfig<F> {
28 pub fn num_cols(&self) -> usize {
29 self.num_entries / self.num_rows()
30 }
31
32 pub fn num_rows(&self) -> usize {
33 self.num_rows
34 }
35}
36
37#[derive(Clone)]
38pub struct TensorMultilinearPCS<F: FieldExt> {
39 config: TensorRSMultilinearPCSConfig<F>,
40}
41
42#[derive(Clone, Serialize, Deserialize)]
43pub struct TensorMLOpening<F: FieldExt> {
44 pub x: Vec<F>,
45 pub y: F,
46 pub base_opening: BaseOpening,
47 pub test_query_leaves: Vec<Vec<F>>,
48 pub eval_query_leaves: Vec<Vec<F>>,
49 pub u_hat_comm: [u8; 32],
50 pub test_u_prime: Vec<F>,
51 pub test_r_prime: Vec<F>,
52 pub eval_r_prime: Vec<F>,
53 pub eval_u_prime: Vec<F>,
54}
55
56impl<F: FieldExt> TensorMultilinearPCS<F> {
57 pub fn new(config: TensorRSMultilinearPCSConfig<F>) -> Self {
58 Self { config }
59 }
60
61 pub fn commit(&self, poly: &SparseMLPoly<F>) -> CommittedTensorCode<F> {
62 // Merkle commit to the evaluations of the polynomial
63 let tensor_code = self.encode_zk(poly);
64 let tree = tensor_code.commit(self.config.num_cols(), self.config.num_rows());
65 tree
66 }
67
68 pub fn open(
69 &self,
70 u_hat_comm: &CommittedTensorCode<F>,
71 poly: &SparseMLPoly<F>,
72 point: &[F],
73 transcript: &mut Transcript<F>,
74 ) -> TensorMLOpening<F> {
75 let num_cols = self.config.num_cols();
76 let num_rows = self.config.num_rows();
77 debug_assert_eq!(poly.num_vars, point.len());
78
79 // ########################################
80 // Testing phase
81 // Prove the consistency between the random linear combination of the evaluation tensor (u_prime)
82 // and the tensor codeword (u_hat)
83 // ########################################
84
85 // Derive the challenge vector;
86 let r_u = transcript.challenge_vec(num_rows);
87
88 let u = (0..num_rows)
89 .map(|i| {
90 poly.evals[(i * num_cols)..((i + 1) * num_cols)]
91 .iter()
92 .map(|entry| entry.1)
93 .collect::<Vec<F>>()
94 })
95 .collect::<Vec<Vec<F>>>();
96
97 // Random linear combination of the rows of the polynomial in a tensor structure
98 let test_u_prime = rlc_rows(u.clone(), &r_u);
99
100 // Random linear combination of the blinder
101 let blinder = u_hat_comm
102 .tensor_codeword
103 .0
104 .iter()
105 .map(|row| row[(row.len() / 2)..].to_vec())
106 .collect::<Vec<Vec<F>>>();
107
108 debug_assert_eq!(blinder[0].len(), u_hat_comm.tensor_codeword.0[0].len() / 2);
109
110 let test_r_prime = rlc_rows(blinder.clone(), &r_u);
111
112 let num_indices = self.config.l;
113 let indices = sample_indices(num_indices, num_cols * 2, transcript);
114
115 let test_queries = self.test_phase(&indices, &u_hat_comm);
116
117 // ########################################
118 // Evaluation phase
119 // Prove the consistency
120 // ########################################
121
122 // Get the evaluation point
123 let mut point_rev = point.to_vec();
124 point_rev.reverse();
125
126 let log2_num_rows = (num_rows as f64).log2() as usize;
127 let q1 = EqPoly::new(point_rev[0..log2_num_rows].to_vec()).evals();
128
129 let eval_r_prime = rlc_rows(blinder, &q1);
130
131 let eval_u_prime = rlc_rows(u.clone(), &q1);
132
133 let eval_queries = self.test_phase(&indices, &u_hat_comm);
134
135 TensorMLOpening {
136 x: point.to_vec(),
137 y: poly.eval(&point_rev),
138 eval_query_leaves: eval_queries,
139 test_query_leaves: test_queries,
140 u_hat_comm: u_hat_comm.committed_tree.root(),
141 test_u_prime,
142 test_r_prime,
143 eval_r_prime,
144 eval_u_prime,
145 base_opening: BaseOpening {
146 hashes: u_hat_comm.committed_tree.column_roots.clone(),
147 },
148 }
149 }
150}
151
152impl<F: FieldExt> TensorMultilinearPCS<F> {
153 pub fn verify(&self, opening: &TensorMLOpening<F>, transcript: &mut Transcript<F>) {
154 let num_rows = self.config.num_rows();
155 let num_cols = self.config.num_cols();
156
157 // Verify the base opening
158 let base_opening = &opening.base_opening;
159 base_opening.verify(opening.u_hat_comm);
160
161 // ########################################
162 // Verify test phase
163 // ########################################
164
165 let r_u = transcript.challenge_vec(num_rows);
166 println!("r_u = {:?}", r_u);
167
168 let test_u_prime_rs_codeword = self
169 .rs_encode(&opening.test_u_prime)
170 .iter()
171 .zip(opening.test_r_prime.iter())
172 .map(|(c, r)| *c + *r)
173 .collect::<Vec<F>>();
174
175 let num_indices = self.config.l;
176 let indices = sample_indices(num_indices, num_cols * 2, transcript);
177
178 debug_assert_eq!(indices.len(), opening.test_query_leaves.len());
179 for (expected_index, leaves) in indices.iter().zip(opening.test_query_leaves.iter()) {
180 // Verify that the hashes of the leaves equals the corresponding column root
181 let leaf_bytes = leaves
182 .iter()
183 .map(|x| x.to_repr())
184 .collect::<Vec<[u8; 32]>>();
185 let column_root = hash_all(&leaf_bytes);
186 let expected_column_root = base_opening.hashes[*expected_index];
187 assert_eq!(column_root, expected_column_root);
188
189 let mut sum = F::ZERO;
190 for (leaf, r_i) in leaves.iter().zip(r_u.iter()) {
191 sum += *r_i * *leaf;
192 }
193 assert_eq!(sum, test_u_prime_rs_codeword[*expected_index]);
194 }
195
196 // ########################################
197 // Verify evaluation phase
198 // ########################################
199
200 let mut x_rev = opening.x.clone();
201 x_rev.reverse();
202
203 let log2_num_rows = (num_rows as f64).log2() as usize;
204 let q1 = EqPoly::new(x_rev[0..log2_num_rows].to_vec()).evals();
205 let q2 = EqPoly::new(x_rev[log2_num_rows..].to_vec()).evals();
206
207 let eval_u_prime_rs_codeword = self
208 .rs_encode(&opening.eval_u_prime)
209 .iter()
210 .zip(opening.eval_r_prime.iter())
211 .map(|(c, r)| *c + *r)
212 .collect::<Vec<F>>();
213
214 debug_assert_eq!(q1.len(), opening.eval_query_leaves[0].len());
215 debug_assert_eq!(indices.len(), opening.test_query_leaves.len());
216 for (expected_index, leaves) in indices.iter().zip(opening.eval_query_leaves.iter()) {
217 // TODO: Don't need to check the leaves again?
218 // Verify that the hashes of the leaves equals the corresponding column root
219 let leaf_bytes = leaves
220 .iter()
221 .map(|x| x.to_repr())
222 .collect::<Vec<[u8; 32]>>();
223 let column_root = hash_all(&leaf_bytes);
224 let expected_column_root = base_opening.hashes[*expected_index];
225 assert_eq!(column_root, expected_column_root);
226
227 let mut sum = F::ZERO;
228 for (leaf, q1_i) in leaves.iter().zip(q1.iter()) {
229 sum += *q1_i * *leaf;
230 }
231 assert_eq!(sum, eval_u_prime_rs_codeword[*expected_index]);
232 }
233
234 let expected_eval = dot_prod(&opening.eval_u_prime, &q2);
235 assert_eq!(expected_eval, opening.y);
236 }
237
238 fn split_encode(&self, message: &[F]) -> Vec<F> {
239 let codeword = self.rs_encode(message);
240
241 let mut rng = rand::thread_rng();
242 let blinder = (0..codeword.len())
243 .map(|_| F::random(&mut rng))
244 .collect::<Vec<F>>();
245
246 let mut randomized_codeword = codeword
247 .iter()
248 .zip(blinder.clone().iter())
249 .map(|(c, b)| *b + *c)
250 .collect::<Vec<F>>();
251
252 randomized_codeword.extend_from_slice(&blinder);
253 debug_assert_eq!(randomized_codeword.len(), codeword.len() * 2);
254 randomized_codeword
255 }
256
257 fn rs_encode(&self, message: &[F]) -> Vec<F> {
258 let codeword = if self.config.fft_domain.is_some() {
259 let fft_domain = self.config.fft_domain.as_ref().unwrap();
260 let mut padded_coeffs = message.clone().to_vec();
261 padded_coeffs.resize(fft_domain.len(), F::ZERO);
262 let codeword = fft(&padded_coeffs, &fft_domain);
263
264 codeword
265 } else if self.config.ecfft_config.is_some() {
266 let ecfft_config = self.config.ecfft_config.as_ref().unwrap();
267 assert_eq!(
268 message.len() * self.config.expansion_factor,
269 ecfft_config.domain[0].len()
270 );
271 let extended_evals = extend(
272 message,
273 &ecfft_config.domain,
274 &ecfft_config.matrices,
275 &ecfft_config.inverse_matrices,
276 0,
277 );
278
279 let codeword = [message.to_vec(), extended_evals].concat();
280 codeword
281 } else {
282 let domain_powers = self.config.domain_powers.as_ref().unwrap();
283 assert_eq!(message.len(), domain_powers[0].len());
284 assert_eq!(
285 message.len() * self.config.expansion_factor,
286 domain_powers.len()
287 );
288
289 let codeword = domain_powers
290 .iter()
291 .map(|powers| {
292 message
293 .iter()
294 .zip(powers.iter())
295 .fold(F::ZERO, |acc, (m, p)| acc + *m * *p)
296 })
297 .collect::<Vec<F>>();
298
299 codeword
300 };
301
302 codeword
303 }
304
305 fn test_phase(&self, indices: &[usize], u_hat_comm: &CommittedTensorCode<F>) -> Vec<Vec<F>> {
306 let num_cols = self.config.num_cols() * 2;
307
308 // Query the columns of u_hat
309 let num_indices = self.config.l;
310
311 let u_hat_openings = indices
312 .iter()
313 .map(|index| u_hat_comm.query_column(*index, num_cols))
314 .collect::<Vec<Vec<F>>>();
315
316 debug_assert_eq!(u_hat_openings.len(), num_indices);
317
318 u_hat_openings
319 }
320
321 fn encode_zk(&self, poly: &SparseMLPoly<F>) -> TensorCode<F> {
322 let num_rows = self.config.num_rows();
323 let num_cols = self.config.num_cols();
324
325 let codewords = (0..num_rows)
326 .map(|i| {
327 poly.evals[i * num_cols..(i + 1) * num_cols]
328 .iter()
329 .map(|entry| entry.1)
330 .collect::<Vec<F>>()
331 })
332 .map(|row| self.split_encode(&row))
333 .collect::<Vec<Vec<F>>>();
334
335 TensorCode(codewords)
336 }
337}
338
339#[cfg(test)]
340mod tests {
341 use super::*;
342 use crate::rs_config::{ecfft, naive, smooth};
343
344 const TEST_NUM_VARS: usize = 10;
345 const TEST_L: usize = 10;
346
347 fn test_poly<F: FieldExt>() -> SparseMLPoly<F> {
348 let num_entries: usize = 2usize.pow(TEST_NUM_VARS as u32);
349
350 let evals = (0..num_entries)
351 .map(|i| (i, F::from(i as u64)))
352 .collect::<Vec<(usize, F)>>();
353
354 let ml_poly = SparseMLPoly::new(evals, TEST_NUM_VARS);
355 ml_poly
356 }
357
358 fn prove_and_verify<F: FieldExt>(ml_poly: SparseMLPoly<F>, pcs: TensorMultilinearPCS<F>) {
359 let comm = pcs.commit(&ml_poly);
360
361 let open_at = (0..ml_poly.num_vars)
362 .map(|i| F::from(i as u64))
363 .collect::<Vec<F>>();
364
365 let mut prover_transcript = Transcript::<F>::new(b"test");
366 prover_transcript.append_bytes(&comm.committed_tree.root);
367 let opening = pcs.open(&comm, &ml_poly, &open_at, &mut prover_transcript);
368
369 let mut verifier_transcript = Transcript::<F>::new(b"test");
370 verifier_transcript.append_bytes(&comm.committed_tree.root);
371 pcs.verify(&opening, &mut verifier_transcript);
372 }
373
374 fn config_base<F: FieldExt>(ml_poly: &SparseMLPoly<F>) -> TensorRSMultilinearPCSConfig<F> {
375 let num_vars = ml_poly.num_vars;
376 let num_evals = 2usize.pow(num_vars as u32);
377 let num_rows = 2usize.pow((num_vars / 2) as u32);
378
379 let expansion_factor = 2;
380
381 TensorRSMultilinearPCSConfig::<F> {
382 expansion_factor,
383 domain_powers: None,
384 fft_domain: None,
385 ecfft_config: None,
386 l: TEST_L,
387 num_entries: num_evals,
388 num_rows,
389 }
390 }
391
392 #[test]
393 fn test_tensor_pcs_fft() {
394 type F = halo2curves::pasta::Fp;
395 // FFT config
396 let ml_poly = test_poly();
397 let mut config = config_base(&ml_poly);
398 config.fft_domain = Some(smooth::gen_config(config.num_cols()));
399
400 // Test FFT PCS
401 let tensor_pcs_fft = TensorMultilinearPCS::<F>::new(config);
402 prove_and_verify(ml_poly, tensor_pcs_fft);
403 }
404
405 #[test]
406 fn test_tensor_pcs_ecfft() {
407 type F = halo2curves::secp256k1::Fp;
408 let ml_poly = test_poly();
409
410 let mut config = config_base(&ml_poly);
411 config.ecfft_config = Some(ecfft::gen_config(config.num_cols()));
412
413 // Test FFT PCS
414 let tensor_pcs_ecf = TensorMultilinearPCS::<F>::new(config);
415 prove_and_verify(ml_poly, tensor_pcs_ecf);
416 }
417
418 #[test]
419 fn test_tensor_pcs_naive() {
420 type F = halo2curves::secp256k1::Fp;
421 // FFT config
422 let ml_poly = test_poly();
423
424 // Naive config
425 let mut config = config_base(&ml_poly);
426 config.domain_powers = Some(naive::gen_config(config.num_cols()));
427
428 // Test FFT PCS
429 let tensor_pcs_naive = TensorMultilinearPCS::<F>::new(config);
430 prove_and_verify(ml_poly, tensor_pcs_naive);
431 }
432}