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