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