M Cargo.toml → Cargo.toml

@@ -4,5 +4,8 @@ version = "0.1.0"
 edition = "2024"
 
 [dependencies]
+plonky2 = { git = "https://github.com/0xPARC/plonky2.git", rev = "109d517d09c210ae4c2cee381d3e3fbc04aa3812" }
+anyhow = "1.0.56"
 
 [dev-dependencies]
+rand = "0.10.0"

A src/circuit.rs

@@ -0,0 +1,675 @@
+use std::{array, sync::Arc};
+
+use plonky2::{
+    field::{extension::Extendable, goldilocks_field::GoldilocksField as F, types::Field},
+    hash::hash_types::RichField,
+    iop::{
+        target::{BoolTarget, Target},
+        witness::{PartialWitness, WitnessWrite},
+    },
+    plonk::circuit_builder::CircuitBuilder,
+};
+
+use crate::{
+    D,
+    constants::{RCON, SBOX},
+    native::{State, rot_word, shift_rows},
+};
+
+#[derive(Debug, Copy, Clone)]
+pub struct StateTarget(pub [[ByteArrayTarget; 4]; 4]);
+
+// TODO: Normalise terminology!
+pub type ByteArrayTarget = [BoolTarget; 8];
+
+pub trait CircuitBuilderAESState<F: RichField + Extendable<D>, const D: usize> {
+    /// Adds state target.
+    fn add_virtual_state(&mut self) -> StateTarget;
+
+    /// AES cipher as in spec.
+    fn aes_cipher<const NR: usize>(
+        &mut self,
+        sbox_lut_idx: usize,
+        mix_matrix: [[ByteArrayTarget; 4]; 4],
+        s: StateTarget,
+
+        w: [[ByteArrayTarget; 4]; 4 * (NR + 1)],
+    ) -> StateTarget;
+
+    /// Applies sub_bytes routine to state.
+    fn state_sub_bytes(&mut self, sbox_lut_idx: usize, s: StateTarget) -> StateTarget;
+
+    /// SubWord
+    fn state_sub_word(
+        &mut self,
+        sbox_lut_idx: usize,
+        word: [ByteArrayTarget; 4],
+    ) -> [ByteArrayTarget; 4];
+
+    /// MixColumns
+    fn state_mix_columns(
+        &mut self,
+        mix_matrix: [[ByteArrayTarget; 4]; 4],
+        s: StateTarget,
+    ) -> StateTarget;
+
+    /// AddRoundKey
+    fn state_add_round_key(
+        &mut self,
+        round_key: &[[ByteArrayTarget; 4]],
+        s: StateTarget,
+    ) -> StateTarget;
+
+    /// KeyExpansion
+    fn key_expansion<const NK: usize, const NR: usize>(
+        &mut self,
+        sbox_lut_idx: usize,
+        key: [ByteArrayTarget; 4 * NK],
+    ) -> [[ByteArrayTarget; 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];
+
+    /// 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];
+
+    /// Bytearray inner product.
+    fn bytearray_ip_bits<const N: usize>(
+        &mut self,
+        x: [[BoolTarget; 8]; N],
+        y: [[BoolTarget; 8]; N],
+    ) -> [BoolTarget; 8];
+
+    /// Bytearray matrix application.
+    fn bytearray_matrix_apply_bits<const M: usize, const N: usize>(
+        &mut self,
+        a: [[ByteArrayTarget; N]; M],
+        x: [ByteArrayTarget; N],
+    ) -> [ByteArrayTarget; M];
+}
+
+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 aes_cipher<const NR: usize>(
+        &mut self,
+        sbox_lut_idx: usize,
+        mix_matrix: [[ByteArrayTarget; 4]; 4],
+        s: StateTarget,
+        w: [[ByteArrayTarget; 4]; 4 * (NR + 1)],
+    ) -> 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 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 state_sub_word(
+        &mut self,
+        sbox_lut_idx: usize,
+        word: [ByteArrayTarget; 4],
+    ) -> [ByteArrayTarget; 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: [[ByteArrayTarget; 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])))
+    }
+
+    fn state_add_round_key(
+        &mut self,
+        round_key: &[[ByteArrayTarget; 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]))
+        }))
+    }
+
+    fn key_expansion<const NK: usize, const NR: usize>(
+        &mut self,
+        sbox_lut_idx: usize,
+        key: [ByteArrayTarget; 4 * NK],
+    ) -> [[ByteArrayTarget; 4]; 4 * (NR + 1)] {
+        let rcon: [ByteArrayTarget; 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 key: [[_; 4]; NK] = array::from_fn(|i| array::from_fn(|j| key[4 * i + j]));
+
+        let key_additions = (NK..4 * (NR + 1))
+            .scan(key.to_vec(), |st, i| {
+                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])
+                        } else {
+                            term[j]
+                        }
+                    })
+                } else if (NK > 6) && (i % NK == 4) {
+                    self.state_sub_word(sbox_lut_idx, st[i - 1])
+                } else {
+                    st[i - 1]
+                };
+
+                let cur_val = array::from_fn(|j| self.gf_2_8_add_bits(st[i - NK][j], offset[j]));
+                st.push(cur_val);
+                Some(cur_val)
+            })
+            .collect::<Vec<_>>();
+        std::array::from_fn(|i| {
+            if i < NK {
+                key[i]
+            } else {
+                key_additions[i - NK]
+            }
+        })
+    }
+
+    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 bytearray_ip_bits<const N: usize>(
+        &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)
+        })
+    }
+
+    fn bytearray_matrix_apply_bits<const M: usize, const N: usize>(
+        &mut self,
+        a: [[ByteArrayTarget; N]; M],
+        x: [ByteArrayTarget; N],
+    ) -> [ByteArrayTarget; M] {
+        std::array::from_fn(|i| self.bytearray_ip_bits(a[i], x))
+    }
+}
+
+pub trait PartialWitnessByteArray {
+    fn set_byte_array_target(&mut self, target: ByteArrayTarget, value: u8) -> anyhow::Result<()>;
+}
+
+impl<F: Field> PartialWitnessByteArray for PartialWitness<F> {
+    fn set_byte_array_target(&mut self, target: ByteArrayTarget, 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))
+    }
+}
+
+pub trait PartialWitnessAESState {
+    fn set_target_state(&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<()> {
+        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))
+        })
+    }
+}
+
+pub fn sbox_lut(builder: &mut CircuitBuilder<F, D>) -> usize {
+    builder.add_lookup_table_from_pairs(Arc::new(
+        SBOX.into_iter()
+            .enumerate()
+            .map(|(i, o)| (i as u16, o as u16))
+            .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));
+
+    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])
+}
+
+#[cfg(test)]
+mod tests {
+    use std::array;
+
+    use anyhow::Result;
+    use plonky2::{
+        field::{goldilocks_field::GoldilocksField as F, types::Field},
+        iop::witness::{PartialWitness, WitnessWrite},
+        plonk::{
+            circuit_builder::CircuitBuilder, circuit_data::CircuitConfig,
+            config::PoseidonGoldilocksConfig,
+        },
+    };
+    use rand::RngExt;
+
+    use crate::{
+        circuit::{
+            ByteArrayTarget, CircuitBuilderAESState, PartialWitnessAESState,
+            PartialWitnessByteArray, sbox_lut, state_mix_matrix_bits,
+        },
+        native::{State, encrypt_block, key_expansion, mix_columns, sub_bytes},
+    };
+
+    use super::D;
+
+    #[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 sbox_lut = sbox_lut(&mut builder);
+        let out_state_target = builder.state_sub_bytes(sbox_lut, state_target);
+
+        let data = builder.build::<PoseidonGoldilocksConfig>();
+        let mut rng = rand::rng();
+
+        let test_states: [State; 10] = array::from_fn(|_| array::from_fn(|_| rng.random()));
+
+        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)?;
+
+            let proof = data.prove(pw)?;
+            data.verify(proof)
+        })
+    }
+
+    #[test]
+    fn test_mix_columns() -> 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 mix_matrix = state_mix_matrix_bits(&mut builder);
+        let out_state_target = builder.state_mix_columns(mix_matrix, state_target);
+
+        let data = builder.build::<PoseidonGoldilocksConfig>();
+        let mut rng = rand::rng();
+
+        let test_states: [State; 10] = array::from_fn(|_| array::from_fn(|_| rng.random()));
+
+        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)?;
+
+            let proof = data.prove(pw)?;
+            data.verify(proof)
+        })
+    }
+
+    #[test]
+    fn test_gf_2_8_mul() -> Result<()> {
+        // 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 a_times_b = builder.gf_2_8_mul(a, b);
+
+        let data = builder.build::<PoseidonGoldilocksConfig>();
+
+        let test_values = [
+            (0x57, 0x01, 0x57),
+            (0x57, 0x02, 0xae),
+            (0x57, 0x04, 0x47),
+            (0x57, 0x08, 0x8e),
+            (0x57, 0x10, 0x07),
+            (0x57, 0x20, 0x0e),
+            (0x57, 0x40, 0x1c),
+            (0x57, 0x80, 0x38),
+            (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_target(a, a_val)?;
+                pw.set_target(b, b_val)?;
+                pw.set_target(a_times_b, expected_result)?;
+
+                let proof = data.prove(pw)?;
+                data.verify(proof)
+            })
+    }
+
+    /// Analogous tests to `native::tests::test_key_expansion`.
+    #[test]
+    fn test_key_expansion_128() -> Result<()> {
+        // Circuit declaration
+        let config = CircuitConfig::standard_recursion_config();
+        let mut builder = CircuitBuilder::<F, D>::new(config);
+
+        let key_target: [ByteArrayTarget; _] =
+            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::<4, 10>(sbox_lut, key_target);
+
+        let data = builder.build::<PoseidonGoldilocksConfig>();
+
+        // AES-128
+        let key = [
+            0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf,
+            0x4f, 0x3c,
+        ];
+        let w = key_expansion::</*Nk,Nb,Nr*/ 4, 4, 10>(&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(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))
+        })?;
+
+        let proof = data.prove(pw)?;
+        data.verify(proof)
+    }
+
+    #[test]
+    fn test_key_expansion_196() -> Result<()> {
+        // Circuit declaration
+        let config = CircuitConfig::standard_recursion_config();
+        let mut builder = CircuitBuilder::<F, D>::new(config);
+
+        let key_target: [ByteArrayTarget; _] =
+            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::<6, 12>(sbox_lut, key_target);
+
+        let data = builder.build::<PoseidonGoldilocksConfig>();
+
+        // AES-196
+        let key = [
+            0x8e, 0x73, 0xb0, 0xf7, 0xda, 0x0e, 0x64, 0x52, 0xc8, 0x10, 0xf3, 0x2b, 0x80, 0x90,
+            0x79, 0xe5, 0x62, 0xf8, 0xea, 0xd2, 0x52, 0x2c, 0x6b, 0x7b,
+        ];
+        let w = key_expansion::</*Nk,Nb,Nr*/ 6, 4, 12>(&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(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))
+        })?;
+
+        let proof = data.prove(pw)?;
+        data.verify(proof)
+    }
+
+    #[test]
+    fn test_key_expansion_256() -> Result<()> {
+        // Circuit declaration
+        let config = CircuitConfig::standard_recursion_config();
+        let mut builder = CircuitBuilder::<F, D>::new(config);
+
+        let key_target: [ByteArrayTarget; _] =
+            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::<8, 14>(sbox_lut, key_target);
+
+        let data = builder.build::<PoseidonGoldilocksConfig>();
+
+        // AES-256
+        let key = [
+            0x60, 0x3d, 0xeb, 0x10, 0x15, 0xca, 0x71, 0xbe, 0x2b, 0x73, 0xae, 0xf0, 0x85, 0x7d,
+            0x77, 0x81, 0x1f, 0x35, 0x2c, 0x07, 0x3b, 0x61, 0x08, 0xd7, 0x2d, 0x98, 0x10, 0xa3,
+            0x09, 0x14, 0xdf, 0xf4,
+        ];
+        let w = key_expansion::</*Nk,Nb,Nr*/ 8, 4, 14>(&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(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))
+        })?;
+
+        let proof = data.prove(pw)?;
+        data.verify(proof)
+    }
+
+    /// test against paper test vector (appendix B)
+    #[test]
+    fn test_encrypt_block_test_vector() -> Result<()> {
+        // Circuit declaration
+        let config = CircuitConfig::standard_recursion_config();
+        let mut builder = CircuitBuilder::<F, D>::new(config);
+
+        let key_target: [ByteArrayTarget; _] =
+            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);
+        let expanded_key_target = builder.key_expansion::<4, 10>(sbox_lut, key_target);
+
+        let input_state_target = builder.add_virtual_state();
+
+        let output = builder.aes_cipher(
+            sbox_lut,
+            mix_matrix,
+            input_state_target,
+            expanded_key_target,
+        );
+
+        let data = builder.build::<PoseidonGoldilocksConfig>();
+
+        let input_state: [[u8; _]; _] = [
+            [0x32, 0x43, 0xf6, 0xa8],
+            [0x88, 0x5a, 0x30, 0x8d],
+            [0x31, 0x31, 0x98, 0xa2],
+            [0xe0, 0x37, 0x07, 0x34],
+        ];
+
+        let key = [
+            0x2b, 0x7e, 0x15, 0x16, 0x28, 0xae, 0xd2, 0xa6, 0xab, 0xf7, 0x15, 0x88, 0x09, 0xcf,
+            0x4f, 0x3c,
+        ];
+
+        let expanded_key = key_expansion::<4, 4, 10>(&key);
+
+        let expected_result: [[u8; _]; _] = [
+            [0x39, 0x02, 0xdc, 0x19],
+            [0x25, 0xdc, 0x11, 0x6a],
+            [0x84, 0x09, 0x85, 0x0b],
+            [0x1d, 0xfb, 0x97, 0x32],
+        ];
+
+        // TODO
+        let expected_result = encrypt_block::<10>(
+            &array::from_fn(|i| input_state[i % 4][i / 4]),
+            &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))?;
+
+        pw.set_target_state(input_state_target, input_state)?;
+
+        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(output.0, expected_result).try_for_each(|(o, e)| {
+            std::iter::zip(o, e).try_for_each(|(o, e)| pw.set_byte_array_target(o, e))
+        })?;
+
+        let proof = data.prove(pw)?;
+        data.verify(proof)
+    }
+}

M src/lib.rs → src/lib.rs

@@ -3,5 +3,9 @@ #![allow(incomplete_features)]
 #![feature(generic_const_exprs)]
 #![allow(dead_code)] // TMP
 
+mod circuit;
 mod constants;
 mod native;
+
+/// D defines the extension degree of the field used in the Plonky2 proofs (quadratic extension).
+pub const D: usize = 2;

M src/native.rs → src/native.rs

@@ -6,14 +6,17 @@ //! in mind the approach that will be done in-circuit; this means that it is not
 //! written in the most rust idiomatic way, but simulating the behavior that we
 //! will do later inside the circuit.
 
+use std::array;
+
 use crate::constants::{RCON, SBOX};
 
-type State = [[u8; 4]; 4];
+pub(crate) type State = [[u8; 4]; 4];
 
 fn flatten_state(s: State) -> [u8; 16] {
     let mut r = [0u8; 16];
     for i in 0..4 {
         for j in 0..4 {
+            // TODO
             r[i * 4 + j] = s[i][j];
         }
     }
@@ -22,12 +25,16 @@ }
 
 /// encrypts an AES block (16 bytes). NR determines the number of rounds,
 /// where AES-128: NR=10, AES-196: NR=12, AES-256: NR=14.
-fn encrypt_block<const NR: usize>(input: &[u8; 16], w: &[[u8; 4]; 4 * (NR + 1)]) -> State {
+pub(crate) fn encrypt_block<const NR: usize>(
+    input: &[u8; 16],
+    w: &[[u8; 4]; 4 * (NR + 1)],
+) -> State {
     assert_eq!(4 * (NR + 1), w.len());
 
     let mut s: State = [[0; 4]; 4];
     for i in 0..4 {
         for j in 0..4 {
+            // TODO
             s[i][j] = input[i + 4 * j];
         }
     }
@@ -48,7 +55,7 @@
     s
 }
 
-fn sub_bytes(s: State) -> State {
+pub(crate) fn sub_bytes(s: State) -> State {
     let mut r = [[0u8; 4]; 4];
     for i in 0..4 {
         for j in 0..4 {
@@ -58,17 +65,11 @@ }
     r
 }
 
-fn shift_rows(s: State) -> State {
-    let mut r = [[0u8; 4]; 4];
-    for i in 0..4 {
-        for j in 0..4 {
-            r[i][j] = s[i][(i + j) % 4];
-        }
-    }
-    r
+pub(crate) fn shift_rows<T: Copy>(s: [[T; 4]; 4]) -> [[T; 4]; 4] {
+    array::from_fn(|i| array::from_fn(|j| s[i][(i + j) % 4]))
 }
 
-fn mix_columns(s: State) -> State {
+pub(crate) fn mix_columns(s: State) -> State {
     let mut r = [[0u8; 4]; 4];
     for c in 0..4 {
         r[0][c] = gf_2_8_mul(0x02, s[0][c]) ^ gf_2_8_mul(0x03, s[1][c]) ^ s[2][c] ^ s[3][c];
@@ -109,7 +110,7 @@ }
     r
 }
 
-fn key_expansion<const NK: usize, const NB: usize, const NR: usize>(
+pub(crate) fn key_expansion<const NK: usize, const NB: usize, const NR: usize>(
     key: &[u8; NK * NB],
 ) -> [[u8; 4]; 4 * (NR + 1)]
 where
@@ -135,10 +136,8 @@ }
     w
 }
 
-fn rot_word(w: [u8; 4]) -> [u8; 4] {
-    let mut r = [0u8; 4];
-    (r[0], r[1], r[2], r[3]) = (w[1], w[2], w[3], w[0]);
-    r
+pub(crate) fn rot_word<T: Copy>(w: [T; 4]) -> [T; 4] {
+    array::from_fn(|i| w[(i + 1) % 4])
 }
 fn sub_word(w: [u8; 4]) -> [u8; 4] {
     let mut r = [0u8; 4];