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