arnaucube
arnaucube/plonky2-crypto-gadgetsPublic Self-hosted repository
plonky2-crypto-gadgets/src/circuit_gcm.rs
src/circuit_gcm.rsRaw
  1//! For LICENSE check out https://github.com/0xPARC/plonky2-aes/blob/main/LICENSE
  2
  3use std::array;
  4
  5use plonky2::{
  6    field::{goldilocks_field::GoldilocksField as F, types::Field},
  7    iop::target::BoolTarget,
  8    plonk::circuit_builder::CircuitBuilder,
  9};
 10
 11use crate::{
 12    circuit_aes::{
 13        le_bits_from_byte, sbox_lut, state_mix_matrix_bits, xor, ByteArrayTarget as ByteTarget,
 14        CircuitBuilderAESState, StateTarget,
 15    },
 16    native_gcm::TAG_LEN,
 17    D,
 18};
 19
 20/// L: max size of plaintext to cipher.
 21pub fn encrypt_target<
 22    const NK: usize,
 23    const NB: usize,
 24    const NR: usize,
 25    const L: usize,
 26    const TAG: bool,
 27>(
 28    builder: &mut CircuitBuilder<F, D>,
 29    key: [ByteTarget; NK * NB],
 30    nonce: &[ByteTarget; 12],
 31    pt: &[ByteTarget; L],
 32) -> ([ByteTarget; L], [ByteTarget; TAG_LEN / 8])
 33where
 34    [(); NK * NB]:,
 35    [(); 4 * (NR + 1)]:,
 36{
 37    let sbox_lut = sbox_lut(builder);
 38    let mix_matrix = state_mix_matrix_bits(builder);
 39
 40    let a: &[ByteTarget] = &[]; // additional authenticated data // TODO maybe as input
 41
 42    let expanded_key = builder.key_expansion::<NK, NB, NR>(sbox_lut, key);
 43
 44    let empty_state = builder.empty_state();
 45
 46    // 1. CIPH_K
 47    let h = builder
 48        .encrypt_block::<NR>(sbox_lut, mix_matrix, empty_state, expanded_key)
 49        .flatten();
 50
 51    // 2. J_0
 52    let zero_byte = builder.zero_byte();
 53    let one_byte = array::from_fn(|i| {
 54        // 0x01
 55        if i == 0 {
 56            builder._true()
 57        } else {
 58            builder._false()
 59        }
 60    });
 61    let j0: [ByteTarget; 16] = if nonce.len() * 8 == 96 {
 62        // J_0 = IV || 0^31 || 1
 63        let mut out = [zero_byte; 16];
 64        out[..12].copy_from_slice(nonce.as_slice());
 65        out[12..16].copy_from_slice(&[zero_byte, zero_byte, zero_byte, one_byte]);
 66        out
 67    } else {
 68        panic!("unsuported at initial version; nonce.len()=12 (96 bits)");
 69    };
 70
 71    // 3. C=GCTR()
 72    let inc32_j0 = inc32_target(builder, j0);
 73    let c = gctr_target::<NR, L>(builder, sbox_lut, mix_matrix, expanded_key, &inc32_j0, &pt);
 74
 75    // the rest of the logic is for the tag; for some use cases we might skip it
 76    // (saving a notable amount of gates, about double)
 77    if !TAG {
 78        let t: [ByteTarget; TAG_LEN / 8] = [zero_byte; TAG_LEN / 8];
 79        return (c, t);
 80    }
 81
 82    // 4. u, v
 83    let u: usize = 16 * (L as f64 / 16_f64).ceil() as usize - L; // L=x.len()
 84    let v: usize = 16 * (a.len() as f64 / 16_f64).ceil() as usize - a.len();
 85
 86    // 5. S = GHASH()
 87    let const_a_len_u8: [u8; 8] = (a.len() * 8).to_be_bytes(); // const
 88    let const_c_len_u8: [u8; 8] = (c.len() * 8).to_be_bytes(); // const
 89    let a_len: [ByteTarget; 8] = array::from_fn(|byte_i| {
 90        let const_bits = le_bits_from_byte(const_a_len_u8[byte_i]);
 91        let byte: ByteTarget = array::from_fn(|bit_i| {
 92            let bit_goldilocks = if const_bits[bit_i] { F::ONE } else { F::ZERO };
 93            BoolTarget::new_unsafe(builder.constant(bit_goldilocks))
 94        });
 95        byte
 96    });
 97    let c_len: [ByteTarget; 8] = array::from_fn(|byte_i| {
 98        let const_bits = le_bits_from_byte(const_c_len_u8[byte_i]);
 99        let byte: ByteTarget = array::from_fn(|bit_i| {
100            let bit_goldilocks = if const_bits[bit_i] { F::ONE } else { F::ZERO };
101            BoolTarget::new_unsafe(builder.constant(bit_goldilocks))
102        });
103        byte
104    });
105    let ghash_input_vec = vec![
106        a.to_vec(),
107        vec![zero_byte; v],
108        c.to_vec(),
109        vec![zero_byte; u],
110        a_len.to_vec(),
111        c_len.to_vec(),
112    ]
113    .concat();
114    let s = ghash_target(builder, h, ghash_input_vec);
115
116    // 6. T=MSB(GCTR()))
117    let msb_input = gctr_target(builder, sbox_lut, mix_matrix, expanded_key, &j0, &s);
118    let t_vec = msb_t_target(builder, TAG_LEN / 8, &msb_input);
119    let mut t: [ByteTarget; TAG_LEN / 8] = [zero_byte; TAG_LEN / 8];
120    t.copy_from_slice(&t_vec);
121    (c, t)
122}
123
124/// L: max size of plaintext to cipher.
125fn gctr_target<const NR: usize, const L: usize>(
126    builder: &mut CircuitBuilder<F, D>,
127    sbox_lut_idx: usize,
128    mix_matrix: [[ByteTarget; 4]; 4],
129    key: [[ByteTarget; 4]; 4 * (NR + 1)],
130    icb: &[ByteTarget; 16],
131    x: &[ByteTarget; L],
132) -> [ByteTarget; L] {
133    let n = ((L * 8) as f64 / 128_f64).ceil() as usize;
134
135    let mut y: [ByteTarget; L] = x.clone();
136    let mut cb_i = *icb;
137
138    let zero_byte = builder.zero_byte();
139    for (i, x_i_raw) in x.chunks(16).enumerate() {
140        if i > 0 {
141            cb_i = inc32_target(builder, cb_i);
142        }
143
144        let l = x_i_raw.len().min(16);
145        let mut x_i = [zero_byte; 16];
146        x_i[..l].copy_from_slice(x_i_raw);
147
148        let ciph_cb_i = builder
149            .encrypt_block::<NR>(sbox_lut_idx, mix_matrix, StateTarget::from_flat(cb_i), key)
150            .flatten();
151
152        let (y_i, n_bytes) = if i < n && x_i_raw.len() == 16 {
153            (xor_blocks(builder, x_i, ciph_cb_i), 16)
154        } else {
155            // last chunk, might be smaller than 16 bytes (L%16 bytes)
156            let msb_res = msb_t_target(builder, L % 16, ciph_cb_i.as_slice());
157            let l = msb_res.len().min(16);
158            let mut m: [ByteTarget; 16] = [zero_byte; 16];
159            m[..l].copy_from_slice(&msb_res);
160            (xor_blocks(builder, x_i, m), l)
161        };
162        y[i * 16..i * 16 + n_bytes].clone_from_slice(&y_i[..n_bytes]);
163    }
164    y
165}
166
167// fn ghash_target<const L: usize>(
168fn ghash_target(
169    builder: &mut CircuitBuilder<F, D>,
170    h: [ByteTarget; 16],
171    // x: &[ByteTarget; L],
172    x: Vec<ByteTarget>,
173) -> [ByteTarget; 16] {
174    // assert!(L.is_multiple_of(16)); // multiple of 128 bits
175    // let m = L / 16;
176    assert!(x.len().is_multiple_of(16)); // multiple of 128 bits
177    let m = x.len() / 16;
178
179    let zero_byte = builder.zero_byte();
180    let mut y = [zero_byte; 16]; // (128 bits)
181    for i in 0..m {
182        let mut xi = [zero_byte; 16];
183        xi.clone_from_slice(&x[i * 16..i * 16 + 16]);
184        let y_xi = xor_blocks(builder, y, xi);
185        y = gf_2_128_mul_target(builder, y_xi, h);
186    }
187    y
188}
189fn gf_2_128_mul_target(
190    builder: &mut CircuitBuilder<F, D>,
191    x: [ByteTarget; 16],
192    y: [ByteTarget; 16],
193) -> [ByteTarget; 16] {
194    let zero_byte = builder.zero_byte();
195    let one_bool = builder._true();
196
197    // R: 10000111 || 0^120 (in little-endian)
198    let mut r = [zero_byte; 16];
199    r[0][0] = one_bool;
200    r[0][5] = one_bool;
201    r[0][6] = one_bool;
202    r[0][7] = one_bool;
203
204    let mut z = [zero_byte; 16];
205    let mut v = y;
206    for i in 0..128 {
207        let byte_index = i / 8;
208        let bit_index = 7 - (i % 8);
209        let xi = x[byte_index][bit_index];
210
211        // set z = if xi==1: z^v, else: z
212        for b in 0..16 {
213            for k in 0..8 {
214                let z_xor_v = xor(builder, z[b][k], v[b][k]);
215                z[b][k] =
216                    BoolTarget::new_unsafe(builder.select(xi, z_xor_v.target, z[b][k].target));
217            }
218        }
219
220        let lsb = v[15][0].clone(); // (little-endian)
221        v = right_shift_one_target(builder, &v);
222
223        // if lsb==1: v=v^R, else: v
224        for b in 0..16 {
225            for k in 0..8 {
226                let v_xor_r = xor(builder, v[b][k], r[b][k]);
227                v[b][k] =
228                    BoolTarget::new_unsafe(builder.select(lsb, v_xor_r.target, v[b][k].target));
229            }
230        }
231    }
232    z
233}
234pub fn right_shift_one_target(
235    builder: &mut CircuitBuilder<F, D>,
236    v: &[ByteTarget; 16],
237) -> [ByteTarget; 16] {
238    let mut r: [ByteTarget; 16] = v.clone();
239
240    let mut carry = builder._false();
241    for i in 0..16 {
242        let current = v[i];
243
244        let next_carry = current[0];
245
246        let mut shifted = [builder._false(); 8];
247        for j in 0..7 {
248            shifted[j] = current[j + 1];
249        }
250        shifted[7] = carry;
251        r[i] = shifted;
252        carry = next_carry;
253    }
254    r
255}
256
257pub fn inc32_target(
258    builder: &mut CircuitBuilder<F, D>,
259    block: [ByteTarget; 16],
260) -> [ByteTarget; 16] {
261    let mut r = block;
262
263    let mut carry = BoolTarget::new_unsafe(builder.one());
264    for byte_index in (12..16).rev() {
265        for bit_index in 0..8 {
266            let a = block[byte_index][bit_index];
267            // builder.assert_bool(a); // TODO we can assume it's valid bool, skip check
268
269            let sum = xor(builder, a, carry);
270            let carry_out = builder.and(a, carry);
271            r[byte_index][bit_index] = sum;
272            carry = carry_out;
273        }
274    }
275    r
276}
277
278fn xor_byte(builder: &mut CircuitBuilder<F, D>, b1: ByteTarget, b2: ByteTarget) -> ByteTarget {
279    array::from_fn(|i| xor(builder, b1[i], b2[i]))
280}
281
282fn xor_blocks(
283    builder: &mut CircuitBuilder<F, D>,
284    b1: [ByteTarget; 16],
285    b2: [ByteTarget; 16],
286) -> [ByteTarget; 16] {
287    array::from_fn(|i| xor_byte(builder, b1[i], b2[i]))
288}
289
290pub fn msb_t_target(
291    builder: &mut CircuitBuilder<F, D>,
292    t_bytes: usize,
293    block: &[ByteTarget],
294) -> Vec<ByteTarget> {
295    let t = t_bytes * 8;
296
297    let full_bytes = t / 8;
298    let rem_bits = t % 8;
299
300    let mut out: Vec<ByteTarget> = Vec::new();
301    for i in 0..full_bytes {
302        out.push(block[i]);
303    }
304
305    let zero_bool = BoolTarget::new_unsafe(builder.zero());
306    if rem_bits != 0 {
307        let mut partial = [zero_bool; 8];
308
309        for bit in 0..8 {
310            if bit < rem_bits {
311                partial[bit] = block[full_bytes][bit];
312            } else {
313                partial[bit] = zero_bool;
314            }
315        }
316        out.push(partial);
317    }
318    out
319}
320
321#[cfg(test)]
322mod tests {
323    use std::array;
324
325    use anyhow::Result;
326    use plonky2::{
327        field::{goldilocks_field::GoldilocksField as F, types::Field},
328        iop::witness::{PartialWitness, WitnessWrite},
329        plonk::{
330            circuit_builder::CircuitBuilder, circuit_data::CircuitConfig,
331            config::PoseidonGoldilocksConfig,
332        },
333    };
334
335    use super::*;
336    use crate::{
337        circuit_aes::PartialWitnessByteArray,
338        native_aes::key_expansion,
339        native_gcm::{gctr, gf_2_128_mul, ghash, right_shift_one},
340    };
341
342    use super::D;
343
344    #[test]
345    fn test_gctr() -> Result<()> {
346        // AES-GCM-128
347        test_gctr_opt::<4, 4, 10, 13>()?;
348        test_gctr_opt::<4, 4, 10, 17>()?;
349
350        // AES-GCM-256
351        test_gctr_opt::<8, 4, 14, 13>()?;
352
353        Ok(())
354    }
355    fn test_gctr_opt<const NK: usize, const NB: usize, const NR: usize, const L: usize>(
356    ) -> Result<()>
357    where
358        [(); 4 * (NR + 1)]:,
359        [(); NK * NB]:,
360    {
361        let key: &[u8; NK * NB] = &[42; NK * NB];
362        let expanded_key = key_expansion::<NK, NB, NR>(key);
363        let nonce: &[u8; 12] = &[111; 12];
364        let icb: &[u8; 16] = &[222; 16];
365        let pt: &[u8; L] = &[42u8; L];
366
367        let expected = gctr(expanded_key, &icb, pt);
368
369        // Circuit declaration
370        let config = CircuitConfig::standard_recursion_config();
371        let mut builder = CircuitBuilder::<F, D>::new(config);
372
373        let key_target: [ByteTarget; NK * NB] =
374            array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe()));
375        let nonce_target: [ByteTarget; 12] =
376            array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe()));
377        let icb_target: [ByteTarget; 16] =
378            array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe()));
379        let pt_target: [ByteTarget; L] =
380            array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe()));
381        let sbox_lut = sbox_lut(&mut builder);
382        let mix_matrix = state_mix_matrix_bits(&mut builder);
383
384        let key_expanded_target = builder.key_expansion::<NK, NB, NR>(sbox_lut, key_target);
385        let gctr_out_target = gctr_target::<NR, L>(
386            &mut builder,
387            sbox_lut,
388            mix_matrix,
389            key_expanded_target,
390            &icb_target,
391            &pt_target,
392        );
393
394        println!(
395            "gctr (NK:{}, NB:{}, NR:{}, L:{}) num_gates: {}",
396            NK,
397            NB,
398            NR,
399            L,
400            builder.num_gates()
401        );
402        let data = builder.build::<PoseidonGoldilocksConfig>();
403
404        // set values to circuit
405        let mut pw = PartialWitness::<F>::new();
406        std::iter::zip(key_target, key).try_for_each(|(t, v)| pw.set_byte_array_target(t, *v))?;
407        std::iter::zip(nonce_target, nonce)
408            .try_for_each(|(t, v)| pw.set_byte_array_target(t, *v))?;
409        std::iter::zip(icb_target, icb).try_for_each(|(t, v)| pw.set_byte_array_target(t, *v))?;
410        std::iter::zip(pt_target, pt).try_for_each(|(t, v)| pw.set_byte_array_target(t, *v))?;
411
412        std::iter::zip(gctr_out_target, expected)
413            .try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?;
414
415        let proof = data.prove(pw)?;
416        data.verify(proof)
417    }
418
419    #[test]
420    fn test_right_shift_one() -> Result<()> {
421        let x: [u8; 16] = [111; 16];
422
423        let mut expected = x.clone();
424        right_shift_one(&mut expected);
425
426        // Circuit declaration
427        let config = CircuitConfig::standard_recursion_config();
428        let mut builder = CircuitBuilder::<F, D>::new(config);
429
430        let x_target: [ByteTarget; 16] =
431            array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe()));
432
433        let out_target: [ByteTarget; 16] = right_shift_one_target(&mut builder, &x_target);
434
435        println!("right_shift_one num_gates: {}", builder.num_gates());
436        let data = builder.build::<PoseidonGoldilocksConfig>();
437
438        // set values to circuit
439        let mut pw = PartialWitness::<F>::new();
440        std::iter::zip(x_target, x).try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?;
441
442        std::iter::zip(out_target, expected)
443            .try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?;
444
445        let proof = data.prove(pw)?;
446        data.verify(proof)?;
447        Ok(())
448    }
449
450    #[test]
451    fn test_gf_mul() -> Result<()> {
452        // AES-GCM-128
453        test_gf_mul_opt::<4, 4, 10, 16>()?;
454        test_gf_mul_opt::<4, 4, 10, 32>()?;
455
456        // AES-GCM-256
457        test_gf_mul_opt::<8, 4, 14, 16>()?;
458
459        Ok(())
460    }
461    fn test_gf_mul_opt<const NK: usize, const NB: usize, const NR: usize, const L: usize>(
462    ) -> Result<()>
463    where
464        [(); 4 * (NR + 1)]:,
465        [(); NK * NB]:,
466    {
467        let x: [u8; 16] = [111; 16];
468        let y: [u8; 16] = [222; 16];
469
470        let expected = gf_2_128_mul(x, y);
471        let x_bits: [[bool; 8]; 16] =
472            array::from_fn(|b| crate::circuit_aes::le_bits_from_byte(x[b]));
473        let y_bits: [[bool; 8]; 16] =
474            array::from_fn(|b| crate::circuit_aes::le_bits_from_byte(y[b]));
475
476        // sanity check
477        let expected2 = crate::native_gcm::gf_2_128_mul_circuit_version(x_bits, y_bits);
478        let expected_bits: [[bool; 8]; 16] =
479            array::from_fn(|b| crate::circuit_aes::le_bits_from_byte(expected[b]));
480        assert_eq!(expected2, expected_bits);
481
482        // Circuit declaration
483        let config = CircuitConfig::standard_recursion_config();
484        let mut builder = CircuitBuilder::<F, D>::new(config);
485
486        let x_target: [ByteTarget; 16] =
487            array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe()));
488        let y_target: [ByteTarget; 16] =
489            array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe()));
490
491        let out_target = gf_2_128_mul_target(&mut builder, x_target, y_target);
492
493        println!("gf_2_128_mul (L:{}) num_gates: {}", L, builder.num_gates());
494        let data = builder.build::<PoseidonGoldilocksConfig>();
495
496        // set values to circuit
497        let mut pw = PartialWitness::<F>::new();
498        std::iter::zip(x_target, x).try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?;
499        std::iter::zip(y_target, y).try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?;
500
501        std::iter::zip(out_target, expected)
502            .try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?;
503
504        let proof = data.prove(pw)?;
505        data.verify(proof)
506    }
507
508    #[test]
509    fn test_ghash() -> Result<()> {
510        // AES-GCM-128
511        test_ghash_opt::<4, 4, 10, 16>()?;
512        test_ghash_opt::<4, 4, 10, 32>()?;
513
514        // AES-GCM-256
515        test_ghash_opt::<8, 4, 14, 16>()?;
516
517        Ok(())
518    }
519    fn test_ghash_opt<const NK: usize, const NB: usize, const NR: usize, const L: usize>(
520    ) -> Result<()>
521    where
522        [(); 4 * (NR + 1)]:,
523        [(); NK * NB]:,
524    {
525        let h: [u8; 16] = [222; 16];
526        let x: &[u8; L] = &[42u8; L];
527
528        let expected = ghash(h, x);
529
530        // Circuit declaration
531        let config = CircuitConfig::standard_recursion_config();
532        let mut builder = CircuitBuilder::<F, D>::new(config);
533
534        let h_target: [ByteTarget; 16] =
535            array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe()));
536        let x_target: [ByteTarget; L] =
537            array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe()));
538
539        // let ghash_out_target = ghash_target::<L>(&mut builder, h_target, &x_target);
540        let ghash_out_target = ghash_target(&mut builder, h_target, x_target.to_vec());
541
542        println!("ghash (L:{}) num_gates: {}", L, builder.num_gates());
543        let data = builder.build::<PoseidonGoldilocksConfig>();
544
545        // set values to circuit
546        let mut pw = PartialWitness::<F>::new();
547        std::iter::zip(h_target, h).try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?;
548        std::iter::zip(x_target, x).try_for_each(|(t, v)| pw.set_byte_array_target(t, *v))?;
549
550        std::iter::zip(ghash_out_target, expected)
551            .try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?;
552
553        let proof = data.prove(pw)?;
554        data.verify(proof)
555    }
556
557    #[test]
558    fn test_encrypt() -> Result<()> {
559        // AES-GCM-128
560        test_encrypt_opt::<4, 4, 10, 13, false>(false)?;
561        test_encrypt_opt::<4, 4, 10, 13, true>(false)?; // with tag
562        test_encrypt_opt::<4, 4, 10, 17, false>(false)?;
563
564        // AES-GCM-256
565        test_encrypt_opt::<8, 4, 14, 13, false>(false)?;
566
567        Ok(())
568    }
569    // run it with: cargo test -- --ignored
570    #[ignore]
571    #[test]
572    fn test_encrypt_report_sizes() -> Result<()> {
573        // AES-GCM-128
574        test_encrypt_opt::<4, 4, 10, 16, false>(true)?;
575        test_encrypt_opt::<4, 4, 10, 17, false>(true)?;
576        test_encrypt_opt::<4, 4, 10, 32, false>(true)?;
577        test_encrypt_opt::<4, 4, 10, 33, false>(true)?;
578        test_encrypt_opt::<4, 4, 10, 64, false>(true)?;
579        test_encrypt_opt::<4, 4, 10, 128, false>(true)?;
580        test_encrypt_opt::<4, 4, 10, 256, false>(true)?;
581        test_encrypt_opt::<4, 4, 10, 512, false>(true)?;
582        test_encrypt_opt::<4, 4, 10, 1024, false>(true)?;
583        test_encrypt_opt::<4, 4, 10, 2048, false>(true)?;
584
585        Ok(())
586    }
587    fn test_encrypt_opt<
588        const NK: usize,
589        const NB: usize,
590        const NR: usize,
591        const L: usize,
592        const TAG: bool,
593    >(
594        only_build: bool,
595    ) -> Result<()>
596    where
597        [(); 4 * (NR + 1)]:,
598        [(); NK * NB]:,
599    {
600        let key: &[u8; NK * NB] = &[42; NK * NB];
601        let nonce: &[u8; 12] = &[111; 12];
602        let pt: &[u8; L] = &[42u8; L];
603
604        let (expected_c, expected_tag) = crate::native_gcm::encrypt::<NK, NB, NR>(key, nonce, pt);
605
606        // Circuit declaration
607        let config = CircuitConfig::standard_recursion_config();
608        let mut builder = CircuitBuilder::<F, D>::new(config);
609
610        let key_target: [ByteTarget; NK * NB] =
611            array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe()));
612        let nonce_target: [ByteTarget; 12] =
613            array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe()));
614        let pt_target: [ByteTarget; L] =
615            array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe()));
616        let (circuit_c, circuit_tag) = encrypt_target::<NK, NB, NR, L, TAG>(
617            &mut builder,
618            key_target,
619            &nonce_target,
620            &pt_target,
621        );
622
623        println!(
624            "encrypt (NK:{}, NB:{}, NR:{}, L:{}, TAG:{}) num_gates: {}",
625            NK,
626            NB,
627            NR,
628            L,
629            TAG,
630            builder.num_gates()
631        );
632        let data = builder.build::<PoseidonGoldilocksConfig>();
633
634        if only_build {
635            return Ok(());
636        }
637
638        // set values to circuit
639        let mut pw = PartialWitness::<F>::new();
640        std::iter::zip(key_target, key).try_for_each(|(t, v)| pw.set_byte_array_target(t, *v))?;
641        std::iter::zip(nonce_target, nonce)
642            .try_for_each(|(t, v)| pw.set_byte_array_target(t, *v))?;
643        std::iter::zip(pt_target, pt).try_for_each(|(t, v)| pw.set_byte_array_target(t, *v))?;
644
645        // expect circuit's ciphertext to match rust native ciphertext
646        std::iter::zip(circuit_c, expected_c)
647            .try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?;
648
649        // if TAG==true, check computed tag, else expect it to be zero
650        if TAG {
651            std::iter::zip(circuit_tag, expected_tag)
652                .try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?;
653        } else {
654            circuit_tag
655                .iter()
656                .try_for_each(|t| pw.set_byte_array_target(*t, 0))?;
657        }
658
659        let proof = data.prove(pw)?;
660        data.verify(proof)
661    }
662}