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