feat: recursion circuit p2 skinny air constraints

recursion/circuit-v2/src/challenger.rs

@@ -96,7 +96,7 @@ impl<C: Config> DuplexChallengerVariable<C> {
         self.sponge_state[0..self.input_buffer.len()].copy_from_slice(self.input_buffer.as_slice());
         self.input_buffer.clear();
 
-        self.sponge_state = builder.poseidon2_permute_v2_wide(self.sponge_state);
+        self.sponge_state = builder.poseidon2_permute_v2(self.sponge_state);
 
         self.output_buffer.clear();
         self.output_buffer.extend_from_slice(&self.sponge_state);
@@ -224,7 +224,7 @@ pub(crate) mod tests {
     use p3_field::AbstractField;
     use sp1_core::stark::StarkGenericConfig;
     use sp1_core::utils::setup_logger;
-    use sp1_core::utils::BabyBearPoseidon2Inner;
+    use sp1_core::utils::BabyBearPoseidon2;
     use sp1_recursion_compiler::asm::AsmBuilder;
     use sp1_recursion_compiler::asm::AsmConfig;
     use sp1_recursion_compiler::circuit::AsmCompiler;
@@ -241,7 +241,7 @@ pub(crate) mod tests {
 
     use sp1_core::utils::run_test_machine;
 
-    type SC = BabyBearPoseidon2Inner;
+    type SC = BabyBearPoseidon2;
     type F = <SC as StarkGenericConfig>::Val;
     type EF = <SC as StarkGenericConfig>::Challenge;
 
@@ -261,10 +261,18 @@ pub(crate) mod tests {
 
         let records = vec![runtime.record];
 
-        let machine = RecursionAir::<_, 3, 0>::machine_with_all_chips(config);
-        let (pk, vk) = machine.setup(&program);
+        // Run with the poseidon2 wide chip.
+        let wide_machine = RecursionAir::<_, 3, 0>::machine_wide(SC::default());
+        let (pk, vk) = wide_machine.setup(&program);
+        let result = run_test_machine(records.clone(), wide_machine, pk, vk);
+        if let Err(e) = result {
+            panic!("Verification failed: {:?}", e);
+        }
 
-        let result = run_test_machine(records.clone(), machine, pk, vk);
+        // Run with the poseidon2 skinny chip.
+        let skinny_machine = RecursionAir::<_, 9, 0>::machine(SC::compressed());
+        let (pk, vk) = skinny_machine.setup(&program);
+        let result = run_test_machine(records.clone(), skinny_machine, pk, vk);
         if let Err(e) = result {
             panic!("Verification failed: {:?}", e);
         }

recursion/compiler/src/circuit/builder.rs

@@ -15,8 +15,7 @@ pub trait CircuitV2Builder<C: Config> {
     fn num2bits_v2_f(&mut self, num: Felt<C::F>, num_bits: usize) -> Vec<Felt<C::F>>;
     fn exp_reverse_bits_v2(&mut self, input: Felt<C::F>, power_bits: Vec<Felt<C::F>>)
         -> Felt<C::F>;
-    fn poseidon2_permute_v2_skinny(&mut self, state: [Felt<C::F>; WIDTH]) -> [Felt<C::F>; WIDTH];
-    fn poseidon2_permute_v2_wide(&mut self, state: [Felt<C::F>; WIDTH]) -> [Felt<C::F>; WIDTH];
+    fn poseidon2_permute_v2(&mut self, state: [Felt<C::F>; WIDTH]) -> [Felt<C::F>; WIDTH];
     fn poseidon2_hash_v2(&mut self, array: &[Felt<C::F>]) -> [Felt<C::F>; DIGEST_SIZE];
     fn poseidon2_compress_v2(
         &mut self,
@@ -70,21 +69,13 @@ impl<C: Config> CircuitV2Builder<C> for Builder<C> {
         output
     }
     /// Applies the Poseidon2 permutation to the given array.
-    fn poseidon2_permute_v2_skinny(&mut self, array: [Felt<C::F>; WIDTH]) -> [Felt<C::F>; WIDTH] {
+    fn poseidon2_permute_v2(&mut self, array: [Felt<C::F>; WIDTH]) -> [Felt<C::F>; WIDTH] {
         let output: [Felt<C::F>; WIDTH] = core::array::from_fn(|_| self.uninit());
         self.operations
-            .push(DslIr::CircuitV2Poseidon2PermuteBabyBearSkinny(
-                output, array,
-            ));
-        output
-    }
-    /// Applies the Poseidon2 permutation to the given array using the wide precompile.
-    fn poseidon2_permute_v2_wide(&mut self, array: [Felt<C::F>; WIDTH]) -> [Felt<C::F>; WIDTH] {
-        let output: [Felt<C::F>; WIDTH] = core::array::from_fn(|_| self.uninit());
-        self.operations
-            .push(DslIr::CircuitV2Poseidon2PermuteBabyBearWide(output, array));
+            .push(DslIr::CircuitV2Poseidon2PermuteBabyBear(output, array));
         output
     }
+
     /// Applies the Poseidon2 permutation to the given array.
     ///
     /// Reference: [p3_symmetric::PaddingFreeSponge]
@@ -93,7 +84,7 @@ impl<C: Config> CircuitV2Builder<C> for Builder<C> {
         let mut state = core::array::from_fn(|_| self.eval(C::F::zero()));
         for input_chunk in input.chunks(HASH_RATE) {
             state[..input_chunk.len()].copy_from_slice(input_chunk);
-            state = self.poseidon2_permute_v2_wide(state);
+            state = self.poseidon2_permute_v2(state);
         }
         let state: [Felt<C::F>; DIGEST_SIZE] = state[..DIGEST_SIZE].try_into().unwrap();
         state
@@ -108,7 +99,7 @@ impl<C: Config> CircuitV2Builder<C> for Builder<C> {
         // debug_assert!(DIGEST_SIZE * N <= WIDTH);
         let mut pre_iter = input.into_iter().chain(repeat(self.eval(C::F::default())));
         let pre = core::array::from_fn(move |_| pre_iter.next().unwrap());
-        let post = self.poseidon2_permute_v2_wide(pre);
+        let post = self.poseidon2_permute_v2(pre);
         let post: [Felt<C::F>; DIGEST_SIZE] = post[..DIGEST_SIZE].try_into().unwrap();
         post
     }

recursion/compiler/src/circuit/compiler.rs

@@ -227,26 +227,12 @@ impl<C: Config> AsmCompiler<C> {
         ]
     }
 
-    fn poseidon2_permute_skinny(
+    fn poseidon2_permute(
         &mut self,
         dst: [impl Reg<C>; WIDTH],
         src: [impl Reg<C>; WIDTH],
     ) -> Instruction<C::F> {
-        Instruction::Poseidon2Skinny(Poseidon2WideInstr {
-            addrs: Poseidon2Io {
-                input: src.map(|r| r.read(self)),
-                output: dst.map(|r| r.write(self)),
-            },
-            mults: [C::F::zero(); WIDTH],
-        })
-    }
-
-    fn poseidon2_permute_wide(
-        &mut self,
-        dst: [impl Reg<C>; WIDTH],
-        src: [impl Reg<C>; WIDTH],
-    ) -> Instruction<C::F> {
-        Instruction::Poseidon2Wide(Poseidon2WideInstr {
+        Instruction::Poseidon2(Poseidon2Instr {
             addrs: Poseidon2Io {
                 input: src.map(|r| r.read(self)),
                 output: dst.map(|r| r.write(self)),
@@ -422,11 +408,8 @@ impl<C: Config> AsmCompiler<C> {
             DslIr::AssertNeFI(lhs, rhs) => self.base_assert_ne(lhs, Imm::F(rhs)),
             DslIr::AssertNeEI(lhs, rhs) => self.ext_assert_ne(lhs, Imm::EF(rhs)),
 
-            DslIr::CircuitV2Poseidon2PermuteBabyBearSkinny(dst, src) => {
-                vec![self.poseidon2_permute_skinny(dst, src)]
-            }
-            DslIr::CircuitV2Poseidon2PermuteBabyBearWide(dst, src) => {
-                vec![self.poseidon2_permute_wide(dst, src)]
+            DslIr::CircuitV2Poseidon2PermuteBabyBear(dst, src) => {
+                vec![self.poseidon2_permute(dst, src)]
             }
             DslIr::CircuitV2ExpReverseBits(dst, base, exp) => {
                 vec![self.exp_reverse_bits(dst, base, exp)]
@@ -521,11 +504,7 @@ impl<C: Config> AsmCompiler<C> {
                     kind: MemAccessKind::Write,
                     ..
                 }) => vec![(mult, inner)],
-                Instruction::Poseidon2Skinny(Poseidon2SkinnyInstr {
-                    addrs: Poseidon2Io { ref output, .. },
-                    mults,
-                }) => mults.iter_mut().zip(output).collect(),
-                Instruction::Poseidon2Wide(Poseidon2WideInstr {
+                Instruction::Poseidon2(Poseidon2SkinnyInstr {
                     addrs: Poseidon2Io { ref output, .. },
                     mults,
                 }) => mults.iter_mut().zip(output).collect(),
@@ -738,7 +717,6 @@ mod tests {
     type SC = BabyBearPoseidon2;
     type F = <SC as StarkGenericConfig>::Val;
     type EF = <SC as StarkGenericConfig>::Challenge;
-    type A = RecursionAir<F, 3, 1>;
     fn test_operations(operations: TracedVec<DslIr<AsmConfig<F, EF>>>) {
         test_operations_with_runner(operations, |program| {
             let mut runtime = Runtime::<F, EF, DiffusionMatrixBabyBear>::new(
@@ -758,39 +736,25 @@ mod tests {
         let program = compiler.compile(operations);
         let record = run(&program);
 
-        let config = SC::new();
-        let machine = A::machine_with_all_chips(config);
-        let (pk, vk) = machine.setup(&program);
-        let result = run_test_machine(vec![record], machine, pk, vk);
+        // Run with the poseidon2 wide chip.
+        let wide_machine = RecursionAir::<_, 3, 0>::machine_wide(BabyBearPoseidon2::default());
+        let (pk, vk) = wide_machine.setup(&program);
+        let result = run_test_machine(vec![record.clone()], wide_machine, pk, vk);
         if let Err(e) = result {
             panic!("Verification failed: {:?}", e);
         }
-    }
-
-    #[test]
-    fn test_poseidon2_skinny() {
-        setup_logger();
-
-        let mut builder = AsmBuilder::<F, EF>::default();
-        let mut rng = StdRng::seed_from_u64(0xCAFEDA7E)
-            .sample_iter::<[F; WIDTH], _>(rand::distributions::Standard);
-        for _ in 0..100 {
-            let input_1: [F; WIDTH] = rng.next().unwrap();
-            let output_1 = inner_perm().permute(input_1);
 
-            let input_1_felts = input_1.map(|x| builder.eval(x));
-            let output_1_felts = builder.poseidon2_permute_v2_skinny(input_1_felts);
-            let expected: [Felt<_>; WIDTH] = output_1.map(|x| builder.eval(x));
-            for (lhs, rhs) in output_1_felts.into_iter().zip(expected) {
-                builder.assert_felt_eq(lhs, rhs);
-            }
+        // Run with the poseidon2 skinny chip.
+        let skinny_machine = RecursionAir::<_, 9, 0>::machine(BabyBearPoseidon2::compressed());
+        let (pk, vk) = skinny_machine.setup(&program);
+        let result = run_test_machine(vec![record.clone()], skinny_machine, pk, vk);
+        if let Err(e) = result {
+            panic!("Verification failed: {:?}", e);
         }
-
-        test_operations(builder.operations);
     }
 
     #[test]
-    fn test_poseidon2_wide() {
+    fn test_poseidon2() {
         setup_logger();
 
         let mut builder = AsmBuilder::<F, EF>::default();
@@ -801,7 +765,7 @@ mod tests {
             let output_1 = inner_perm().permute(input_1);
 
             let input_1_felts = input_1.map(|x| builder.eval(x));
-            let output_1_felts = builder.poseidon2_permute_v2_wide(input_1_felts);
+            let output_1_felts = builder.poseidon2_permute_v2(input_1_felts);
             let expected: [Felt<_>; WIDTH] = output_1.map(|x| builder.eval(x));
             for (lhs, rhs) in output_1_felts.into_iter().zip(expected) {
                 builder.assert_felt_eq(lhs, rhs);

recursion/compiler/src/ir/instructions.rs

@@ -210,9 +210,7 @@ pub enum DslIr<C: Config> {
     /// Permutates an array of BabyBear elements in the circuit.
     CircuitPoseidon2PermuteBabyBear([Felt<C::F>; 16]),
     /// Permutates an array of BabyBear elements in the circuit using the skinny precompile.
-    CircuitV2Poseidon2PermuteBabyBearSkinny([Felt<C::F>; 16], [Felt<C::F>; 16]),
-    /// Permutates an array of BabyBear elements in the circuit using the wide precompile.
-    CircuitV2Poseidon2PermuteBabyBearWide([Felt<C::F>; 16], [Felt<C::F>; 16]),
+    CircuitV2Poseidon2PermuteBabyBear([Felt<C::F>; 16], [Felt<C::F>; 16]),
 
     // Miscellaneous instructions.
     /// Decompose hint operation of a usize into an array. (output = num2bits(usize)).

recursion/core-v2/src/chips/alu_base.rs

@@ -201,18 +201,13 @@ where
 
 #[cfg(test)]
 mod tests {
-    use machine::RecursionAir;
-    use p3_baby_bear::{BabyBear, DiffusionMatrixBabyBear};
+    use machine::tests::run_recursion_test_machines;
+    use p3_baby_bear::BabyBear;
     use p3_field::AbstractField;
     use p3_matrix::dense::RowMajorMatrix;
 
     use rand::{rngs::StdRng, Rng, SeedableRng};
-    use sp1_core::{
-        air::MachineAir,
-        stark::StarkGenericConfig,
-        utils::{run_test_machine, BabyBearPoseidon2},
-    };
-    use sp1_recursion_core::stark::config::BabyBearPoseidon2Outer;
+    use sp1_core::{air::MachineAir, stark::StarkGenericConfig, utils::BabyBearPoseidon2};
 
     use super::*;
 
@@ -237,10 +232,8 @@ mod tests {
 
     #[test]
     pub fn four_ops() {
-        type SC = BabyBearPoseidon2Outer;
+        type SC = BabyBearPoseidon2;
         type F = <SC as StarkGenericConfig>::Val;
-        type EF = <SC as StarkGenericConfig>::Challenge;
-        type A = RecursionAir<F, 3, 1>;
 
         let mut rng = StdRng::seed_from_u64(0xDEADBEEF);
         let mut random_felt = move || -> F { rng.sample(rand::distributions::Standard) };
@@ -274,16 +267,6 @@ mod tests {
             traces: Default::default(),
         };
 
-        let config = SC::new();
-
-        let mut runtime =
-            Runtime::<F, EF, DiffusionMatrixBabyBear>::new(&program, BabyBearPoseidon2::new().perm);
-        runtime.run().unwrap();
-        let machine = A::machine(config);
-        let (pk, vk) = machine.setup(&program);
-        let result = run_test_machine(vec![runtime.record], machine, pk, vk);
-        if let Err(e) = result {
-            panic!("Verification failed: {:?}", e);
-        }
+        run_recursion_test_machines(program);
     }
 }

recursion/core-v2/src/chips/alu_ext.rs

@@ -183,17 +183,13 @@ where
 
 #[cfg(test)]
 mod tests {
-    use machine::RecursionAir;
-    use p3_baby_bear::{BabyBear, DiffusionMatrixBabyBear};
+    use machine::tests::run_recursion_test_machines;
+    use p3_baby_bear::BabyBear;
     use p3_field::{extension::BinomialExtensionField, AbstractExtensionField, AbstractField};
     use p3_matrix::dense::RowMajorMatrix;
 
     use rand::{rngs::StdRng, Rng, SeedableRng};
-    use sp1_core::{
-        air::MachineAir,
-        stark::StarkGenericConfig,
-        utils::{run_test_machine, BabyBearPoseidon2Inner},
-    };
+    use sp1_core::{air::MachineAir, stark::StarkGenericConfig};
     use sp1_recursion_core::stark::config::BabyBearPoseidon2Outer;
 
     use super::*;
@@ -221,8 +217,6 @@ mod tests {
     pub fn four_ops() {
         type SC = BabyBearPoseidon2Outer;
         type F = <SC as StarkGenericConfig>::Val;
-        type EF = <SC as StarkGenericConfig>::Challenge;
-        type A = RecursionAir<F, 3, 1>;
 
         let mut rng = StdRng::seed_from_u64(0xDEADBEEF);
         let mut random_extfelt = move || {
@@ -258,18 +252,7 @@ mod tests {
             instructions,
             traces: Default::default(),
         };
-        let mut runtime = Runtime::<F, EF, DiffusionMatrixBabyBear>::new(
-            &program,
-            BabyBearPoseidon2Inner::new().perm,
-        );
-        runtime.run().unwrap();
-
-        let config = SC::new();
-        let machine = A::machine(config);
-        let (pk, vk) = machine.setup(&program);
-        let result = run_test_machine(vec![runtime.record], machine, pk, vk);
-        if let Err(e) = result {
-            panic!("Verification failed: {:?}", e);
-        }
+
+        run_recursion_test_machines(program);
     }
 }

recursion/core-v2/src/chips/exp_reverse_bits.rs

@@ -346,35 +346,29 @@ mod tests {
     use rand::Rng;
     use rand::SeedableRng;
     use sp1_core::air::MachineAir;
-    use sp1_core::utils::run_test_machine;
     use sp1_core::utils::setup_logger;
-    use sp1_core::utils::BabyBearPoseidon2;
     use sp1_recursion_core::stark::config::BabyBearPoseidon2Outer;
     use std::iter::once;
 
     use p3_baby_bear::BabyBear;
-    use p3_baby_bear::DiffusionMatrixBabyBear;
     use p3_field::{AbstractField, PrimeField32};
     use p3_matrix::dense::RowMajorMatrix;
     use sp1_core::stark::StarkGenericConfig;
 
     use crate::chips::exp_reverse_bits::ExpReverseBitsLenChip;
-    use crate::machine::RecursionAir;
+    use crate::machine::tests::run_recursion_test_machines;
     use crate::runtime::instruction as instr;
     use crate::runtime::ExecutionRecord;
     use crate::ExpReverseBitsEvent;
     use crate::Instruction;
     use crate::MemAccessKind;
     use crate::RecursionProgram;
-    use crate::Runtime;
 
     #[test]
     fn prove_babybear_circuit_erbl() {
         setup_logger();
         type SC = BabyBearPoseidon2Outer;
         type F = <SC as StarkGenericConfig>::Val;
-        type EF = <SC as StarkGenericConfig>::Challenge;
-        type A = RecursionAir<F, 3, 1>;
 
         let mut rng = StdRng::seed_from_u64(0xDEADBEEF);
         let mut random_felt = move || -> F { F::from_canonical_u32(rng.gen_range(0..1 << 16)) };
@@ -434,17 +428,7 @@ mod tests {
             traces: Default::default(),
         };
 
-        let config = SC::new();
-
-        let mut runtime =
-            Runtime::<F, EF, DiffusionMatrixBabyBear>::new(&program, BabyBearPoseidon2::new().perm);
-        runtime.run().unwrap();
-        let machine = A::machine(config);
-        let (pk, vk) = machine.setup(&program);
-        let result = run_test_machine(vec![runtime.record], machine, pk, vk);
-        if let Err(e) = result {
-            panic!("Verification failed: {:?}", e);
-        }
+        run_recursion_test_machines(program);
     }
 
     #[test]