plonky2-crypto-gadgets/src/circuit_aes.rs
src/circuit_aes.rsRaw
1//! For LICENSE check out https://github.com/0xPARC/plonky2-aes/blob/main/LICENSE
2
3use std::{array, sync::Arc};
4
5use plonky2::{
6 field::{extension::Extendable, goldilocks_field::GoldilocksField as F, types::Field},
7 hash::hash_types::RichField,
8 iop::{
9 target::{BoolTarget, Target},
10 witness::{PartialWitness, WitnessWrite},
11 },
12 plonk::circuit_builder::CircuitBuilder,
13};
14
15use crate::{
16 constants::{RCON, SBOX},
17 native_aes::{rot_word, shift_rows, State},
18 D,
19};
20
21#[derive(Debug, Copy, Clone)]
22pub struct StateTarget(pub [[ByteArrayTarget; 4]; 4]);
23
24// TODO: Normalise terminology!
25pub type ByteArrayTarget = [BoolTarget; 8];
26
27pub trait CircuitBuilderAESState<F: RichField + Extendable<D>, const D: usize> {
28 /// Adds state target.
29 fn add_virtual_state(&mut self) -> StateTarget;
30
31 /// AES cipher as in spec.
32 fn encrypt_block<const NR: usize>(
33 &mut self,
34 sbox_lut_idx: usize,
35 mix_matrix: [[ByteArrayTarget; 4]; 4],
36 s: StateTarget,
37
38 w: [[ByteArrayTarget; 4]; 4 * (NR + 1)],
39 ) -> StateTarget;
40
41 /// Applies sub_bytes routine to state.
42 fn state_sub_bytes(&mut self, sbox_lut_idx: usize, s: StateTarget) -> StateTarget;
43
44 /// SubWord
45 fn state_sub_word(
46 &mut self,
47 sbox_lut_idx: usize,
48 word: [ByteArrayTarget; 4],
49 ) -> [ByteArrayTarget; 4];
50
51 /// MixColumns
52 fn state_mix_columns(
53 &mut self,
54 mix_matrix: [[ByteArrayTarget; 4]; 4],
55 s: StateTarget,
56 ) -> StateTarget;
57
58 /// AddRoundKey
59 fn state_add_round_key(
60 &mut self,
61 round_key: &[[ByteArrayTarget; 4]],
62 s: StateTarget,
63 ) -> StateTarget;
64
65 /// KeyExpansion
66 fn key_expansion<const NK: usize, const NB: usize, const NR: usize>(
67 &mut self,
68 sbox_lut_idx: usize,
69 key: [ByteArrayTarget; NK * NB],
70 ) -> [[ByteArrayTarget; 4]; 4 * (NR + 1)];
71
72 /// GF(2^8) addition
73 fn gf_2_8_add(&mut self, x: Target, y: Target) -> Target;
74
75 /// GF(2^8) addition actinb on bit arrays
76 fn gf_2_8_add_bits(
77 &mut self,
78 x_bits: [BoolTarget; 8],
79 y_bits: [BoolTarget; 8],
80 ) -> [BoolTarget; 8];
81
82 /// GF(2^8) multiplication
83 fn gf_2_8_mul(&mut self, x: Target, y: Target) -> Target;
84
85 /// GF(2^8) multiplication acting on bit arrays
86 fn gf_2_8_mul_bits(&mut self, x: [BoolTarget; 8], y: [BoolTarget; 8]) -> [BoolTarget; 8];
87
88 /// Bytearray inner product.
89 fn bytearray_ip_bits<const N: usize>(
90 &mut self,
91 x: [[BoolTarget; 8]; N],
92 y: [[BoolTarget; 8]; N],
93 ) -> [BoolTarget; 8];
94
95 /// Bytearray matrix application.
96 fn bytearray_matrix_apply_bits<const M: usize, const N: usize>(
97 &mut self,
98 a: [[ByteArrayTarget; N]; M],
99 x: [ByteArrayTarget; N],
100 ) -> [ByteArrayTarget; M];
101}
102
103impl CircuitBuilderAESState<F, D> for CircuitBuilder<F, D> {
104 fn add_virtual_state(&mut self) -> StateTarget {
105 StateTarget(array::from_fn(|_| {
106 array::from_fn(|_| array::from_fn(|_| self.add_virtual_bool_target_safe()))
107 }))
108 }
109
110 fn encrypt_block<const NR: usize>(
111 &mut self,
112 sbox_lut_idx: usize,
113 mix_matrix: [[ByteArrayTarget; 4]; 4], // constant mix_matrix
114 s: StateTarget,
115 w: [[ByteArrayTarget; 4]; 4 * (NR + 1)], // expanded key
116 ) -> StateTarget {
117 let mut s = s;
118 s = self.state_add_round_key(&w[0..4], s);
119 (1..NR).for_each(|i| {
120 s = self.state_sub_bytes(sbox_lut_idx, s);
121 s = StateTarget(shift_rows(s.0));
122 s = self.state_mix_columns(mix_matrix, s);
123 s = self.state_add_round_key(&w[4 * i..4 * (i + 1)], s);
124 });
125 s = self.state_sub_bytes(sbox_lut_idx, s);
126 s = StateTarget(shift_rows(s.0));
127 self.state_add_round_key(&w[4 * NR..4 * (NR + 1)], s)
128 }
129
130 fn state_sub_bytes(&mut self, sbox_lut_idx: usize, s: StateTarget) -> StateTarget {
131 StateTarget(array::from_fn(|i| {
132 self.state_sub_word(sbox_lut_idx, s.0[i])
133 }))
134 }
135
136 fn state_sub_word(
137 &mut self,
138 sbox_lut_idx: usize,
139 word: [ByteArrayTarget; 4],
140 ) -> [ByteArrayTarget; 4] {
141 array::from_fn(|i| {
142 let byte_target = target_from_bitarray(self, &word[i]);
143 let out_target = self.add_lookup_from_index(byte_target, sbox_lut_idx);
144 bitarray_from_bytetarget(self, out_target)
145 })
146 }
147
148 fn state_mix_columns(
149 &mut self,
150 mix_matrix: [[ByteArrayTarget; 4]; 4],
151 s: StateTarget,
152 ) -> StateTarget {
153 let cols: [_; 4] = array::from_fn(|i| array::from_fn(|j| s.0[j][i]));
154 let out_cols: [_; 4] =
155 array::from_fn(|i| self.bytearray_matrix_apply_bits(mix_matrix, cols[i]));
156 StateTarget(array::from_fn(|i| array::from_fn(|j| out_cols[j][i])))
157 }
158
159 fn state_add_round_key(
160 &mut self,
161 round_key: &[[ByteArrayTarget; 4]],
162 s: StateTarget,
163 ) -> StateTarget {
164 StateTarget(array::from_fn(|i| {
165 array::from_fn(|j| self.gf_2_8_add_bits(s.0[i][j], round_key[j][i]))
166 }))
167 }
168
169 fn key_expansion<const NK: usize, const NB: usize, const NR: usize>(
170 &mut self,
171 sbox_lut_idx: usize,
172 key: [ByteArrayTarget; NK * NB],
173 ) -> [[ByteArrayTarget; 4]; 4 * (NR + 1)] {
174 let rcon: [ByteArrayTarget; 11] = array::from_fn(|i| {
175 let rcon_bits = le_bits_from_byte(RCON[i]);
176 array::from_fn(|j| self.constant_bool(rcon_bits[j]))
177 });
178
179 let key: [[_; 4]; NK] = array::from_fn(|i| array::from_fn(|j| key[4 * i + j]));
180
181 let key_additions = (NK..4 * (NR + 1))
182 .scan(key.to_vec(), |st, i| {
183 let offset = if i % NK == 0 {
184 let term = self.state_sub_word(sbox_lut_idx, rot_word(st[i - 1]));
185 array::from_fn(|j| {
186 if j == 0 {
187 self.gf_2_8_add_bits(term[j], rcon[i / NK])
188 } else {
189 term[j]
190 }
191 })
192 } else if (NK > 6) && (i % NK == 4) {
193 self.state_sub_word(sbox_lut_idx, st[i - 1])
194 } else {
195 st[i - 1]
196 };
197
198 let cur_val = array::from_fn(|j| self.gf_2_8_add_bits(st[i - NK][j], offset[j]));
199 st.push(cur_val);
200 Some(cur_val)
201 })
202 .collect::<Vec<_>>();
203 std::array::from_fn(|i| {
204 if i < NK {
205 key[i]
206 } else {
207 key_additions[i - NK]
208 }
209 })
210 }
211
212 fn gf_2_8_add(&mut self, x: Target, y: Target) -> Target {
213 let x_bits = bitarray_from_bytetarget(self, x);
214 let y_bits = bitarray_from_bytetarget(self, y);
215
216 let sum_bits = self.gf_2_8_add_bits(x_bits, y_bits);
217 target_from_bitarray(self, &sum_bits)
218 }
219
220 fn gf_2_8_add_bits(
221 &mut self,
222 x_bits: [BoolTarget; 8],
223 y_bits: [BoolTarget; 8],
224 ) -> [BoolTarget; 8] {
225 let sum_bits = std::iter::zip(x_bits, y_bits)
226 .map(|(a, b)| xor(self, a, b))
227 .collect::<Vec<_>>();
228
229 array::from_fn(|i| sum_bits[i])
230 }
231
232 fn gf_2_8_mul(&mut self, x: Target, y: Target) -> Target {
233 let x_bits = bitarray_from_bytetarget(self, x);
234 let y_bits = bitarray_from_bytetarget(self, y);
235
236 let prod_bits = self.gf_2_8_mul_bits(x_bits, y_bits);
237 target_from_bitarray(self, &prod_bits)
238 }
239
240 fn gf_2_8_mul_bits(
241 &mut self,
242 x_bits: [BoolTarget; 8],
243 y_bits: [BoolTarget; 8],
244 ) -> [BoolTarget; 8] {
245 let zero_bits = vec![self._false(); 8];
246 assert!(x_bits.len() == 8 && y_bits.len() == 8);
247
248 let x_times = |builder: &mut CircuitBuilder<_, _>, y: &[BoolTarget]| -> [_; 8] {
249 let zero = builder._false();
250 let one = builder._true();
251
252 let potential_offset = [one, one, zero, one, one, zero, zero, zero];
253 let y_shifted: [_; 8] = array::from_fn(|i| if i == 0 { zero } else { y[i - 1] });
254
255 let offset: [_; 8] = array::from_fn(|i| builder.and(y[7], potential_offset[i]));
256 array::from_fn(|i| xor(builder, y_shifted[i], offset[i]))
257 };
258
259 let powers_times_x = (0..8)
260 .scan(x_bits.to_vec(), |st, i| {
261 if i > 0 {
262 *st = x_times(self, st).to_vec();
263 }
264 Some(st.clone())
265 })
266 .collect::<Vec<_>>();
267
268 let prod_bits = std::iter::zip(y_bits, powers_times_x)
269 .map(|(c, o)| o.into_iter().map(|b| self.and(c, b)).collect::<Vec<_>>())
270 .collect::<Vec<_>>()
271 .into_iter()
272 .fold(zero_bits, |acc, term| {
273 std::iter::zip(acc, term)
274 .map(|(a, b)| xor(self, a, b))
275 .collect::<Vec<_>>()
276 });
277
278 array::from_fn(|i| prod_bits[i])
279 }
280
281 fn bytearray_ip_bits<const N: usize>(
282 &mut self,
283 x: [[BoolTarget; 8]; N],
284 y: [[BoolTarget; 8]; N],
285 ) -> [BoolTarget; 8] {
286 let zero = self.zero();
287 let zero_bits = bitarray_from_bytetarget(self, zero);
288
289 std::iter::zip(x, y).fold(zero_bits, |acc, (a, b)| {
290 let prod = self.gf_2_8_mul_bits(a, b);
291 self.gf_2_8_add_bits(acc, prod)
292 })
293 }
294
295 fn bytearray_matrix_apply_bits<const M: usize, const N: usize>(
296 &mut self,
297 a: [[ByteArrayTarget; N]; M],
298 x: [ByteArrayTarget; N],
299 ) -> [ByteArrayTarget; M] {
300 std::array::from_fn(|i| self.bytearray_ip_bits(a[i], x))
301 }
302}
303
304pub trait PartialWitnessByteArray {
305 fn set_byte_array_target(&mut self, target: ByteArrayTarget, value: u8) -> anyhow::Result<()>;
306}
307
308impl<F: Field> PartialWitnessByteArray for PartialWitness<F> {
309 fn set_byte_array_target(&mut self, target: ByteArrayTarget, value: u8) -> anyhow::Result<()> {
310 let value_bits = le_bits_from_byte(value);
311 std::iter::zip(target, value_bits).try_for_each(|(t, v)| self.set_bool_target(t, v))
312 }
313}
314
315pub trait PartialWitnessAESState {
316 fn set_target_state(&mut self, target: StateTarget, value: State) -> anyhow::Result<()>;
317}
318
319impl<F: Field> PartialWitnessAESState for PartialWitness<F> {
320 fn set_target_state(&mut self, target: StateTarget, value: State) -> anyhow::Result<()> {
321 std::iter::zip(target.0, value).try_for_each(|(t, v)| {
322 std::iter::zip(t, v).try_for_each(|(t, v)| self.set_byte_array_target(t, v))
323 })
324 }
325}
326
327pub fn sbox_lut(builder: &mut CircuitBuilder<F, D>) -> usize {
328 builder.add_lookup_table_from_pairs(Arc::new(
329 SBOX.into_iter()
330 .enumerate()
331 .map(|(i, o)| (i as u16, o as u16))
332 .collect::<Vec<_>>(),
333 ))
334}
335
336pub fn state_mix_matrix_bits(builder: &mut CircuitBuilder<F, D>) -> [[[BoolTarget; 8]; 4]; 4] {
337 let one = builder.one();
338 let two = builder.two();
339 let three = builder.constant(F::from_canonical_u64(3));
340
341 let matrix = [
342 [two, three, one, one],
343 [one, two, three, one],
344 [one, one, two, three],
345 [three, one, one, two],
346 ];
347
348 array::from_fn(|i| array::from_fn(|j| bitarray_from_bytetarget(builder, matrix[i][j])))
349}
350
351pub fn xor(builder: &mut CircuitBuilder<F, D>, x: BoolTarget, y: BoolTarget) -> BoolTarget {
352 let x_or_y = builder.or(x, y);
353 BoolTarget::new_unsafe(builder.arithmetic(
354 F::NEG_ONE,
355 F::ONE,
356 x.target,
357 y.target,
358 x_or_y.target,
359 ))
360}
361
362// LE bit array conversions
363pub fn target_from_bitarray(builder: &mut CircuitBuilder<F, D>, x: &[BoolTarget]) -> Target {
364 let zero = builder.zero();
365 let two = builder.two();
366 x.iter()
367 .rev()
368 .fold(zero, |acc, b| builder.mul_add(two, acc, b.target))
369}
370
371pub fn bitarray_from_bytetarget(builder: &mut CircuitBuilder<F, D>, x: Target) -> [BoolTarget; 8] {
372 let x_bits = builder.split_le(x, 8);
373 array::from_fn(|i| x_bits[i])
374}
375
376pub fn le_bits_from_byte(v: u8) -> [bool; 8] {
377 let v_bits = (0..8)
378 .scan(v, |st, _| {
379 let cur_bit = (*st % 2) != 0;
380 *st >>= 1;
381 Some(cur_bit)
382 })
383 .collect::<Vec<_>>();
384 array::from_fn(|i| v_bits[i])
385}
386
387#[cfg(test)]
388mod tests {
389 use std::array;
390
391 use anyhow::Result;
392 use plonky2::{
393 field::{goldilocks_field::GoldilocksField as F, types::Field},
394 iop::witness::{PartialWitness, WitnessWrite},
395 plonk::{
396 circuit_builder::CircuitBuilder, circuit_data::CircuitConfig,
397 config::PoseidonGoldilocksConfig,
398 },
399 };
400 use rand::RngExt;
401
402 use super::{
403 sbox_lut, state_mix_matrix_bits, ByteArrayTarget, CircuitBuilderAESState,
404 PartialWitnessAESState, PartialWitnessByteArray,
405 };
406 use crate::native_aes::{encrypt_block, key_expansion, mix_columns, sub_bytes, State};
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_aes::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}