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