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