plonky2-crypto-gadgets/aes-gcm/src/circuit_aes.rs
aes-gcm/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 D,
17 constants::{RCON, SBOX},
18 native_aes::{State, rot_word, shift_rows},
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 ByteArrayTarget, CircuitBuilderAESState, D, PartialWitnessAESState,
404 PartialWitnessByteArray, sbox_lut, state_mix_matrix_bits,
405 };
406 use crate::native_aes::{State, encrypt_block, key_expansion, mix_columns, sub_bytes};
407
408 #[test]
409 fn test_sub_bytes() -> Result<()> {
410 // Circuit declaration
411 let config = CircuitConfig::standard_recursion_config();
412 let mut builder = CircuitBuilder::<F, D>::new(config);
413
414 let state_target = builder.add_virtual_state();
415 let sbox_lut = sbox_lut(&mut builder);
416 let out_state_target = builder.state_sub_bytes(sbox_lut, state_target);
417
418 let data = builder.build::<PoseidonGoldilocksConfig>();
419
420 // set values to circuit
421 let mut rng = rand::rng();
422 let test_states: [State; 10] = array::from_fn(|_| array::from_fn(|_| rng.random()));
423
424 test_states.into_iter().try_for_each(|s| {
425 let expected_result = sub_bytes(s);
426
427 let mut pw = PartialWitness::<F>::new();
428 pw.set_target_state(state_target, s)?;
429 pw.set_target_state(out_state_target, expected_result)?;
430
431 let proof = data.prove(pw)?;
432 data.verify(proof)
433 })
434 }
435
436 #[test]
437 fn test_mix_columns() -> Result<()> {
438 // Circuit declaration
439 let config = CircuitConfig::standard_recursion_config();
440 let mut builder = CircuitBuilder::<F, D>::new(config);
441
442 let state_target = builder.add_virtual_state();
443 let mix_matrix = state_mix_matrix_bits(&mut builder);
444 let out_state_target = builder.state_mix_columns(mix_matrix, state_target);
445
446 let data = builder.build::<PoseidonGoldilocksConfig>();
447
448 // set values to circuit
449 let mut rng = rand::rng();
450 let test_states: [State; 10] = array::from_fn(|_| array::from_fn(|_| rng.random()));
451
452 test_states.into_iter().try_for_each(|s| {
453 let expected_result = mix_columns(s);
454
455 let mut pw = PartialWitness::<F>::new();
456 pw.set_target_state(state_target, s)?;
457 pw.set_target_state(out_state_target, expected_result)?;
458
459 let proof = data.prove(pw)?;
460 data.verify(proof)
461 })
462 }
463
464 #[test]
465 fn test_gf_2_8_mul() -> Result<()> {
466 // Circuit declaration
467 let config = CircuitConfig::standard_recursion_config();
468 let mut builder = CircuitBuilder::<F, D>::new(config);
469
470 let a = builder.add_virtual_target();
471 let b = builder.add_virtual_target();
472
473 let a_times_b = builder.gf_2_8_mul(a, b);
474
475 let data = builder.build::<PoseidonGoldilocksConfig>();
476
477 // set values to circuit
478 let test_values = [
479 (0x57, 0x01, 0x57),
480 (0x57, 0x02, 0xae),
481 (0x57, 0x04, 0x47),
482 (0x57, 0x08, 0x8e),
483 (0x57, 0x10, 0x07),
484 (0x57, 0x20, 0x0e),
485 (0x57, 0x40, 0x1c),
486 (0x57, 0x80, 0x38),
487 (0x57, 0x13, 0xfe),
488 ];
489 test_values
490 .into_iter()
491 .map(|(a, b, c)| {
492 (
493 F::from_canonical_u8(a),
494 F::from_canonical_u8(b),
495 F::from_canonical_u8(c),
496 )
497 })
498 .try_for_each(|(a_val, b_val, expected_result)| {
499 let mut pw = PartialWitness::<F>::new();
500 pw.set_target(a, a_val)?;
501 pw.set_target(b, b_val)?;
502 pw.set_target(a_times_b, expected_result)?;
503
504 let proof = data.prove(pw)?;
505 data.verify(proof)
506 })
507 }
508
509 #[test]
510 fn test_key_expansion() -> Result<()> {
511 // AES-128
512 let key = [
513 0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf,
514 0x4f, 0x3c,
515 ];
516 test_key_expansion_op::</*Nk,Nb,Nr*/ 4, 4, 10>(key)?;
517
518 // AES-196
519 let key = [
520 0x8e, 0x73, 0xb0, 0xf7, 0xda, 0x0e, 0x64, 0x52, 0xc8, 0x10, 0xf3, 0x2b, 0x80, 0x90,
521 0x79, 0xe5, 0x62, 0xf8, 0xea, 0xd2, 0x52, 0x2c, 0x6b, 0x7b,
522 ];
523 test_key_expansion_op::</*Nk,Nb,Nr*/ 6, 4, 12>(key)?;
524
525 // AES-256
526 let key = [
527 0x60, 0x3d, 0xeb, 0x10, 0x15, 0xca, 0x71, 0xbe, 0x2b, 0x73, 0xae, 0xf0, 0x85, 0x7d,
528 0x77, 0x81, 0x1f, 0x35, 0x2c, 0x07, 0x3b, 0x61, 0x08, 0xd7, 0x2d, 0x98, 0x10, 0xa3,
529 0x09, 0x14, 0xdf, 0xf4,
530 ];
531 test_key_expansion_op::</*Nk,Nb,Nr*/ 8, 4, 14>(key)?;
532
533 Ok(())
534 }
535
536 /// Analogous tests to `native_aes::tests::test_key_expansion`.
537 fn test_key_expansion_op<const NK: usize, const NB: usize, const NR: usize>(
538 key: [u8; NK * NB],
539 ) -> Result<()>
540 where
541 [(); 4 * (NR + 1)]:,
542 [(); 4 * NK]:,
543 {
544 // Circuit declaration
545 let config = CircuitConfig::standard_recursion_config();
546 let mut builder = CircuitBuilder::<F, D>::new(config);
547
548 let key_target: [ByteArrayTarget; NK * NB] =
549 array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe()));
550 let sbox_lut = sbox_lut(&mut builder);
551 let expanded_key_target = builder.key_expansion::<NK, NB, NR>(sbox_lut, key_target);
552
553 println!(
554 "key_expansion(NK:{}, NB:{}, NR:{}) num_gates: {}",
555 NK,
556 NB,
557 NR,
558 builder.num_gates()
559 );
560 let data = builder.build::<PoseidonGoldilocksConfig>();
561
562 // set values to circuit
563 let w = key_expansion::<NK, NB, NR>(&key);
564
565 let mut pw = PartialWitness::<F>::new();
566
567 std::iter::zip(key_target, key).try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?;
568
569 std::iter::zip(expanded_key_target, w).try_for_each(|(t, v)| {
570 std::iter::zip(t, v).try_for_each(|(t, v)| pw.set_byte_array_target(t, v))
571 })?;
572
573 let proof = data.prove(pw)?;
574 data.verify(proof)
575 }
576
577 /// test against native version
578 #[test]
579 fn test_encrypt_block_test_vector() -> Result<()> {
580 // AES-128
581 let input_state: [u8; 16] = [
582 0x32, 0x43, 0xf6, 0xa8, 0x88, 0x5a, 0x30, 0x8d, 0x31, 0x31, 0x98, 0xa2, 0xe0, 0x37,
583 0x07, 0x34,
584 ];
585 let key = [
586 0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf,
587 0x4f, 0x3c,
588 ];
589 test_encrypt_block_test_vector_op::<4, 4, 10>(input_state, key)?;
590
591 // AES-196
592 let input_state: [u8; 16] = [
593 0x32, 0x43, 0xf6, 0xa8, 0x88, 0x5a, 0x30, 0x8d, 0x31, 0x31, 0x98, 0xa2, 0xe0, 0x37,
594 0x07, 0x34,
595 ];
596 let key = [
597 0x8e, 0x73, 0xb0, 0xf7, 0xda, 0x0e, 0x64, 0x52, 0xc8, 0x10, 0xf3, 0x2b, 0x80, 0x90,
598 0x79, 0xe5, 0x62, 0xf8, 0xea, 0xd2, 0x52, 0x2c, 0x6b, 0x7b,
599 ];
600 test_encrypt_block_test_vector_op::<6, 4, 12>(input_state, key)?;
601
602 // AES-256
603 let input_state: [u8; 16] = [
604 0x32, 0x43, 0xf6, 0xa8, 0x88, 0x5a, 0x30, 0x8d, 0x31, 0x31, 0x98, 0xa2, 0xe0, 0x37,
605 0x07, 0x34,
606 ];
607 let key = [
608 0x60, 0x3d, 0xeb, 0x10, 0x15, 0xca, 0x71, 0xbe, 0x2b, 0x73, 0xae, 0xf0, 0x85, 0x7d,
609 0x77, 0x81, 0x1f, 0x35, 0x2c, 0x07, 0x3b, 0x61, 0x08, 0xd7, 0x2d, 0x98, 0x10, 0xa3,
610 0x09, 0x14, 0xdf, 0xf4,
611 ];
612 test_encrypt_block_test_vector_op::<8, 4, 14>(input_state, key)?;
613
614 Ok(())
615 }
616
617 fn test_encrypt_block_test_vector_op<const NK: usize, const NB: usize, const NR: usize>(
618 input_state: [u8; 16],
619 key: [u8; NK * NB],
620 ) -> Result<()>
621 where
622 [(); 4 * (NR + 1)]:,
623 [(); 4 * NK]:,
624 // [(); NK * NB]:,
625 {
626 // Circuit declaration
627 let config = CircuitConfig::standard_recursion_config();
628 let mut builder = CircuitBuilder::<F, D>::new(config);
629
630 let key_target: [ByteArrayTarget; NK * NB] =
631 array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe()));
632 let sbox_lut = sbox_lut(&mut builder);
633 let mix_matrix = state_mix_matrix_bits(&mut builder);
634 let expanded_key_target: [[ByteArrayTarget; 4]; 4 * (NR + 1)] =
635 builder.key_expansion::<NK, NB, NR>(sbox_lut, key_target);
636
637 let input_state_target = builder.add_virtual_state();
638
639 let output = builder.encrypt_block(
640 sbox_lut,
641 mix_matrix,
642 input_state_target,
643 expanded_key_target,
644 );
645
646 println!(
647 "encrypt_block (NK:{}, NB:{}, NR:{}) num_gates: {}",
648 NK,
649 NB,
650 NR,
651 builder.num_gates()
652 );
653 let data = builder.build::<PoseidonGoldilocksConfig>();
654
655 // set values to circuit
656 let mut input_state_matrix: State = [[0; 4]; 4];
657 for i in 0..4 {
658 for j in 0..4 {
659 input_state_matrix[i][j] = input_state[i + 4 * j];
660 }
661 }
662
663 let expanded_key: [[u8; 4]; 4 * (NR + 1)] = key_expansion::<NK, NB, NR>(&key);
664 let native_ciphertext = encrypt_block::<NR>(&input_state, &expanded_key);
665
666 let mut pw = PartialWitness::<F>::new();
667
668 std::iter::zip(key_target, key).try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?;
669
670 pw.set_target_state(input_state_target, input_state_matrix)?;
671
672 std::iter::zip(expanded_key_target, expanded_key).try_for_each(|(t, v)| {
673 std::iter::zip(t, v).try_for_each(|(t, v)| pw.set_byte_array_target(t, v))
674 })?;
675
676 std::iter::zip(output.0, native_ciphertext).try_for_each(|(o, e)| {
677 std::iter::zip(o, e).try_for_each(|(o, e)| pw.set_byte_array_target(o, e))
678 })?;
679
680 let proof = data.prove(pw)?;
681 data.verify(proof)
682 }
683}