plonky2-crypto-gadgets/src/circuit.rs
src/circuit.rsRaw
1use std::{array, sync::Arc};
2
3use plonky2::{
4 field::{extension::Extendable, goldilocks_field::GoldilocksField as F, types::Field},
5 hash::hash_types::RichField,
6 iop::{
7 target::{BoolTarget, Target},
8 witness::{PartialWitness, WitnessWrite},
9 },
10 plonk::circuit_builder::CircuitBuilder,
11};
12
13use crate::{
14 constants::{RCON, SBOX},
15 native::{rot_word, shift_rows, State},
16 D,
17};
18
19#[derive(Debug, Copy, Clone)]
20pub struct StateTarget(pub [[ByteArrayTarget; 4]; 4]);
21
22// TODO: Normalise terminology!
23pub type ByteArrayTarget = [BoolTarget; 8];
24
25pub trait CircuitBuilderAESState<F: RichField + Extendable<D>, const D: usize> {
26 /// Adds state target.
27 fn add_virtual_state(&mut self) -> StateTarget;
28
29 /// AES cipher as in spec.
30 fn encrypt_block<const NR: usize>(
31 &mut self,
32 sbox_lut_idx: usize,
33 mix_matrix: [[ByteArrayTarget; 4]; 4],
34 s: StateTarget,
35
36 w: [[ByteArrayTarget; 4]; 4 * (NR + 1)],
37 ) -> StateTarget;
38
39 /// Applies sub_bytes routine to state.
40 fn state_sub_bytes(&mut self, sbox_lut_idx: usize, s: StateTarget) -> StateTarget;
41
42 /// SubWord
43 fn state_sub_word(
44 &mut self,
45 sbox_lut_idx: usize,
46 word: [ByteArrayTarget; 4],
47 ) -> [ByteArrayTarget; 4];
48
49 /// MixColumns
50 fn state_mix_columns(
51 &mut self,
52 mix_matrix: [[ByteArrayTarget; 4]; 4],
53 s: StateTarget,
54 ) -> StateTarget;
55
56 /// AddRoundKey
57 fn state_add_round_key(
58 &mut self,
59 round_key: &[[ByteArrayTarget; 4]],
60 s: StateTarget,
61 ) -> StateTarget;
62
63 /// KeyExpansion
64 fn key_expansion<const NK: usize, const NB: usize, const NR: usize>(
65 &mut self,
66 sbox_lut_idx: usize,
67 key: [ByteArrayTarget; NK * NB],
68 ) -> [[ByteArrayTarget; 4]; 4 * (NR + 1)];
69
70 /// GF(2^8) addition
71 fn gf_2_8_add(&mut self, x: Target, y: Target) -> Target;
72
73 /// GF(2^8) addition actinb on bit arrays
74 fn gf_2_8_add_bits(
75 &mut self,
76 x_bits: [BoolTarget; 8],
77 y_bits: [BoolTarget; 8],
78 ) -> [BoolTarget; 8];
79
80 /// GF(2^8) multiplication
81 fn gf_2_8_mul(&mut self, x: Target, y: Target) -> Target;
82
83 /// GF(2^8) multiplication acting on bit arrays
84 fn gf_2_8_mul_bits(&mut self, x: [BoolTarget; 8], y: [BoolTarget; 8]) -> [BoolTarget; 8];
85
86 /// Bytearray inner product.
87 fn bytearray_ip_bits<const N: usize>(
88 &mut self,
89 x: [[BoolTarget; 8]; N],
90 y: [[BoolTarget; 8]; N],
91 ) -> [BoolTarget; 8];
92
93 /// Bytearray matrix application.
94 fn bytearray_matrix_apply_bits<const M: usize, const N: usize>(
95 &mut self,
96 a: [[ByteArrayTarget; N]; M],
97 x: [ByteArrayTarget; N],
98 ) -> [ByteArrayTarget; M];
99}
100
101impl CircuitBuilderAESState<F, D> for CircuitBuilder<F, D> {
102 fn add_virtual_state(&mut self) -> StateTarget {
103 StateTarget(array::from_fn(|_| {
104 array::from_fn(|_| array::from_fn(|_| self.add_virtual_bool_target_safe()))
105 }))
106 }
107
108 fn encrypt_block<const NR: usize>(
109 &mut self,
110 sbox_lut_idx: usize,
111 mix_matrix: [[ByteArrayTarget; 4]; 4], // constant mix_matrix
112 s: StateTarget,
113 w: [[ByteArrayTarget; 4]; 4 * (NR + 1)], // expanded key
114 ) -> StateTarget {
115 let mut s = s;
116 s = self.state_add_round_key(&w[0..4], s);
117 (1..NR).for_each(|i| {
118 s = self.state_sub_bytes(sbox_lut_idx, s);
119 s = StateTarget(shift_rows(s.0));
120 s = self.state_mix_columns(mix_matrix, s);
121 s = self.state_add_round_key(&w[4 * i..4 * (i + 1)], s);
122 });
123 s = self.state_sub_bytes(sbox_lut_idx, s);
124 s = StateTarget(shift_rows(s.0));
125 self.state_add_round_key(&w[4 * NR..4 * (NR + 1)], s)
126 }
127
128 fn state_sub_bytes(&mut self, sbox_lut_idx: usize, s: StateTarget) -> StateTarget {
129 StateTarget(array::from_fn(|i| {
130 self.state_sub_word(sbox_lut_idx, s.0[i])
131 }))
132 }
133
134 fn state_sub_word(
135 &mut self,
136 sbox_lut_idx: usize,
137 word: [ByteArrayTarget; 4],
138 ) -> [ByteArrayTarget; 4] {
139 array::from_fn(|i| {
140 let byte_target = target_from_bitarray(self, &word[i]);
141 let out_target = self.add_lookup_from_index(byte_target, sbox_lut_idx);
142 bitarray_from_bytetarget(self, out_target)
143 })
144 }
145
146 fn state_mix_columns(
147 &mut self,
148 mix_matrix: [[ByteArrayTarget; 4]; 4],
149 s: StateTarget,
150 ) -> StateTarget {
151 let cols: [_; 4] = array::from_fn(|i| array::from_fn(|j| s.0[j][i]));
152 let out_cols: [_; 4] =
153 array::from_fn(|i| self.bytearray_matrix_apply_bits(mix_matrix, cols[i]));
154 StateTarget(array::from_fn(|i| array::from_fn(|j| out_cols[j][i])))
155 }
156
157 fn state_add_round_key(
158 &mut self,
159 round_key: &[[ByteArrayTarget; 4]],
160 s: StateTarget,
161 ) -> StateTarget {
162 StateTarget(array::from_fn(|i| {
163 array::from_fn(|j| self.gf_2_8_add_bits(s.0[i][j], round_key[j][i]))
164 }))
165 }
166
167 fn key_expansion<const NK: usize, const NB: usize, const NR: usize>(
168 &mut self,
169 sbox_lut_idx: usize,
170 key: [ByteArrayTarget; NK * NB],
171 ) -> [[ByteArrayTarget; 4]; 4 * (NR + 1)] {
172 let rcon: [ByteArrayTarget; 11] = array::from_fn(|i| {
173 let rcon_bits = le_bits_from_byte(RCON[i]);
174 array::from_fn(|j| self.constant_bool(rcon_bits[j]))
175 });
176
177 let key: [[_; 4]; NK] = array::from_fn(|i| array::from_fn(|j| key[4 * i + j]));
178
179 let key_additions = (NK..4 * (NR + 1))
180 .scan(key.to_vec(), |st, i| {
181 let offset = if i % NK == 0 {
182 let term = self.state_sub_word(sbox_lut_idx, rot_word(st[i - 1]));
183 array::from_fn(|j| {
184 if j == 0 {
185 self.gf_2_8_add_bits(term[j], rcon[i / NK])
186 } else {
187 term[j]
188 }
189 })
190 } else if (NK > 6) && (i % NK == 4) {
191 self.state_sub_word(sbox_lut_idx, st[i - 1])
192 } else {
193 st[i - 1]
194 };
195
196 let cur_val = array::from_fn(|j| self.gf_2_8_add_bits(st[i - NK][j], offset[j]));
197 st.push(cur_val);
198 Some(cur_val)
199 })
200 .collect::<Vec<_>>();
201 std::array::from_fn(|i| {
202 if i < NK {
203 key[i]
204 } else {
205 key_additions[i - NK]
206 }
207 })
208 }
209
210 fn gf_2_8_add(&mut self, x: Target, y: Target) -> Target {
211 let x_bits = bitarray_from_bytetarget(self, x);
212 let y_bits = bitarray_from_bytetarget(self, y);
213
214 let sum_bits = self.gf_2_8_add_bits(x_bits, y_bits);
215 target_from_bitarray(self, &sum_bits)
216 }
217
218 fn gf_2_8_add_bits(
219 &mut self,
220 x_bits: [BoolTarget; 8],
221 y_bits: [BoolTarget; 8],
222 ) -> [BoolTarget; 8] {
223 let sum_bits = std::iter::zip(x_bits, y_bits)
224 .map(|(a, b)| xor(self, a, b))
225 .collect::<Vec<_>>();
226
227 array::from_fn(|i| sum_bits[i])
228 }
229
230 fn gf_2_8_mul(&mut self, x: Target, y: Target) -> Target {
231 let x_bits = bitarray_from_bytetarget(self, x);
232 let y_bits = bitarray_from_bytetarget(self, y);
233
234 let prod_bits = self.gf_2_8_mul_bits(x_bits, y_bits);
235 target_from_bitarray(self, &prod_bits)
236 }
237
238 fn gf_2_8_mul_bits(
239 &mut self,
240 x_bits: [BoolTarget; 8],
241 y_bits: [BoolTarget; 8],
242 ) -> [BoolTarget; 8] {
243 let zero_bits = vec![self._false(); 8];
244 assert!(x_bits.len() == 8 && y_bits.len() == 8);
245
246 let x_times = |builder: &mut CircuitBuilder<_, _>, y: &[BoolTarget]| -> [_; 8] {
247 let zero = builder._false();
248 let one = builder._true();
249
250 let potential_offset = [one, one, zero, one, one, zero, zero, zero];
251 let y_shifted: [_; 8] = array::from_fn(|i| if i == 0 { zero } else { y[i - 1] });
252
253 let offset: [_; 8] = array::from_fn(|i| builder.and(y[7], potential_offset[i]));
254 array::from_fn(|i| xor(builder, y_shifted[i], offset[i]))
255 };
256
257 let powers_times_x = (0..8)
258 .scan(x_bits.to_vec(), |st, i| {
259 if i > 0 {
260 *st = x_times(self, st).to_vec();
261 }
262 Some(st.clone())
263 })
264 .collect::<Vec<_>>();
265
266 let prod_bits = std::iter::zip(y_bits, powers_times_x)
267 .map(|(c, o)| o.into_iter().map(|b| self.and(c, b)).collect::<Vec<_>>())
268 .collect::<Vec<_>>()
269 .into_iter()
270 .fold(zero_bits, |acc, term| {
271 std::iter::zip(acc, term)
272 .map(|(a, b)| xor(self, a, b))
273 .collect::<Vec<_>>()
274 });
275
276 array::from_fn(|i| prod_bits[i])
277 }
278
279 fn bytearray_ip_bits<const N: usize>(
280 &mut self,
281 x: [[BoolTarget; 8]; N],
282 y: [[BoolTarget; 8]; N],
283 ) -> [BoolTarget; 8] {
284 let zero = self.zero();
285 let zero_bits = bitarray_from_bytetarget(self, zero);
286
287 std::iter::zip(x, y).fold(zero_bits, |acc, (a, b)| {
288 let prod = self.gf_2_8_mul_bits(a, b);
289 self.gf_2_8_add_bits(acc, prod)
290 })
291 }
292
293 fn bytearray_matrix_apply_bits<const M: usize, const N: usize>(
294 &mut self,
295 a: [[ByteArrayTarget; N]; M],
296 x: [ByteArrayTarget; N],
297 ) -> [ByteArrayTarget; M] {
298 std::array::from_fn(|i| self.bytearray_ip_bits(a[i], x))
299 }
300}
301
302pub trait PartialWitnessByteArray {
303 fn set_byte_array_target(&mut self, target: ByteArrayTarget, value: u8) -> anyhow::Result<()>;
304}
305
306impl<F: Field> PartialWitnessByteArray for PartialWitness<F> {
307 fn set_byte_array_target(&mut self, target: ByteArrayTarget, value: u8) -> anyhow::Result<()> {
308 let value_bits = le_bits_from_byte(value);
309 std::iter::zip(target, value_bits).try_for_each(|(t, v)| self.set_bool_target(t, v))
310 }
311}
312
313pub trait PartialWitnessAESState {
314 fn set_target_state(&mut self, target: StateTarget, value: State) -> anyhow::Result<()>;
315}
316
317impl<F: Field> PartialWitnessAESState for PartialWitness<F> {
318 fn set_target_state(&mut self, target: StateTarget, value: State) -> anyhow::Result<()> {
319 std::iter::zip(target.0, value).try_for_each(|(t, v)| {
320 std::iter::zip(t, v).try_for_each(|(t, v)| self.set_byte_array_target(t, v))
321 })
322 }
323}
324
325pub fn sbox_lut(builder: &mut CircuitBuilder<F, D>) -> usize {
326 builder.add_lookup_table_from_pairs(Arc::new(
327 SBOX.into_iter()
328 .enumerate()
329 .map(|(i, o)| (i as u16, o as u16))
330 .collect::<Vec<_>>(),
331 ))
332}
333
334pub fn state_mix_matrix_bits(builder: &mut CircuitBuilder<F, D>) -> [[[BoolTarget; 8]; 4]; 4] {
335 let one = builder.one();
336 let two = builder.two();
337 let three = builder.constant(F::from_canonical_u64(3));
338
339 let matrix = [
340 [two, three, one, one],
341 [one, two, three, one],
342 [one, one, two, three],
343 [three, one, one, two],
344 ];
345
346 array::from_fn(|i| array::from_fn(|j| bitarray_from_bytetarget(builder, matrix[i][j])))
347}
348
349pub fn xor(builder: &mut CircuitBuilder<F, D>, x: BoolTarget, y: BoolTarget) -> BoolTarget {
350 let x_or_y = builder.or(x, y);
351 BoolTarget::new_unsafe(builder.arithmetic(
352 F::NEG_ONE,
353 F::ONE,
354 x.target,
355 y.target,
356 x_or_y.target,
357 ))
358}
359
360// LE bit array conversions
361pub fn target_from_bitarray(builder: &mut CircuitBuilder<F, D>, x: &[BoolTarget]) -> Target {
362 let zero = builder.zero();
363 let two = builder.two();
364 x.iter()
365 .rev()
366 .fold(zero, |acc, b| builder.mul_add(two, acc, b.target))
367}
368
369pub fn bitarray_from_bytetarget(builder: &mut CircuitBuilder<F, D>, x: Target) -> [BoolTarget; 8] {
370 let x_bits = builder.split_le(x, 8);
371 array::from_fn(|i| x_bits[i])
372}
373
374pub fn le_bits_from_byte(v: u8) -> [bool; 8] {
375 let v_bits = (0..8)
376 .scan(v, |st, _| {
377 let cur_bit = (*st % 2) != 0;
378 *st >>= 1;
379 Some(cur_bit)
380 })
381 .collect::<Vec<_>>();
382 array::from_fn(|i| v_bits[i])
383}
384
385#[cfg(test)]
386mod tests {
387 use std::array;
388
389 use anyhow::Result;
390 use plonky2::{
391 field::{goldilocks_field::GoldilocksField as F, types::Field},
392 iop::witness::{PartialWitness, WitnessWrite},
393 plonk::{
394 circuit_builder::CircuitBuilder, circuit_data::CircuitConfig,
395 config::PoseidonGoldilocksConfig,
396 },
397 };
398 use rand::RngExt;
399
400 use crate::{
401 circuit::{
402 sbox_lut, state_mix_matrix_bits, ByteArrayTarget, CircuitBuilderAESState,
403 PartialWitnessAESState, PartialWitnessByteArray,
404 },
405 native::{encrypt_block, key_expansion, mix_columns, sub_bytes, State},
406 };
407
408 use super::D;
409
410 #[test]
411 fn test_sub_bytes() -> Result<()> {
412 // Circuit declaration
413 let config = CircuitConfig::standard_recursion_config();
414 let mut builder = CircuitBuilder::<F, D>::new(config);
415
416 let state_target = builder.add_virtual_state();
417 let sbox_lut = sbox_lut(&mut builder);
418 let out_state_target = builder.state_sub_bytes(sbox_lut, state_target);
419
420 let data = builder.build::<PoseidonGoldilocksConfig>();
421
422 // set values to circuit
423 let mut rng = rand::rng();
424 let test_states: [State; 10] = array::from_fn(|_| array::from_fn(|_| rng.random()));
425
426 test_states.into_iter().try_for_each(|s| {
427 let expected_result = sub_bytes(s);
428
429 let mut pw = PartialWitness::<F>::new();
430 pw.set_target_state(state_target, s)?;
431 pw.set_target_state(out_state_target, expected_result)?;
432
433 let proof = data.prove(pw)?;
434 data.verify(proof)
435 })
436 }
437
438 #[test]
439 fn test_mix_columns() -> Result<()> {
440 // Circuit declaration
441 let config = CircuitConfig::standard_recursion_config();
442 let mut builder = CircuitBuilder::<F, D>::new(config);
443
444 let state_target = builder.add_virtual_state();
445 let mix_matrix = state_mix_matrix_bits(&mut builder);
446 let out_state_target = builder.state_mix_columns(mix_matrix, state_target);
447
448 let data = builder.build::<PoseidonGoldilocksConfig>();
449
450 // set values to circuit
451 let mut rng = rand::rng();
452 let test_states: [State; 10] = array::from_fn(|_| array::from_fn(|_| rng.random()));
453
454 test_states.into_iter().try_for_each(|s| {
455 let expected_result = mix_columns(s);
456
457 let mut pw = PartialWitness::<F>::new();
458 pw.set_target_state(state_target, s)?;
459 pw.set_target_state(out_state_target, expected_result)?;
460
461 let proof = data.prove(pw)?;
462 data.verify(proof)
463 })
464 }
465
466 #[test]
467 fn test_gf_2_8_mul() -> Result<()> {
468 // Circuit declaration
469 let config = CircuitConfig::standard_recursion_config();
470 let mut builder = CircuitBuilder::<F, D>::new(config);
471
472 let a = builder.add_virtual_target();
473 let b = builder.add_virtual_target();
474
475 let a_times_b = builder.gf_2_8_mul(a, b);
476
477 let data = builder.build::<PoseidonGoldilocksConfig>();
478
479 // set values to circuit
480 let test_values = [
481 (0x57, 0x01, 0x57),
482 (0x57, 0x02, 0xae),
483 (0x57, 0x04, 0x47),
484 (0x57, 0x08, 0x8e),
485 (0x57, 0x10, 0x07),
486 (0x57, 0x20, 0x0e),
487 (0x57, 0x40, 0x1c),
488 (0x57, 0x80, 0x38),
489 (0x57, 0x13, 0xfe),
490 ];
491 test_values
492 .into_iter()
493 .map(|(a, b, c)| {
494 (
495 F::from_canonical_u8(a),
496 F::from_canonical_u8(b),
497 F::from_canonical_u8(c),
498 )
499 })
500 .try_for_each(|(a_val, b_val, expected_result)| {
501 let mut pw = PartialWitness::<F>::new();
502 pw.set_target(a, a_val)?;
503 pw.set_target(b, b_val)?;
504 pw.set_target(a_times_b, expected_result)?;
505
506 let proof = data.prove(pw)?;
507 data.verify(proof)
508 })
509 }
510
511 #[test]
512 fn test_key_expansion() -> Result<()> {
513 // AES-128
514 let key = [
515 0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf,
516 0x4f, 0x3c,
517 ];
518 test_key_expansion_op::</*Nk,Nb,Nr*/ 4, 4, 10>(key)?;
519
520 // AES-196
521 let key = [
522 0x8e, 0x73, 0xb0, 0xf7, 0xda, 0x0e, 0x64, 0x52, 0xc8, 0x10, 0xf3, 0x2b, 0x80, 0x90,
523 0x79, 0xe5, 0x62, 0xf8, 0xea, 0xd2, 0x52, 0x2c, 0x6b, 0x7b,
524 ];
525 test_key_expansion_op::</*Nk,Nb,Nr*/ 6, 4, 12>(key)?;
526
527 // AES-256
528 let key = [
529 0x60, 0x3d, 0xeb, 0x10, 0x15, 0xca, 0x71, 0xbe, 0x2b, 0x73, 0xae, 0xf0, 0x85, 0x7d,
530 0x77, 0x81, 0x1f, 0x35, 0x2c, 0x07, 0x3b, 0x61, 0x08, 0xd7, 0x2d, 0x98, 0x10, 0xa3,
531 0x09, 0x14, 0xdf, 0xf4,
532 ];
533 test_key_expansion_op::</*Nk,Nb,Nr*/ 8, 4, 14>(key)?;
534
535 Ok(())
536 }
537
538 /// Analogous tests to `native::tests::test_key_expansion`.
539 fn test_key_expansion_op<const NK: usize, const NB: usize, const NR: usize>(
540 key: [u8; NK * NB],
541 ) -> Result<()>
542 where
543 [(); 4 * (NR + 1)]:,
544 [(); 4 * NK]:,
545 {
546 // Circuit declaration
547 let config = CircuitConfig::standard_recursion_config();
548 let mut builder = CircuitBuilder::<F, D>::new(config);
549
550 let key_target: [ByteArrayTarget; NK * NB] =
551 array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe()));
552 let sbox_lut = sbox_lut(&mut builder);
553 let expanded_key_target = builder.key_expansion::<NK, NB, NR>(sbox_lut, key_target);
554
555 println!(
556 "key_expansion(NK:{}, NB:{}, NR:{}) num_gates: {}",
557 NK,
558 NB,
559 NR,
560 builder.num_gates()
561 );
562 let data = builder.build::<PoseidonGoldilocksConfig>();
563
564 // set values to circuit
565 let w = key_expansion::<NK, NB, NR>(&key);
566
567 let mut pw = PartialWitness::<F>::new();
568
569 std::iter::zip(key_target, key).try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?;
570
571 std::iter::zip(expanded_key_target, w).try_for_each(|(t, v)| {
572 std::iter::zip(t, v).try_for_each(|(t, v)| pw.set_byte_array_target(t, v))
573 })?;
574
575 let proof = data.prove(pw)?;
576 data.verify(proof)
577 }
578
579 /// test against native version
580 #[test]
581 fn test_encrypt_block_test_vector() -> Result<()> {
582 // AES-128
583 let input_state: [u8; 16] = [
584 0x32, 0x43, 0xf6, 0xa8, 0x88, 0x5a, 0x30, 0x8d, 0x31, 0x31, 0x98, 0xa2, 0xe0, 0x37,
585 0x07, 0x34,
586 ];
587 let key = [
588 0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf,
589 0x4f, 0x3c,
590 ];
591 test_encrypt_block_test_vector_op::<4, 4, 10>(input_state, key)?;
592
593 // AES-196
594 let input_state: [u8; 16] = [
595 0x32, 0x43, 0xf6, 0xa8, 0x88, 0x5a, 0x30, 0x8d, 0x31, 0x31, 0x98, 0xa2, 0xe0, 0x37,
596 0x07, 0x34,
597 ];
598 let key = [
599 0x8e, 0x73, 0xb0, 0xf7, 0xda, 0x0e, 0x64, 0x52, 0xc8, 0x10, 0xf3, 0x2b, 0x80, 0x90,
600 0x79, 0xe5, 0x62, 0xf8, 0xea, 0xd2, 0x52, 0x2c, 0x6b, 0x7b,
601 ];
602 test_encrypt_block_test_vector_op::<6, 4, 12>(input_state, key)?;
603
604 // AES-256
605 let input_state: [u8; 16] = [
606 0x32, 0x43, 0xf6, 0xa8, 0x88, 0x5a, 0x30, 0x8d, 0x31, 0x31, 0x98, 0xa2, 0xe0, 0x37,
607 0x07, 0x34,
608 ];
609 let key = [
610 0x60, 0x3d, 0xeb, 0x10, 0x15, 0xca, 0x71, 0xbe, 0x2b, 0x73, 0xae, 0xf0, 0x85, 0x7d,
611 0x77, 0x81, 0x1f, 0x35, 0x2c, 0x07, 0x3b, 0x61, 0x08, 0xd7, 0x2d, 0x98, 0x10, 0xa3,
612 0x09, 0x14, 0xdf, 0xf4,
613 ];
614 test_encrypt_block_test_vector_op::<8, 4, 14>(input_state, key)?;
615
616 Ok(())
617 }
618
619 fn test_encrypt_block_test_vector_op<const NK: usize, const NB: usize, const NR: usize>(
620 input_state: [u8; 16],
621 key: [u8; NK * NB],
622 ) -> Result<()>
623 where
624 [(); 4 * (NR + 1)]:,
625 [(); 4 * NK]:,
626 // [(); NK * NB]:,
627 {
628 // Circuit declaration
629 let config = CircuitConfig::standard_recursion_config();
630 let mut builder = CircuitBuilder::<F, D>::new(config);
631
632 let key_target: [ByteArrayTarget; NK * NB] =
633 array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe()));
634 let sbox_lut = sbox_lut(&mut builder);
635 let mix_matrix = state_mix_matrix_bits(&mut builder);
636 let expanded_key_target: [[ByteArrayTarget; 4]; 4 * (NR + 1)] =
637 builder.key_expansion::<NK, NB, NR>(sbox_lut, key_target);
638
639 let input_state_target = builder.add_virtual_state();
640
641 let output = builder.encrypt_block(
642 sbox_lut,
643 mix_matrix,
644 input_state_target,
645 expanded_key_target,
646 );
647
648 println!(
649 "encrypt_block (NK:{}, NB:{}, NR:{}) num_gates: {}",
650 NK,
651 NB,
652 NR,
653 builder.num_gates()
654 );
655 let data = builder.build::<PoseidonGoldilocksConfig>();
656
657 // set values to circuit
658 let mut input_state_matrix: State = [[0; 4]; 4];
659 for i in 0..4 {
660 for j in 0..4 {
661 input_state_matrix[i][j] = input_state[i + 4 * j];
662 }
663 }
664
665 let expanded_key: [[u8; 4]; 4 * (NR + 1)] = key_expansion::<NK, NB, NR>(&key);
666 let native_ciphertext = encrypt_block::<NR>(&input_state, &expanded_key);
667
668 let mut pw = PartialWitness::<F>::new();
669
670 std::iter::zip(key_target, key).try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?;
671
672 pw.set_target_state(input_state_target, input_state_matrix)?;
673
674 std::iter::zip(expanded_key_target, expanded_key).try_for_each(|(t, v)| {
675 std::iter::zip(t, v).try_for_each(|(t, v)| pw.set_byte_array_target(t, v))
676 })?;
677
678 std::iter::zip(output.0, native_ciphertext).try_for_each(|(o, e)| {
679 std::iter::zip(o, e).try_for_each(|(o, e)| pw.set_byte_array_target(o, e))
680 })?;
681
682 let proof = data.prove(pw)?;
683 data.verify(proof)
684 }
685}