M src/circuit.rs → src/circuit.rs

@@ -11,9 +11,9 @@ plonk::circuit_builder::CircuitBuilder,
 };
 
 use crate::{
-    D,
     constants::{RCON, SBOX},
-    native::{State, rot_word, shift_rows},
+    native::{rot_word, shift_rows, State},
+    D,
 };
 
 #[derive(Debug, Copy, Clone)]
@@ -399,10 +399,10 @@ use rand::RngExt;
 
     use crate::{
         circuit::{
-            ByteArrayTarget, CircuitBuilderAESState, PartialWitnessAESState,
-            PartialWitnessByteArray, sbox_lut, state_mix_matrix_bits,
+            sbox_lut, state_mix_matrix_bits, ByteArrayTarget, CircuitBuilderAESState,
+            PartialWitnessAESState, PartialWitnessByteArray,
         },
-        native::{State, encrypt_block, key_expansion, mix_columns, sub_bytes},
+        native::{encrypt_block, key_expansion, mix_columns, sub_bytes, State},
     };
 
     use super::D;
@@ -506,73 +506,41 @@ 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>();
-
+    fn test_key_expansion() -> Result<()> {
         // 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>();
+        test_key_expansion_op::</*Nk,Nb,Nr*/ 4, 4, 10>(key)?;
 
         // 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);
+        test_key_expansion_op::</*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))
-        })?;
+        // 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,
+        ];
+        test_key_expansion_op::</*Nk,Nb,Nr*/ 8, 4, 14>(key)?;
 
-        let proof = data.prove(pw)?;
-        data.verify(proof)
+        Ok(())
     }
 
-    #[test]
-    fn test_key_expansion_256() -> Result<()> {
+    /// Analogous tests to `native::tests::test_key_expansion`.
+    fn test_key_expansion_op<const NK: usize, const NB: usize, const NR: usize>(
+        key: [u8; NK * NB],
+    ) -> Result<()>
+    where
+        [(); 4 * (NR + 1)]:,
+        [(); 4 * NK]:,
+    {
         // Circuit declaration
         let config = CircuitConfig::standard_recursion_config();
         let mut builder = CircuitBuilder::<F, D>::new(config);
@@ -580,17 +548,18 @@
         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 expanded_key_target = builder.key_expansion::<NK, NR>(sbox_lut, key_target);
 
+        println!(
+            "key_expansion(NK:{}, NB:{}, NR:{}) num_gates: {}",
+            NK,
+            NB,
+            NR,
+            builder.num_gates()
+        );
         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 w = key_expansion::<NK, NB, NR>(&key);
 
         let mut pw = PartialWitness::<F>::new();
 
@@ -625,6 +594,8 @@ mix_matrix,
             input_state_target,
             expanded_key_target,
         );
+
+        println!("aes_cipher num_gates: {}", builder.num_gates());
 
         let data = builder.build::<PoseidonGoldilocksConfig>();
 

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

@@ -115,7 +115,6 @@ key: &[u8; NK * NB],
 ) -> [[u8; 4]; 4 * (NR + 1)]
 where
     [(); 4 * (NR + 1)]:,
-    [(); 4 * 4 * (NR + 1)]:,
 {
     let mut w = [[0u8; 4]; 4 * (NR + 1)]; // expanded key