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