chore(aes): optimise circuits (#9) * Factor out byte XOR * WIP * include u8::max in the lookups generation (avoiding missing generators) * Complete optimisations * Add range checks --------- Co-authored-by: arnaucube <git@arnaucube.com>
Ahmad Afuni root@ahmadafuni.com
Wed, 25 Mar 2026 16:28:22 +1000
3 files changed,
519 insertions(+),
424 deletions(-)
M
aes-gcm/src/circuit_aes.rs
→
aes-gcm/src/circuit_aes.rs
@@ -9,7 +9,7 @@ use plonky2::{
field::{extension::Extendable, goldilocks_field::GoldilocksField as F, types::Field}, hash::hash_types::RichField, iop::{ - target::{BoolTarget, Target}, + target::Target, witness::{PartialWitness, WitnessWrite}, }, plonk::circuit_builder::CircuitBuilder,@@ -18,7 +18,7 @@
use crate::{ D, constants::{RCON, SBOX}, - native_aes::{State, rot_word, shift_rows}, + native_aes::{State, gf_2_8_mul, rot_word, shift_rows}, }; #[derive(Debug, Copy, Clone)]@@ -34,24 +34,73 @@ StateTarget(array::from_fn(|i| array::from_fn(|j| b[j * 4 + i])))
} } -pub type ByteTarget = [BoolTarget; 8]; +/// Type for byte targets +#[derive(Debug, Copy, Clone)] +pub struct ByteTarget(pub(crate) Target); pub trait CircuitBuilderAESState<F: RichField + Extendable<D>, const D: usize> { - /// Adds state target. - fn add_virtual_state(&mut self) -> StateTarget; + /// Adds byte target (unsafe). + fn add_virtual_byte_target_unsafe(&mut self) -> ByteTarget; + + /// Adds byte target safely assuming an appropriate lookup table + /// (see `assert_byte`). + fn add_virtual_byte_target(&mut self, u8_table_idx: usize) -> ByteTarget { + let byte_target = self.add_virtual_byte_target_unsafe(); + self.assert_byte(byte_target.0, u8_table_idx); + byte_target + } + + /// Asserts that a given target is a byte. + /// Requires a lookup table of the form {(x, y): x ∈ u8}. + fn assert_byte(&mut self, x: Target, u8_table_idx: usize); + + /// Adds byte constant + fn byte_constant(&mut self, c: u8) -> ByteTarget; + + /// Adds state target (unsafe). + fn add_virtual_state_target_unsafe(&mut self) -> StateTarget { + StateTarget(array::from_fn(|_| { + array::from_fn(|_| self.add_virtual_byte_target_unsafe()) + })) + } + + /// Adds state target under the same assumptions as + /// `add_virtual_byte_target`. + fn add_virtual_state_target(&mut self, u8_table_idx: usize) -> StateTarget { + StateTarget(array::from_fn(|_| { + array::from_fn(|_| self.add_virtual_byte_target(u8_table_idx)) + })) + } /// AES cipher as in spec. fn encrypt_block<const NR: usize>( &mut self, + xor_lut_idx: usize, + gf_2_8_mul_lut_idx: usize, sbox_lut_idx: usize, - mix_matrix: [[ByteTarget; 4]; 4], + mix_matrix: [[ByteTarget; 4]; 4], // constant mix_matrix s: StateTarget, - - w: [[ByteTarget; 4]; 4 * (NR + 1)], - ) -> StateTarget; + w: [[ByteTarget; 4]; 4 * (NR + 1)], // expanded key + ) -> StateTarget { + let mut s = s; + s = self.state_add_round_key(xor_lut_idx, &w[0..4], s); + (1..NR).for_each(|i| { + s = self.state_sub_bytes(sbox_lut_idx, s); + s = StateTarget(shift_rows(s.0)); + s = self.state_mix_columns(xor_lut_idx, gf_2_8_mul_lut_idx, mix_matrix, s); + s = self.state_add_round_key(xor_lut_idx, &w[4 * i..4 * (i + 1)], s); + }); + s = self.state_sub_bytes(sbox_lut_idx, s); + s = StateTarget(shift_rows(s.0)); + self.state_add_round_key(xor_lut_idx, &w[4 * NR..4 * (NR + 1)], s) + } /// Applies sub_bytes routine to state. - fn state_sub_bytes(&mut self, sbox_lut_idx: usize, s: StateTarget) -> StateTarget; + fn state_sub_bytes(&mut self, sbox_lut_idx: usize, s: StateTarget) -> StateTarget { + StateTarget(array::from_fn(|i| { + self.state_sub_word(sbox_lut_idx, s.0[i]) + })) + } /// SubWord fn state_sub_word(&mut self, sbox_lut_idx: usize, word: [ByteTarget; 4]) -> [ByteTarget; 4];@@ -59,50 +108,64 @@
/// MixColumns fn state_mix_columns( &mut self, + xor_lut_idx: usize, + gf_2_8_mul_lut_idx: usize, mix_matrix: [[ByteTarget; 4]; 4], s: StateTarget, - ) -> StateTarget; + ) -> StateTarget { + let cols: [_; 4] = array::from_fn(|i| array::from_fn(|j| s.0[j][i])); + let out_cols: [_; 4] = array::from_fn(|i| { + self.byte_matrix_apply(xor_lut_idx, gf_2_8_mul_lut_idx, mix_matrix, cols[i]) + }); + StateTarget(array::from_fn(|i| array::from_fn(|j| out_cols[j][i]))) + } /// AddRoundKey - fn state_add_round_key(&mut self, round_key: &[[ByteTarget; 4]], s: StateTarget) - -> StateTarget; + fn state_add_round_key( + &mut self, + xor_lut_idx: usize, + round_key: &[[ByteTarget; 4]], + s: StateTarget, + ) -> StateTarget { + StateTarget(array::from_fn(|i| { + array::from_fn(|j| self.gf_2_8_add(xor_lut_idx, s.0[i][j], round_key[j][i])) + })) + } /// KeyExpansion fn key_expansion<const NK: usize, const NB: usize, const NR: usize>( &mut self, + xor_lut_idx: usize, sbox_lut_idx: usize, key: [ByteTarget; NK * NB], ) -> [[ByteTarget; 4]; 4 * (NR + 1)]; /// GF(2^8) addition - fn gf_2_8_add(&mut self, x: Target, y: Target) -> Target; - - /// GF(2^8) addition actinb on bit arrays - fn gf_2_8_add_bits( - &mut self, - x_bits: [BoolTarget; 8], - y_bits: [BoolTarget; 8], - ) -> [BoolTarget; 8]; + fn gf_2_8_add(&mut self, xor_lut_idx: usize, x: ByteTarget, y: ByteTarget) -> ByteTarget; /// GF(2^8) multiplication - fn gf_2_8_mul(&mut self, x: Target, y: Target) -> Target; - - /// GF(2^8) multiplication acting on bit arrays - fn gf_2_8_mul_bits(&mut self, x: [BoolTarget; 8], y: [BoolTarget; 8]) -> [BoolTarget; 8]; + fn gf_2_8_mul(&mut self, gf_2_8_mul_lut_idx: usize, x: ByteTarget, y: ByteTarget) + -> ByteTarget; /// Bytearray inner product. - fn bytearray_ip_bits<const N: usize>( + fn bytearray_ip<const N: usize>( &mut self, - x: [[BoolTarget; 8]; N], - y: [[BoolTarget; 8]; N], - ) -> [BoolTarget; 8]; + xor_lut_idx: usize, + gf_2_8_mul_lut_idx: usize, + x: [ByteTarget; N], + y: [ByteTarget; N], + ) -> ByteTarget; - /// Bytearray matrix application. - fn bytearray_matrix_apply_bits<const M: usize, const N: usize>( + /// Byte matrix application. + fn byte_matrix_apply<const M: usize, const N: usize>( &mut self, + xor_lut_idx: usize, + gf_2_8_mul_lut_idx: usize, a: [[ByteTarget; N]; M], x: [ByteTarget; N], - ) -> [ByteTarget; M]; + ) -> [ByteTarget; M] { + std::array::from_fn(|i| self.bytearray_ip(xor_lut_idx, gf_2_8_mul_lut_idx, a[i], x)) + } /// returns a 0u8 in the shape of a ByteTarget. fn zero_byte(&mut self) -> ByteTarget;@@ -111,76 +174,33 @@ fn empty_state(&mut self) -> StateTarget;
} impl CircuitBuilderAESState<F, D> for CircuitBuilder<F, D> { - fn add_virtual_state(&mut self) -> StateTarget { - StateTarget(array::from_fn(|_| { - array::from_fn(|_| array::from_fn(|_| self.add_virtual_bool_target_safe())) - })) + fn add_virtual_byte_target_unsafe(&mut self) -> ByteTarget { + ByteTarget(self.add_virtual_target()) } - fn encrypt_block<const NR: usize>( - &mut self, - sbox_lut_idx: usize, - mix_matrix: [[ByteTarget; 4]; 4], // constant mix_matrix - s: StateTarget, - w: [[ByteTarget; 4]; 4 * (NR + 1)], // expanded key - ) -> StateTarget { - let mut s = s; - s = self.state_add_round_key(&w[0..4], s); - (1..NR).for_each(|i| { - s = self.state_sub_bytes(sbox_lut_idx, s); - s = StateTarget(shift_rows(s.0)); - s = self.state_mix_columns(mix_matrix, s); - s = self.state_add_round_key(&w[4 * i..4 * (i + 1)], s); - }); - s = self.state_sub_bytes(sbox_lut_idx, s); - s = StateTarget(shift_rows(s.0)); - self.state_add_round_key(&w[4 * NR..4 * (NR + 1)], s) + fn assert_byte(&mut self, x: Target, u8_table_idx: usize) { + self.add_lookup_from_index(x, u8_table_idx); } - fn state_sub_bytes(&mut self, sbox_lut_idx: usize, s: StateTarget) -> StateTarget { - StateTarget(array::from_fn(|i| { - self.state_sub_word(sbox_lut_idx, s.0[i]) - })) + fn byte_constant(&mut self, c: u8) -> ByteTarget { + ByteTarget(self.constant(F::from_canonical_u8(c))) } fn state_sub_word(&mut self, sbox_lut_idx: usize, word: [ByteTarget; 4]) -> [ByteTarget; 4] { array::from_fn(|i| { - let byte_target = target_from_bitarray(self, &word[i]); - let out_target = self.add_lookup_from_index(byte_target, sbox_lut_idx); - bitarray_from_bytetarget(self, out_target) - }) - } - - fn state_mix_columns( - &mut self, - mix_matrix: [[ByteTarget; 4]; 4], - s: StateTarget, - ) -> StateTarget { - let cols: [_; 4] = array::from_fn(|i| array::from_fn(|j| s.0[j][i])); - let out_cols: [_; 4] = - array::from_fn(|i| self.bytearray_matrix_apply_bits(mix_matrix, cols[i])); - StateTarget(array::from_fn(|i| array::from_fn(|j| out_cols[j][i]))) - } + let byte_target = word[i]; - fn state_add_round_key( - &mut self, - round_key: &[[ByteTarget; 4]], - s: StateTarget, - ) -> StateTarget { - StateTarget(array::from_fn(|i| { - array::from_fn(|j| self.gf_2_8_add_bits(s.0[i][j], round_key[j][i])) - })) + ByteTarget(self.add_lookup_from_index(byte_target.0, sbox_lut_idx)) + }) } fn key_expansion<const NK: usize, const NB: usize, const NR: usize>( &mut self, + xor_lut_idx: usize, sbox_lut_idx: usize, key: [ByteTarget; NK * NB], ) -> [[ByteTarget; 4]; 4 * (NR + 1)] { - let rcon: [ByteTarget; 11] = array::from_fn(|i| { - let rcon_bits = le_bits_from_byte(RCON[i]); - array::from_fn(|j| self.constant_bool(rcon_bits[j])) - }); + let rcon: [ByteTarget; 11] = array::from_fn(|i| self.byte_constant(RCON[i])); let key: [[_; 4]; NK] = array::from_fn(|i| array::from_fn(|j| key[4 * i + j]));@@ -190,7 +210,7 @@ let offset = if i % NK == 0 {
let term = self.state_sub_word(sbox_lut_idx, rot_word(st[i - 1])); array::from_fn(|j| { if j == 0 { - self.gf_2_8_add_bits(term[j], rcon[i / NK]) + self.gf_2_8_add(xor_lut_idx, term[j], rcon[i / NK]) } else { term[j] }@@ -201,7 +221,8 @@ } else {
st[i - 1] }; - let cur_val = array::from_fn(|j| self.gf_2_8_add_bits(st[i - NK][j], offset[j])); + let cur_val = + array::from_fn(|j| self.gf_2_8_add(xor_lut_idx, st[i - NK][j], offset[j])); st.push(cur_val); Some(cur_val) })@@ -215,125 +236,62 @@ }
}) } - fn gf_2_8_add(&mut self, x: Target, y: Target) -> Target { - let x_bits = bitarray_from_bytetarget(self, x); - let y_bits = bitarray_from_bytetarget(self, y); - - let sum_bits = self.gf_2_8_add_bits(x_bits, y_bits); - target_from_bitarray(self, &sum_bits) - } - - fn gf_2_8_add_bits( - &mut self, - x_bits: [BoolTarget; 8], - y_bits: [BoolTarget; 8], - ) -> [BoolTarget; 8] { - let sum_bits = std::iter::zip(x_bits, y_bits) - .map(|(a, b)| xor(self, a, b)) - .collect::<Vec<_>>(); - - array::from_fn(|i| sum_bits[i]) - } - - fn gf_2_8_mul(&mut self, x: Target, y: Target) -> Target { - let x_bits = bitarray_from_bytetarget(self, x); - let y_bits = bitarray_from_bytetarget(self, y); - - let prod_bits = self.gf_2_8_mul_bits(x_bits, y_bits); - target_from_bitarray(self, &prod_bits) - } - - fn gf_2_8_mul_bits( - &mut self, - x_bits: [BoolTarget; 8], - y_bits: [BoolTarget; 8], - ) -> [BoolTarget; 8] { - let zero_bits = vec![self._false(); 8]; - assert!(x_bits.len() == 8 && y_bits.len() == 8); - - let x_times = |builder: &mut CircuitBuilder<_, _>, y: &[BoolTarget]| -> [_; 8] { - let zero = builder._false(); - let one = builder._true(); - - let potential_offset = [one, one, zero, one, one, zero, zero, zero]; - let y_shifted: [_; 8] = array::from_fn(|i| if i == 0 { zero } else { y[i - 1] }); - - let offset: [_; 8] = array::from_fn(|i| builder.and(y[7], potential_offset[i])); - array::from_fn(|i| xor(builder, y_shifted[i], offset[i])) - }; - - let powers_times_x = (0..8) - .scan(x_bits.to_vec(), |st, i| { - if i > 0 { - *st = x_times(self, st).to_vec(); - } - Some(st.clone()) - }) - .collect::<Vec<_>>(); - - let prod_bits = std::iter::zip(y_bits, powers_times_x) - .map(|(c, o)| o.into_iter().map(|b| self.and(c, b)).collect::<Vec<_>>()) - .collect::<Vec<_>>() - .into_iter() - .fold(zero_bits, |acc, term| { - std::iter::zip(acc, term) - .map(|(a, b)| xor(self, a, b)) - .collect::<Vec<_>>() - }); - - array::from_fn(|i| prod_bits[i]) + fn gf_2_8_add(&mut self, xor_lut_idx: usize, x: ByteTarget, y: ByteTarget) -> ByteTarget { + byte_xor(self, xor_lut_idx, x, y) } - fn bytearray_ip_bits<const N: usize>( + fn gf_2_8_mul( &mut self, - x: [[BoolTarget; 8]; N], - y: [[BoolTarget; 8]; N], - ) -> [BoolTarget; 8] { - let zero = self.zero(); - let zero_bits = bitarray_from_bytetarget(self, zero); - - std::iter::zip(x, y).fold(zero_bits, |acc, (a, b)| { - let prod = self.gf_2_8_mul_bits(a, b); - self.gf_2_8_add_bits(acc, prod) - }) + gf_2_8_mul_lut_idx: usize, + x: ByteTarget, + y: ByteTarget, + ) -> ByteTarget { + let lookup_idx = self.mul_const_add(F::from_canonical_u64(1 << 8), x.0, y.0); + ByteTarget(self.add_lookup_from_index(lookup_idx, gf_2_8_mul_lut_idx)) } - fn bytearray_matrix_apply_bits<const M: usize, const N: usize>( + fn bytearray_ip<const N: usize>( &mut self, - a: [[ByteTarget; N]; M], + xor_lut_idx: usize, + gf_2_8_mul_lut_idx: usize, x: [ByteTarget; N], - ) -> [ByteTarget; M] { - std::array::from_fn(|i| self.bytearray_ip_bits(a[i], x)) + y: [ByteTarget; N], + ) -> ByteTarget { + let zero = ByteTarget(self.zero()); + std::iter::zip(x, y).fold(zero, |acc, (a, b)| { + let prod = self.gf_2_8_mul(gf_2_8_mul_lut_idx, a, b); + self.gf_2_8_add(xor_lut_idx, acc, prod) + }) } fn zero_byte(&mut self) -> ByteTarget { - array::from_fn(|_| self._false()) + ByteTarget(self.zero()) } + fn empty_state(&mut self) -> StateTarget { - let zero_byte: [BoolTarget; 8] = array::from_fn(|_| self._false()); + let zero_byte: ByteTarget = self.zero_byte(); StateTarget(array::from_fn(|_| array::from_fn(|_| zero_byte))) } } pub trait PartialWitnessByteArray { - fn set_byte_array_target(&mut self, target: ByteTarget, value: u8) -> anyhow::Result<()>; + fn set_byte_target(&mut self, target: ByteTarget, value: u8) -> anyhow::Result<()>; } impl<F: Field> PartialWitnessByteArray for PartialWitness<F> { - fn set_byte_array_target(&mut self, target: ByteTarget, value: u8) -> anyhow::Result<()> { - let value_bits = le_bits_from_byte(value); - std::iter::zip(target, value_bits).try_for_each(|(t, v)| self.set_bool_target(t, v)) + fn set_byte_target(&mut self, target: ByteTarget, value: u8) -> anyhow::Result<()> { + self.set_target(target.0, F::from_canonical_u8(value)) } } pub trait PartialWitnessAESState { - fn set_target_state(&mut self, target: StateTarget, value: State) -> anyhow::Result<()>; + fn set_state_target(&mut self, target: StateTarget, value: State) -> anyhow::Result<()>; } impl<F: Field> PartialWitnessAESState for PartialWitness<F> { - fn set_target_state(&mut self, target: StateTarget, value: State) -> anyhow::Result<()> { + fn set_state_target(&mut self, target: StateTarget, value: State) -> anyhow::Result<()> { std::iter::zip(target.0, value).try_for_each(|(t, v)| { - std::iter::zip(t, v).try_for_each(|(t, v)| self.set_byte_array_target(t, v)) + std::iter::zip(t, v).try_for_each(|(t, v)| self.set_byte_target(t, v)) }) } }@@ -347,55 +305,56 @@ .collect::<Vec<_>>(),
)) } -pub fn state_mix_matrix_bits(builder: &mut CircuitBuilder<F, D>) -> [[[BoolTarget; 8]; 4]; 4] { - let one = builder.one(); - let two = builder.two(); - let three = builder.constant(F::from_canonical_u64(3)); +/// Lookup table for GF(2^8) multiplication, a.k.a. XOR on bytes. +pub fn byte_xor_lut(builder: &mut CircuitBuilder<F, D>) -> usize { + let xor_table: Vec<(u16, u16)> = (0..=u8::MAX as usize) + .flat_map(|x| { + (0..=u8::MAX as usize) + .map(|y| (((x as u16) << 8) + y as u16, (x as u16) ^ (y as u16))) + .collect::<Vec<_>>() + }) + .collect(); + builder.add_lookup_table_from_pairs(Arc::new(xor_table)) +} + +/// Lookup table for GF(2^8) multiplication. +pub fn gf_2_8_mul_lut(builder: &mut CircuitBuilder<F, D>) -> usize { + let gf_2_8_mul_table: Vec<(u16, u16)> = (0..=u8::MAX as usize) + .flat_map(|x| { + (0..=u8::MAX as usize) + .map(|y| { + ( + ((x as u16) << 8) + y as u16, + gf_2_8_mul(x as u8, y as u8) as u16, + ) + }) + .collect::<Vec<_>>() + }) + .collect(); + builder.add_lookup_table_from_pairs(Arc::new(gf_2_8_mul_table)) +} + +pub fn state_mix_matrix(builder: &mut CircuitBuilder<F, D>) -> [[ByteTarget; 4]; 4] { + let one = builder.byte_constant(1); + let two = builder.byte_constant(2); + let three = builder.byte_constant(3); - let matrix = [ + [ [two, three, one, one], [one, two, three, one], [one, one, two, three], [three, one, one, two], - ]; - - array::from_fn(|i| array::from_fn(|j| bitarray_from_bytetarget(builder, matrix[i][j]))) -} - -pub fn xor(builder: &mut CircuitBuilder<F, D>, x: BoolTarget, y: BoolTarget) -> BoolTarget { - let x_or_y = builder.or(x, y); - BoolTarget::new_unsafe(builder.arithmetic( - F::NEG_ONE, - F::ONE, - x.target, - y.target, - x_or_y.target, - )) -} - -// LE bit array conversions -pub fn target_from_bitarray(builder: &mut CircuitBuilder<F, D>, x: &[BoolTarget]) -> Target { - let zero = builder.zero(); - let two = builder.two(); - x.iter() - .rev() - .fold(zero, |acc, b| builder.mul_add(two, acc, b.target)) + ] } -pub fn bitarray_from_bytetarget(builder: &mut CircuitBuilder<F, D>, x: Target) -> [BoolTarget; 8] { - let x_bits = builder.split_le(x, 8); - array::from_fn(|i| x_bits[i]) -} - -pub fn le_bits_from_byte(v: u8) -> [bool; 8] { - let v_bits = (0..8) - .scan(v, |st, _| { - let cur_bit = (*st % 2) != 0; - *st >>= 1; - Some(cur_bit) - }) - .collect::<Vec<_>>(); - array::from_fn(|i| v_bits[i]) +pub fn byte_xor( + builder: &mut CircuitBuilder<F, D>, + xor_lut_idx: usize, + x: ByteTarget, + y: ByteTarget, +) -> ByteTarget { + let lookup_idx = builder.mul_const_add(F::from_canonical_u64(1 << 8), x.0, y.0); + ByteTarget(builder.add_lookup_from_index(lookup_idx, xor_lut_idx)) } #[cfg(test)]@@ -404,7 +363,10 @@ use std::array;
use anyhow::Result; use plonky2::{ - field::{goldilocks_field::GoldilocksField as F, types::Field}, + field::{ + goldilocks_field::GoldilocksField as F, + types::{Field, Field64}, + }, iop::witness::{PartialWitness, WitnessWrite}, plonk::{ circuit_builder::CircuitBuilder, circuit_data::CircuitConfig,@@ -415,17 +377,46 @@ use rand::RngExt;
use super::{ ByteTarget, CircuitBuilderAESState, D, PartialWitnessAESState, PartialWitnessByteArray, - sbox_lut, state_mix_matrix_bits, + byte_xor_lut, gf_2_8_mul_lut, sbox_lut, state_mix_matrix, }; use crate::native_aes::{State, encrypt_block, key_expansion, mix_columns, sub_bytes}; #[test] + fn test_assert_byte() -> Result<()> { + // Circuit declaration + let config = CircuitConfig::standard_recursion_config(); + let mut builder = CircuitBuilder::<F, D>::new(config); + + let u8_table_idx = sbox_lut(&mut builder); + + let byte_target = builder.add_virtual_byte_target(u8_table_idx); + + let data = builder.build::<PoseidonGoldilocksConfig>(); + + let mut rng = rand::rng(); + let test_values: [u64; 10] = array::from_fn(|_| rng.random_range(0..F::ORDER)); + + test_values.into_iter().try_for_each(|tv| { + let mut pw = PartialWitness::<F>::new(); + pw.set_target(byte_target.0, F::from_canonical_u64(tv))?; + + if tv > u8::MAX as u64 { + assert!(data.prove(pw).is_err()); + Ok(()) + } else { + let proof = data.prove(pw)?; + data.verify(proof) + } + }) + } + + #[test] fn test_sub_bytes() -> Result<()> { // Circuit declaration let config = CircuitConfig::standard_recursion_config(); let mut builder = CircuitBuilder::<F, D>::new(config); - let state_target = builder.add_virtual_state(); + let state_target = builder.add_virtual_state_target_unsafe(); let sbox_lut = sbox_lut(&mut builder); let out_state_target = builder.state_sub_bytes(sbox_lut, state_target);@@ -439,8 +430,8 @@ test_states.into_iter().try_for_each(|s| {
let expected_result = sub_bytes(s); let mut pw = PartialWitness::<F>::new(); - pw.set_target_state(state_target, s)?; - pw.set_target_state(out_state_target, expected_result)?; + pw.set_state_target(state_target, s)?; + pw.set_state_target(out_state_target, expected_result)?; let proof = data.prove(pw)?; data.verify(proof)@@ -453,9 +444,12 @@ // Circuit declaration
let config = CircuitConfig::standard_recursion_config(); let mut builder = CircuitBuilder::<F, D>::new(config); - let state_target = builder.add_virtual_state(); - let mix_matrix = state_mix_matrix_bits(&mut builder); - let out_state_target = builder.state_mix_columns(mix_matrix, state_target); + let xor_lut_idx = byte_xor_lut(&mut builder); + let gf_2_8_mul_lut_idx = gf_2_8_mul_lut(&mut builder); + let state_target = builder.add_virtual_state_target_unsafe(); + let mix_matrix = state_mix_matrix(&mut builder); + let out_state_target = + builder.state_mix_columns(xor_lut_idx, gf_2_8_mul_lut_idx, mix_matrix, state_target); let data = builder.build::<PoseidonGoldilocksConfig>();@@ -467,8 +461,8 @@ test_states.into_iter().try_for_each(|s| {
let expected_result = mix_columns(s); let mut pw = PartialWitness::<F>::new(); - pw.set_target_state(state_target, s)?; - pw.set_target_state(out_state_target, expected_result)?; + pw.set_state_target(state_target, s)?; + pw.set_state_target(out_state_target, expected_result)?; let proof = data.prove(pw)?; data.verify(proof)@@ -481,10 +475,12 @@ // Circuit declaration
let config = CircuitConfig::standard_recursion_config(); let mut builder = CircuitBuilder::<F, D>::new(config); - let a = builder.add_virtual_target(); - let b = builder.add_virtual_target(); + let gf_2_8_mul_lut_idx = gf_2_8_mul_lut(&mut builder); - let a_times_b = builder.gf_2_8_mul(a, b); + let a = builder.add_virtual_byte_target_unsafe(); + let b = builder.add_virtual_byte_target_unsafe(); + + let a_times_b = builder.gf_2_8_mul(gf_2_8_mul_lut_idx, a, b); let data = builder.build::<PoseidonGoldilocksConfig>();@@ -502,18 +498,46 @@ (0x57, 0x13, 0xfe),
]; test_values .into_iter() - .map(|(a, b, c)| { - ( - F::from_canonical_u8(a), - F::from_canonical_u8(b), - F::from_canonical_u8(c), - ) + .try_for_each(|(a_val, b_val, expected_result)| { + let mut pw = PartialWitness::<F>::new(); + pw.set_byte_target(a, a_val)?; + pw.set_byte_target(b, b_val)?; + pw.set_byte_target(a_times_b, expected_result)?; + + let proof = data.prove(pw)?; + data.verify(proof) }) + } + + #[test] + fn test_gf_2_8_add() -> Result<()> { + // Circuit declaration + let config = CircuitConfig::standard_recursion_config(); + let mut builder = CircuitBuilder::<F, D>::new(config); + + let xor_lut_idx = byte_xor_lut(&mut builder); + + let a = builder.add_virtual_byte_target_unsafe(); + let b = builder.add_virtual_byte_target_unsafe(); + + let a_plus_b = builder.gf_2_8_add(xor_lut_idx, a, b); + + let data = builder.build::<PoseidonGoldilocksConfig>(); + + // set values to circuit + let mut rng = rand::rng(); + let test_values: [(u8, u8, u8); 20] = std::array::from_fn(|_| { + let x = rng.random(); + let y = rng.random(); + (x, y, x ^ y) + }); + test_values + .into_iter() .try_for_each(|(a_val, b_val, expected_result)| { let mut pw = PartialWitness::<F>::new(); - pw.set_target(a, a_val)?; - pw.set_target(b, b_val)?; - pw.set_target(a_times_b, expected_result)?; + pw.set_byte_target(a, a_val)?; + pw.set_byte_target(b, b_val)?; + pw.set_byte_target(a_plus_b, expected_result)?; let proof = data.prove(pw)?; data.verify(proof)@@ -560,9 +584,11 @@ let config = CircuitConfig::standard_recursion_config();
let mut builder = CircuitBuilder::<F, D>::new(config); let key_target: [ByteTarget; NK * NB] = - array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe())); - let sbox_lut = sbox_lut(&mut builder); - let expanded_key_target = builder.key_expansion::<NK, NB, NR>(sbox_lut, key_target); + array::from_fn(|_| builder.add_virtual_byte_target_unsafe()); + let xor_lut_idx = byte_xor_lut(&mut builder); + let sbox_lut_idx = sbox_lut(&mut builder); + let expanded_key_target = + builder.key_expansion::<NK, NB, NR>(xor_lut_idx, sbox_lut_idx, key_target); println!( "key_expansion(NK:{}, NB:{}, NR:{}) num_gates: {}",@@ -578,10 +604,10 @@ let w = key_expansion::<NK, NB, NR>(&key);
let mut pw = PartialWitness::<F>::new(); - std::iter::zip(key_target, key).try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?; + std::iter::zip(key_target, key).try_for_each(|(t, v)| pw.set_byte_target(t, v))?; std::iter::zip(expanded_key_target, w).try_for_each(|(t, v)| { - std::iter::zip(t, v).try_for_each(|(t, v)| pw.set_byte_array_target(t, v)) + std::iter::zip(t, v).try_for_each(|(t, v)| pw.set_byte_target(t, v)) })?; let proof = data.prove(pw)?;@@ -642,16 +668,20 @@ let config = CircuitConfig::standard_recursion_config();
let mut builder = CircuitBuilder::<F, D>::new(config); let key_target: [ByteTarget; NK * NB] = - array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe())); - let sbox_lut = sbox_lut(&mut builder); - let mix_matrix = state_mix_matrix_bits(&mut builder); + array::from_fn(|_| builder.add_virtual_byte_target_unsafe()); + let xor_lut_idx = byte_xor_lut(&mut builder); + let gf_2_8_mul_lut_idx = gf_2_8_mul_lut(&mut builder); + let sbox_lut_idx = sbox_lut(&mut builder); + let mix_matrix = state_mix_matrix(&mut builder); let expanded_key_target: [[ByteTarget; 4]; 4 * (NR + 1)] = - builder.key_expansion::<NK, NB, NR>(sbox_lut, key_target); + builder.key_expansion::<NK, NB, NR>(xor_lut_idx, sbox_lut_idx, key_target); - let input_state_target = builder.add_virtual_state(); + let input_state_target = builder.add_virtual_state_target(sbox_lut_idx); let output = builder.encrypt_block( - sbox_lut, + xor_lut_idx, + gf_2_8_mul_lut_idx, + sbox_lut_idx, mix_matrix, input_state_target, expanded_key_target,@@ -679,16 +709,16 @@ let native_ciphertext = encrypt_block::<NR>(&input_state, &expanded_key);
let mut pw = PartialWitness::<F>::new(); - std::iter::zip(key_target, key).try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?; + std::iter::zip(key_target, key).try_for_each(|(t, v)| pw.set_byte_target(t, v))?; - pw.set_target_state(input_state_target, input_state_matrix)?; + pw.set_state_target(input_state_target, input_state_matrix)?; std::iter::zip(expanded_key_target, expanded_key).try_for_each(|(t, v)| { - std::iter::zip(t, v).try_for_each(|(t, v)| pw.set_byte_array_target(t, v)) + std::iter::zip(t, v).try_for_each(|(t, v)| pw.set_byte_target(t, v)) })?; std::iter::zip(output.0, native_ciphertext).try_for_each(|(o, e)| { - std::iter::zip(o, e).try_for_each(|(o, e)| pw.set_byte_array_target(o, e)) + std::iter::zip(o, e).try_for_each(|(o, e)| pw.set_byte_target(o, e)) })?; let proof = data.prove(pw)?;
M
aes-gcm/src/circuit_gcm.rs
→
aes-gcm/src/circuit_gcm.rs
@@ -3,7 +3,7 @@ //!
//! Plonky2 circuit implementation of //! [AES-GCM (Galois Counter Mode)](https://nvlpubs.nist.gov/nistpubs/Legacy/SP/nistspecialpublication800-38d.pdf). -use std::array; +use std::{array, sync::Arc}; use anyhow::Result; use plonky2::{@@ -15,8 +15,8 @@
use crate::{ D, circuit_aes::{ - ByteTarget, CircuitBuilderAESState, PartialWitnessByteArray, StateTarget, - le_bits_from_byte, sbox_lut, state_mix_matrix_bits, xor, + ByteTarget, CircuitBuilderAESState, PartialWitnessByteArray, StateTarget, byte_xor, + byte_xor_lut, gf_2_8_mul_lut, sbox_lut, state_mix_matrix, }, constants::TAG_LEN, };@@ -47,41 +47,43 @@ [(); NK * NB]:,
[(); 4 * (NR + 1)]:, { pub fn build(builder: &mut CircuitBuilder<F, D>) -> Self { + // Declare LUTs + let sbox_lut = sbox_lut(builder); + let xor_lut = byte_xor_lut(builder); + let gf_2_8_mul_lut = gf_2_8_mul_lut(builder); + // add targets let key: [ByteTarget; NK * NB] = - array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe())); - let nonce: [ByteTarget; 12] = - array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe())); - let pt: [ByteTarget; L] = - array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe())); + array::from_fn(|_| builder.add_virtual_byte_target(sbox_lut)); + let nonce: [ByteTarget; 12] = array::from_fn(|_| builder.add_virtual_byte_target(sbox_lut)); + let pt: [ByteTarget; L] = array::from_fn(|_| builder.add_virtual_byte_target(sbox_lut)); let tag: [ByteTarget; TAG_LEN / 8] = - array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe())); + array::from_fn(|_| builder.add_virtual_byte_target(sbox_lut)); - let sbox_lut = sbox_lut(builder); - let mix_matrix = state_mix_matrix_bits(builder); + let mix_matrix = state_mix_matrix(builder); let a: &[ByteTarget] = &[]; // additional authenticated data - let expanded_key = builder.key_expansion::<NK, NB, NR>(sbox_lut, key); + let expanded_key = builder.key_expansion::<NK, NB, NR>(xor_lut, sbox_lut, key); let empty_state = builder.empty_state(); // 1. CIPH_K let h = builder - .encrypt_block::<NR>(sbox_lut, mix_matrix, empty_state, expanded_key) + .encrypt_block::<NR>( + xor_lut, + gf_2_8_mul_lut, + sbox_lut, + mix_matrix, + empty_state, + expanded_key, + ) .flatten(); // 2. J_0 let zero_byte = builder.zero_byte(); - let one_byte = array::from_fn(|i| { - // 0x01 - if i == 0 { - builder._true() - } else { - builder._false() - } - }); + let one_byte = builder.byte_constant(1); let j0: [ByteTarget; 16] = if nonce.len() * 8 == 96 { // J_0 = IV || 0^31 || 1 let mut out = [zero_byte; 16];@@ -94,7 +96,14 @@ };
// 3. C=GCTR() let inc32_j0 = inc32_target(builder, j0); - let ct = gctr_target::<NR, L>(builder, sbox_lut, mix_matrix, expanded_key, &inc32_j0, &pt); + let ct = gctr_target::<NR, L>( + builder, + (xor_lut, gf_2_8_mul_lut, sbox_lut), + mix_matrix, + expanded_key, + &inc32_j0, + &pt, + ); // the rest of the logic is for the tag; for some use cases we might skip it // (saving a notable amount of gates, about double)@@ -115,22 +124,10 @@
// 5. S = GHASH() let const_a_len_u8: [u8; 8] = (a.len() * 8).to_be_bytes(); // const let const_c_len_u8: [u8; 8] = (ct.len() * 8).to_be_bytes(); // const - let a_len: [ByteTarget; 8] = array::from_fn(|byte_i| { - let const_bits = le_bits_from_byte(const_a_len_u8[byte_i]); - let byte: ByteTarget = array::from_fn(|bit_i| { - let bit_goldilocks = if const_bits[bit_i] { F::ONE } else { F::ZERO }; - BoolTarget::new_unsafe(builder.constant(bit_goldilocks)) - }); - byte - }); - let c_len: [ByteTarget; 8] = array::from_fn(|byte_i| { - let const_bits = le_bits_from_byte(const_c_len_u8[byte_i]); - let byte: ByteTarget = array::from_fn(|bit_i| { - let bit_goldilocks = if const_bits[bit_i] { F::ONE } else { F::ZERO }; - BoolTarget::new_unsafe(builder.constant(bit_goldilocks)) - }); - byte - }); + let a_len: [ByteTarget; 8] = + array::from_fn(|byte_i| builder.byte_constant(const_a_len_u8[byte_i])); + let c_len: [ByteTarget; 8] = + array::from_fn(|byte_i| builder.byte_constant(const_c_len_u8[byte_i])); let ghash_input_vec = [ a.to_vec(), vec![zero_byte; v],@@ -140,18 +137,30 @@ a_len.to_vec(),
c_len.to_vec(), ] .concat(); - let s = ghash_target(builder, h, ghash_input_vec); + + let u8_unit_right_shift_lut_idx = u8_unit_right_shift_lut(builder); + let u8_bitref_lut_idx = u8_bitref_lut(builder); + let s = ghash_target( + builder, + xor_lut, + u8_unit_right_shift_lut_idx, + u8_bitref_lut_idx, + h, + ghash_input_vec, + ); // 6. T=MSB(GCTR())) - let msb_input = gctr_target(builder, sbox_lut, mix_matrix, expanded_key, &j0, &s); - let t_vec = msb_t_target(builder, TAG_LEN / 8, &msb_input); + let msb_input = gctr_target( + builder, + (xor_lut, gf_2_8_mul_lut, sbox_lut), + mix_matrix, + expanded_key, + &j0, + &s, + ); + let t_vec = msb_t_target(TAG_LEN / 8, &msb_input); - #[allow(clippy::needless_range_loop)] - for byte in 0..TAG_LEN / 8 { - for bit in 0..8 { - builder.connect(tag[byte][bit].target, t_vec[byte][bit].target); - } - } + std::iter::zip(tag, t_vec).for_each(|(a, b)| builder.connect(a.0, b.0)); Self { key,@@ -183,18 +192,17 @@ ct_arr.copy_from_slice(ct);
let mut tag_arr: [u8; TAG_LEN / 8] = [0u8; TAG_LEN / 8]; tag_arr.copy_from_slice(tag); - std::iter::zip(self.key, key).try_for_each(|(t, v)| pw.set_byte_array_target(t, *v))?; - std::iter::zip(self.nonce, nonce).try_for_each(|(t, v)| pw.set_byte_array_target(t, *v))?; - std::iter::zip(self.pt, pt_arr).try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?; + std::iter::zip(self.key, key).try_for_each(|(t, v)| pw.set_byte_target(t, *v))?; + std::iter::zip(self.nonce, nonce).try_for_each(|(t, v)| pw.set_byte_target(t, *v))?; + std::iter::zip(self.pt, pt_arr).try_for_each(|(t, v)| pw.set_byte_target(t, v))?; - std::iter::zip(self.ct, ct_arr).try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?; + std::iter::zip(self.ct, ct_arr).try_for_each(|(t, v)| pw.set_byte_target(t, v))?; if TAG { - std::iter::zip(self.tag, tag_arr) - .try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?; + std::iter::zip(self.tag, tag_arr).try_for_each(|(t, v)| pw.set_byte_target(t, v))?; } else { self.tag .iter() - .try_for_each(|t| pw.set_byte_array_target(*t, 0u8))?; + .try_for_each(|t| pw.set_byte_target(*t, 0u8))?; } Ok(()) }@@ -203,12 +211,13 @@
/// L: max size of plaintext to cipher. fn gctr_target<const NR: usize, const L: usize>( builder: &mut CircuitBuilder<F, D>, - sbox_lut_idx: usize, + lut_indices: (usize, usize, usize), mix_matrix: [[ByteTarget; 4]; 4], key: [[ByteTarget; 4]; 4 * (NR + 1)], icb: &[ByteTarget; 16], x: &[ByteTarget; L], ) -> [ByteTarget; L] { + let (xor_lut_idx, gf_2_8_mul_lut_idx, sbox_lut_idx) = lut_indices; let n = ((L * 8) as f64 / 128_f64).ceil() as usize; let mut y: [ByteTarget; L] = *x;@@ -225,27 +234,36 @@ let mut x_i = [zero_byte; 16];
x_i[..l].copy_from_slice(x_i_raw); let ciph_cb_i = builder - .encrypt_block::<NR>(sbox_lut_idx, mix_matrix, StateTarget::from_flat(cb_i), key) + .encrypt_block::<NR>( + xor_lut_idx, + gf_2_8_mul_lut_idx, + sbox_lut_idx, + mix_matrix, + StateTarget::from_flat(cb_i), + key, + ) .flatten(); let (y_i, n_bytes) = if i < n && x_i_raw.len() == 16 { - (xor_blocks(builder, x_i, ciph_cb_i), 16) + (xor_blocks(builder, xor_lut_idx, x_i, ciph_cb_i), 16) } else { // last chunk, might be smaller than 16 bytes (L%16 bytes) - let msb_res = msb_t_target(builder, L % 16, ciph_cb_i.as_slice()); + let msb_res = msb_t_target(L % 16, ciph_cb_i.as_slice()); let l = msb_res.len().min(16); let mut m: [ByteTarget; 16] = [zero_byte; 16]; m[..l].copy_from_slice(&msb_res); - (xor_blocks(builder, x_i, m), l) + (xor_blocks(builder, xor_lut_idx, x_i, m), l) }; y[i * 16..i * 16 + n_bytes].clone_from_slice(&y_i[..n_bytes]); } y } -// fn ghash_target<const L: usize>( fn ghash_target( builder: &mut CircuitBuilder<F, D>, + xor_lut_idx: usize, + u8_unit_right_shift_lut_idx: usize, + u8_bitref_lut_idx: usize, h: [ByteTarget; 16], x: Vec<ByteTarget>, ) -> [ByteTarget; 16] {@@ -257,71 +275,72 @@ let mut y = [zero_byte; 16]; // (128 bits)
for i in 0..m { let mut xi = [zero_byte; 16]; xi.clone_from_slice(&x[i * 16..i * 16 + 16]); - let y_xi = xor_blocks(builder, y, xi); - y = gf_2_128_mul_target(builder, y_xi, h); + let y_xi = xor_blocks(builder, xor_lut_idx, y, xi); + y = gf_2_128_mul_target( + builder, + xor_lut_idx, + u8_unit_right_shift_lut_idx, + u8_bitref_lut_idx, + y_xi, + h, + ); } y } fn gf_2_128_mul_target( builder: &mut CircuitBuilder<F, D>, + xor_lut_idx: usize, + u8_unit_right_shift_lut_idx: usize, + u8_bitref_lut_idx: usize, x: [ByteTarget; 16], y: [ByteTarget; 16], ) -> [ByteTarget; 16] { - let zero_byte = builder.zero_byte(); - let one_bool = builder._true(); + let zero = builder.zero_byte(); - // R: 10000111 || 0^120 (in little-endian) - let mut r = [zero_byte; 16]; - r[0][0] = one_bool; - r[0][5] = one_bool; - r[0][6] = one_bool; - r[0][7] = one_bool; + // R: 10000111 (= 225) || 0^120 (in little-endian) + let r_first = builder.byte_constant(225); - let mut z = [zero_byte; 16]; + let mut z = [zero; 16]; let mut v = y; for i in 0..128 { let byte_index = i / 8; let bit_index = 7 - (i % 8); - let xi = x[byte_index][bit_index]; + let bit_idx_target = builder.byte_constant(bit_index as u8); + let xi = u8_bitref(builder, u8_bitref_lut_idx, x[byte_index], bit_idx_target); // set z = if xi==1: z^v, else: z for b in 0..16 { - for k in 0..8 { - let z_xor_v = xor(builder, z[b][k], v[b][k]); - z[b][k] = - BoolTarget::new_unsafe(builder.select(xi, z_xor_v.target, z[b][k].target)); - } + let z_xor_v = byte_xor(builder, xor_lut_idx, z[b], v[b]); + z[b] = ByteTarget(builder.select(xi, z_xor_v.0, z[b].0)); } - let lsb = v[15][0]; // (little-endian) - v = right_shift_one_target(builder, &v); + let lsb = u8_bitref(builder, u8_bitref_lut_idx, v[15], zero); // (little-endian) + + v = right_shift_one_target(builder, u8_unit_right_shift_lut_idx, &v); // if lsb==1: v=v^R, else: v - for b in 0..16 { - for k in 0..8 { - let v_xor_r = xor(builder, v[b][k], r[b][k]); - v[b][k] = - BoolTarget::new_unsafe(builder.select(lsb, v_xor_r.target, v[b][k].target)); - } - } + let v_xor_r = byte_xor(builder, xor_lut_idx, v[0], r_first); + v[0] = ByteTarget(builder.select(lsb, v_xor_r.0, v[0].0)); } z } pub fn right_shift_one_target( builder: &mut CircuitBuilder<F, D>, + u8_unit_right_shift_lut_idx: usize, v: &[ByteTarget; 16], ) -> [ByteTarget; 16] { let mut r: [ByteTarget; 16] = *v; - let mut carry = builder._false(); + let mut carry = builder.zero(); for i in 0..16 { let current = v[i]; + let mut shifted = u8_unit_right_shift(builder, u8_unit_right_shift_lut_idx, current); + let next_carry = + builder.mul_const_add(F::from_canonical_u64(2) * F::NEG_ONE, shifted.0, current.0); - let next_carry = current[0]; + shifted = + ByteTarget(builder.mul_const_add(F::from_canonical_u64(1 << 7), carry, shifted.0)); - let mut shifted = [builder._false(); 8]; - shifted[..7].copy_from_slice(¤t[1..(7 + 1)]); - shifted[7] = carry; r[i] = shifted; carry = next_carry; }@@ -333,63 +352,87 @@ builder: &mut CircuitBuilder<F, D>,
block: [ByteTarget; 16], ) -> [ByteTarget; 16] { let mut r = block; + let zero = builder.zero(); + let u8_max = builder.constant(F::from_canonical_u8(u8::MAX)); - let mut carry = BoolTarget::new_unsafe(builder.one()); + let mut carry = builder.one(); for byte_index in (12..16).rev() { - for bit_index in 0..8 { - let a = block[byte_index][bit_index]; - // builder.assert_bool(a); // Note: we can assume it's valid bool, skip check - - let sum = xor(builder, a, carry); - let carry_out = builder.and(a, carry); - r[byte_index][bit_index] = sum; - carry = carry_out; - } + let a = block[byte_index]; + let sum = builder.add(a.0, carry); + let a_is_u8_max = builder.is_equal(a.0, u8_max); + let carry_out = builder.mul(carry, a_is_u8_max.target); + r[byte_index] = ByteTarget(builder.select(a_is_u8_max, zero, sum)); + carry = carry_out; } r } -fn xor_byte(builder: &mut CircuitBuilder<F, D>, b1: ByteTarget, b2: ByteTarget) -> ByteTarget { - array::from_fn(|i| xor(builder, b1[i], b2[i])) -} - fn xor_blocks( builder: &mut CircuitBuilder<F, D>, + xor_lut_idx: usize, b1: [ByteTarget; 16], b2: [ByteTarget; 16], ) -> [ByteTarget; 16] { - array::from_fn(|i| xor_byte(builder, b1[i], b2[i])) + array::from_fn(|i| byte_xor(builder, xor_lut_idx, b1[i], b2[i])) } -pub fn msb_t_target( - builder: &mut CircuitBuilder<F, D>, - t_bytes: usize, - block: &[ByteTarget], -) -> Vec<ByteTarget> { - let t = t_bytes * 8; - - let full_bytes = t / 8; - let rem_bits = t % 8; +pub fn msb_t_target(t_bytes: usize, block: &[ByteTarget]) -> Vec<ByteTarget> { + // Always dealing with full bytes in circuit. + let full_bytes = t_bytes; let mut out: Vec<ByteTarget> = Vec::new(); for block_i in block.iter().take(full_bytes) { out.push(*block_i); } - let zero_bool = BoolTarget::new_unsafe(builder.zero()); - if rem_bits != 0 { - let mut partial = [zero_bool; 8]; + out +} + +/// Lookup table for unit right-shift of elements of u8. +pub fn u8_unit_right_shift_lut(builder: &mut CircuitBuilder<F, D>) -> usize { + let u8_unit_right_shift_table: Vec<(u16, u16)> = (0..=u8::MAX as usize) + .map(|x| (x as u16, (x as u16) >> 1)) + .collect(); + builder.add_lookup_table_from_pairs(Arc::new(u8_unit_right_shift_table)) +} +pub fn u8_unit_right_shift( + builder: &mut CircuitBuilder<F, D>, + u8_unit_right_shift_lut_idx: usize, + x: ByteTarget, +) -> ByteTarget { + ByteTarget(builder.add_lookup_from_index(x.0, u8_unit_right_shift_lut_idx)) +} + +/// Lookup table for u8 bit referencing. +pub fn u8_bitref_lut(builder: &mut CircuitBuilder<F, D>) -> usize { + let u8_bitref_table: Vec<(u16, u16)> = (0..=u8::MAX as usize) + .flat_map(|x| { + (0..8) + .map(|i| (((x as u16) << 3) + i as u16, ((x as u16) >> i) & 1)) + .collect::<Vec<_>>() + }) + .collect(); + builder.add_lookup_table_from_pairs(Arc::new(u8_bitref_table)) +} +pub fn u8_bitref( + builder: &mut CircuitBuilder<F, D>, + u8_bitref_lut_idx: usize, + x: ByteTarget, + i: ByteTarget, +) -> BoolTarget { + let lookup_idx = builder.mul_const_add(F::from_canonical_u64(8), x.0, i.0); + BoolTarget::new_unsafe(builder.add_lookup_from_index(lookup_idx, u8_bitref_lut_idx)) +} - for bit in 0..8 { - if bit < rem_bits { - partial[bit] = block[full_bytes][bit]; - } else { - partial[bit] = zero_bool; - } - } - out.push(partial); - } - out +fn le_bits_from_byte(v: u8) -> [bool; 8] { + let v_bits = (0..8) + .scan(v, |st, _| { + let cur_bit = (*st % 2) != 0; + *st >>= 1; + Some(cur_bit) + }) + .collect::<Vec<_>>(); + array::from_fn(|i| v_bits[i]) } #[cfg(test)]@@ -408,7 +451,7 @@ };
use super::{D, *}; use crate::{ - circuit_aes::PartialWitnessByteArray, + circuit_aes::{PartialWitnessByteArray, byte_xor_lut, gf_2_8_mul_lut}, native_aes::key_expansion, native_gcm::{gctr, gf_2_128_mul, ghash, right_shift_one}, };@@ -443,20 +486,23 @@ let config = CircuitConfig::standard_recursion_config();
let mut builder = CircuitBuilder::<F, D>::new(config); let key_target: [ByteTarget; NK * NB] = - array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe())); + array::from_fn(|_| builder.add_virtual_byte_target_unsafe()); let nonce_target: [ByteTarget; 12] = - array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe())); + array::from_fn(|_| builder.add_virtual_byte_target_unsafe()); let icb_target: [ByteTarget; 16] = - array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe())); + array::from_fn(|_| builder.add_virtual_byte_target_unsafe()); let pt_target: [ByteTarget; L] = - array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe())); + array::from_fn(|_| builder.add_virtual_byte_target_unsafe()); let sbox_lut = sbox_lut(&mut builder); - let mix_matrix = state_mix_matrix_bits(&mut builder); + let xor_lut = byte_xor_lut(&mut builder); + let gf_2_8_mul_lut = gf_2_8_mul_lut(&mut builder); + let mix_matrix = state_mix_matrix(&mut builder); - let key_expanded_target = builder.key_expansion::<NK, NB, NR>(sbox_lut, key_target); + let key_expanded_target = + builder.key_expansion::<NK, NB, NR>(xor_lut, sbox_lut, key_target); let gctr_out_target = gctr_target::<NR, L>( &mut builder, - sbox_lut, + (xor_lut, gf_2_8_mul_lut, sbox_lut), mix_matrix, key_expanded_target, &icb_target,@@ -475,14 +521,13 @@ let data = builder.build::<PoseidonGoldilocksConfig>();
// set values to circuit let mut pw = PartialWitness::<F>::new(); - std::iter::zip(key_target, key).try_for_each(|(t, v)| pw.set_byte_array_target(t, *v))?; - std::iter::zip(nonce_target, nonce) - .try_for_each(|(t, v)| pw.set_byte_array_target(t, *v))?; - std::iter::zip(icb_target, icb).try_for_each(|(t, v)| pw.set_byte_array_target(t, *v))?; - std::iter::zip(pt_target, pt).try_for_each(|(t, v)| pw.set_byte_array_target(t, *v))?; + std::iter::zip(key_target, key).try_for_each(|(t, v)| pw.set_byte_target(t, *v))?; + std::iter::zip(nonce_target, nonce).try_for_each(|(t, v)| pw.set_byte_target(t, *v))?; + std::iter::zip(icb_target, icb).try_for_each(|(t, v)| pw.set_byte_target(t, *v))?; + std::iter::zip(pt_target, pt).try_for_each(|(t, v)| pw.set_byte_target(t, *v))?; std::iter::zip(gctr_out_target, expected) - .try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?; + .try_for_each(|(t, v)| pw.set_byte_target(t, v))?; let proof = data.prove(pw)?; data.verify(proof)@@ -499,20 +544,22 @@ // Circuit declaration
let config = CircuitConfig::standard_recursion_config(); let mut builder = CircuitBuilder::<F, D>::new(config); + let u8_unit_right_shift_lut_idx = u8_unit_right_shift_lut(&mut builder); + let x_target: [ByteTarget; 16] = - array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe())); + array::from_fn(|_| builder.add_virtual_byte_target_unsafe()); - let out_target: [ByteTarget; 16] = right_shift_one_target(&mut builder, &x_target); + let out_target: [ByteTarget; 16] = + right_shift_one_target(&mut builder, u8_unit_right_shift_lut_idx, &x_target); println!("right_shift_one num_gates: {}", builder.num_gates()); let data = builder.build::<PoseidonGoldilocksConfig>(); // set values to circuit let mut pw = PartialWitness::<F>::new(); - std::iter::zip(x_target, x).try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?; + std::iter::zip(x_target, x).try_for_each(|(t, v)| pw.set_byte_target(t, v))?; - std::iter::zip(out_target, expected) - .try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?; + std::iter::zip(out_target, expected).try_for_each(|(t, v)| pw.set_byte_target(t, v))?; let proof = data.prove(pw)?; data.verify(proof)?;@@ -540,38 +587,45 @@ let x: [u8; 16] = [111; 16];
let y: [u8; 16] = [222; 16]; let expected = gf_2_128_mul(x, y); - let x_bits: [[bool; 8]; 16] = - array::from_fn(|b| crate::circuit_aes::le_bits_from_byte(x[b])); - let y_bits: [[bool; 8]; 16] = - array::from_fn(|b| crate::circuit_aes::le_bits_from_byte(y[b])); + let x_bits: [[bool; 8]; 16] = array::from_fn(|b| super::le_bits_from_byte(x[b])); + let y_bits: [[bool; 8]; 16] = array::from_fn(|b| super::le_bits_from_byte(y[b])); // sanity check let expected2 = crate::native_gcm::gf_2_128_mul_circuit_version(x_bits, y_bits); let expected_bits: [[bool; 8]; 16] = - array::from_fn(|b| crate::circuit_aes::le_bits_from_byte(expected[b])); + array::from_fn(|b| super::le_bits_from_byte(expected[b])); assert_eq!(expected2, expected_bits); // Circuit declaration let config = CircuitConfig::standard_recursion_config(); let mut builder = CircuitBuilder::<F, D>::new(config); + let xor_lut_idx = byte_xor_lut(&mut builder); + let u8_unit_right_shift_lut_idx = u8_unit_right_shift_lut(&mut builder); + let u8_bitref_lut_idx = u8_bitref_lut(&mut builder); let x_target: [ByteTarget; 16] = - array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe())); + array::from_fn(|_| builder.add_virtual_byte_target_unsafe()); let y_target: [ByteTarget; 16] = - array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe())); + array::from_fn(|_| builder.add_virtual_byte_target_unsafe()); - let out_target = gf_2_128_mul_target(&mut builder, x_target, y_target); + let out_target = gf_2_128_mul_target( + &mut builder, + xor_lut_idx, + u8_unit_right_shift_lut_idx, + u8_bitref_lut_idx, + x_target, + y_target, + ); println!("gf_2_128_mul (L:{}) num_gates: {}", L, builder.num_gates()); let data = builder.build::<PoseidonGoldilocksConfig>(); // set values to circuit let mut pw = PartialWitness::<F>::new(); - std::iter::zip(x_target, x).try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?; - std::iter::zip(y_target, y).try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?; + std::iter::zip(x_target, x).try_for_each(|(t, v)| pw.set_byte_target(t, v))?; + std::iter::zip(y_target, y).try_for_each(|(t, v)| pw.set_byte_target(t, v))?; - std::iter::zip(out_target, expected) - .try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?; + std::iter::zip(out_target, expected).try_for_each(|(t, v)| pw.set_byte_target(t, v))?; let proof = data.prove(pw)?; data.verify(proof)@@ -603,24 +657,35 @@ // Circuit declaration
let config = CircuitConfig::standard_recursion_config(); let mut builder = CircuitBuilder::<F, D>::new(config); + let xor_lut_idx = byte_xor_lut(&mut builder); + let u8_unit_right_shift_lut_idx = u8_unit_right_shift_lut(&mut builder); + let u8_bitref_lut_idx = u8_bitref_lut(&mut builder); + let h_target: [ByteTarget; 16] = - array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe())); + array::from_fn(|_| builder.add_virtual_byte_target_unsafe()); let x_target: [ByteTarget; L] = - array::from_fn(|_| array::from_fn(|_| builder.add_virtual_bool_target_safe())); + array::from_fn(|_| builder.add_virtual_byte_target_unsafe()); // let ghash_out_target = ghash_target::<L>(&mut builder, h_target, &x_target); - let ghash_out_target = ghash_target(&mut builder, h_target, x_target.to_vec()); + let ghash_out_target = ghash_target( + &mut builder, + xor_lut_idx, + u8_unit_right_shift_lut_idx, + u8_bitref_lut_idx, + h_target, + x_target.to_vec(), + ); println!("ghash (L:{}) num_gates: {}", L, builder.num_gates()); let data = builder.build::<PoseidonGoldilocksConfig>(); // set values to circuit let mut pw = PartialWitness::<F>::new(); - std::iter::zip(h_target, h).try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?; - std::iter::zip(x_target, x).try_for_each(|(t, v)| pw.set_byte_array_target(t, *v))?; + std::iter::zip(h_target, h).try_for_each(|(t, v)| pw.set_byte_target(t, v))?; + std::iter::zip(x_target, x).try_for_each(|(t, v)| pw.set_byte_target(t, *v))?; std::iter::zip(ghash_out_target, expected) - .try_for_each(|(t, v)| pw.set_byte_array_target(t, v))?; + .try_for_each(|(t, v)| pw.set_byte_target(t, v))?; let proof = data.prove(pw)?; data.verify(proof)
M
aes-gcm/src/native_aes.rs
→
aes-gcm/src/native_aes.rs
@@ -79,7 +79,7 @@ r
} /// multiplication in GF(2^8), section 4.2 -fn gf_2_8_mul(a_raw: u8, b_raw: u8) -> u8 { +pub(crate) fn gf_2_8_mul(a_raw: u8, b_raw: u8) -> u8 { let mut r = 0u8; let mut a = a_raw; let mut b = b_raw;