Arrabbiata: new poseidon gadget

arrabbiata/src/column.rs

@@ -18,12 +18,19 @@ pub enum Gadget {
     // Elliptic curve related gadgets
     EllipticCurveAddition,
     EllipticCurveScaling,
-    /// This gadget implement the Poseidon hash instance described in the
-    /// top-level documentation. This implementation does use the "next row"
-    /// to allow the computation of one additional round per row. In the current
-    /// setup, with [crate::NUMBER_OF_COLUMNS] columns, we can compute 5 full
+    /// The following gadgets implement the Poseidon hash instance described in
+    /// the top-level documentation. In the current setup, with
+    /// [crate::NUMBER_OF_COLUMNS] columns, we can compute 5 full
     /// rounds per row.
     Poseidon,
+    /// We provide a new Poseidon gadget that allows computing 5 rounds, without
+    /// using "public inputs".
+    ///
+    /// We split the Poseidon gadget in 13 sub-gadgets, one for each set of 5
+    /// permutations and one for the absorbtion. The parameteris the starting
+    /// round of Poseidon. It is expected to be a multiple of five.
+    PoseidonPermutation(usize),
+    PoseidonSpongeAbsorb,
 }
 
 #[derive(Debug, Clone, Copy, PartialEq)]
@@ -56,6 +63,19 @@ impl From<Column> for usize {
 
 pub type E<Fp> = Expr<ConstantExpr<Fp, ChallengeTerm>, Column>;
 
+impl From<Gadget> for usize {
+    fn from(val: Gadget) -> usize {
+        match val {
+            Gadget::App => 0,
+            Gadget::EllipticCurveAddition => 1,
+            Gadget::EllipticCurveScaling => 2,
+            Gadget::Poseidon => 3,
+            Gadget::PoseidonSpongeAbsorb => 4,
+            Gadget::PoseidonPermutation(starting_round) => 5 + starting_round / 5,
+        }
+    }
+}
+
 // Code to allow for pretty printing of the expressions
 impl FormattedOutput for Column {
     fn latex(&self, _cache: &mut HashMap<CacheId, Self>) -> String {
@@ -65,6 +85,10 @@ impl FormattedOutput for Column {
                 Gadget::EllipticCurveAddition => "q_ec_add".to_string(),
                 Gadget::EllipticCurveScaling => "q_ec_mul".to_string(),
                 Gadget::Poseidon => "q_pos".to_string(),
+                Gadget::PoseidonSpongeAbsorb => "q_pos_sponge_absorb".to_string(),
+                Gadget::PoseidonPermutation(starting_round) => {
+                    format!("q_pos_permutation{}", starting_round)
+                }
             },
             Column::PublicInput(i) => format!("pi_{{{i}}}").to_string(),
             Column::X(i) => format!("x_{{{i}}}").to_string(),
@@ -77,7 +101,11 @@ impl FormattedOutput for Column {
                 Gadget::App => "q_app".to_string(),
                 Gadget::EllipticCurveAddition => "q_ec_add".to_string(),
                 Gadget::EllipticCurveScaling => "q_ec_mul".to_string(),
-                Gadget::Poseidon => "q_pos_next_row".to_string(),
+                Gadget::Poseidon => "q_pos".to_string(),
+                Gadget::PoseidonSpongeAbsorb => "q_pos_sponge_absorb".to_string(),
+                Gadget::PoseidonPermutation(starting_round) => {
+                    format!("q_pos_permutation{}", starting_round)
+                }
             },
             Column::PublicInput(i) => format!("pi[{i}]"),
             Column::X(i) => format!("x[{i}]"),

arrabbiata/src/constraint.rs

@@ -193,6 +193,12 @@ where
         Expr::Atom(ExprInner::Cell(Variable { col, row }))
     }
 
+    fn get_poseidon_round_constant_as_constant(&self, round: usize, i: usize) -> Self::Variable {
+        let cst = C::sponge_params().round_constants[round][i];
+        let cst_inner = Operations::from(Literal(cst));
+        Self::Variable::constant(cst_inner)
+    }
+
     fn get_poseidon_mds_matrix(&mut self, i: usize, j: usize) -> Self::Variable {
         let v = C::sponge_params().mds[i][j];
         let v_inner = Operations::from(Literal(v));

arrabbiata/src/interpreter.rs

@@ -596,12 +596,17 @@ use num_bigint::BigInt;
 #[derive(Copy, Clone, Debug)]
 pub enum Instruction {
     /// This gadget implement the Poseidon hash instance described in the
-    /// top-level documentation. Compared to the previous one (that might be
-    /// deprecated in the future), this implementation does use the "next row"
-    /// to allow the computation of one additional round per row. In the current
-    /// setup, with [NUMBER_OF_COLUMNS] columns, we can compute 5 full rounds
-    /// per row.
+    /// top-level documentation. In the current setup, with [NUMBER_OF_COLUMNS]
+    /// columns, we can compute 5 full rounds per row.
     Poseidon(usize),
+    /// We provide a new Poseidon gadget that allows computing 5 rounds, without
+    /// using "public inputs".
+    ///
+    /// We split the Poseidon gadget in 13 sub-gadgets, one for each set of 5
+    /// permutations and one for the absorbtion. The parameteris the starting
+    /// round of Poseidon. It is expected to be a multiple of five.
+    PoseidonPermutation(usize),
+    PoseidonSpongeAbsorb,
     EllipticCurveScaling(usize, u64),
     EllipticCurveAddition(usize),
     // The NoOp will simply do nothing
@@ -730,6 +735,11 @@ pub trait InterpreterEnv {
         i: usize,
     ) -> Self::Variable;
 
+    /// Return the Poseidon round constants as a constant.
+    /// It aims to replace [InterpreterEnv::get_poseidon_round_constant] when
+    /// we get rid of the gadget [Gadget::Poseidon].
+    fn get_poseidon_round_constant_as_constant(&self, round: usize, i: usize) -> Self::Variable;
+
     /// Return the requested MDS matrix coefficient
     fn get_poseidon_mds_matrix(&mut self, i: usize, j: usize) -> Self::Variable;
 
@@ -1166,6 +1176,94 @@ pub fn run_ivc<E: InterpreterEnv>(env: &mut E, instr: Instruction) {
                 );
             }
         }
+        Instruction::PoseidonPermutation(starting_round) => {
+            assert!(
+                starting_round < PlonkSpongeConstants::PERM_ROUNDS_FULL,
+                "Invalid round index. Only values below {} are allowed.",
+                PlonkSpongeConstants::PERM_ROUNDS_FULL
+            );
+            assert!(
+                starting_round % 5 == 0,
+                "Invalid round index. Only values that are multiple of 5 are allowed."
+            );
+            debug!(
+                "Executing instruction Poseidon starting from round {starting_round} to {}",
+                starting_round + 5
+            );
+
+            env.activate_gadget(Gadget::PoseidonPermutation(starting_round));
+
+            let round_input_positions: Vec<E::Position> = (0..PlonkSpongeConstants::SPONGE_WIDTH)
+                .map(|_i| env.allocate())
+                .collect();
+
+            let round_output_positions: Vec<E::Position> = (0..PlonkSpongeConstants::SPONGE_WIDTH)
+                .map(|_i| env.allocate_next_row())
+                .collect();
+
+            let state: Vec<E::Variable> = if starting_round == 0 {
+                round_input_positions
+                    .iter()
+                    .enumerate()
+                    .map(|(i, pos)| env.load_poseidon_state(*pos, i))
+                    .collect()
+            } else {
+                round_input_positions
+                    .iter()
+                    .map(|pos| env.read_position(*pos))
+                    .collect()
+            };
+
+            // 5 is the number of rounds we treat per row
+            (0..5).fold(state, |state, idx_round| {
+                let state: Vec<E::Variable> =
+                    state.iter().map(|x| env.compute_x5(x.clone())).collect();
+
+                let round = starting_round + idx_round;
+
+                let rcs: Vec<E::Variable> = (0..PlonkSpongeConstants::SPONGE_WIDTH)
+                    .map(|i| env.get_poseidon_round_constant_as_constant(round, i))
+                    .collect();
+
+                let state: Vec<E::Variable> = rcs
+                    .iter()
+                    .enumerate()
+                    .map(|(i, rc)| {
+                        let acc: E::Variable =
+                            state.iter().enumerate().fold(env.zero(), |acc, (j, x)| {
+                                acc + env.get_poseidon_mds_matrix(i, j) * x.clone()
+                            });
+                        // The last iteration is written on the next row.
+                        if idx_round == 4 {
+                            env.write_column(round_output_positions[i], acc + rc.clone())
+                        } else {
+                            // Otherwise, we simply allocate a new position
+                            // in the circuit.
+                            let pos = env.allocate();
+                            env.write_column(pos, acc + rc.clone())
+                        }
+                    })
+                    .collect();
+
+                // If we are at the last round, we save the state in the
+                // environment.
+                // FIXME/IMPROVEME: we might want to execute more Poseidon
+                // full hash in sequentially, and then save one row. For
+                // now, we will save the state at the end of the last round
+                // and reload it at the beginning of the next Poseidon full
+                // hash.
+                if round == PlonkSpongeConstants::PERM_ROUNDS_FULL - 1 {
+                    state.iter().enumerate().for_each(|(i, x)| {
+                        unsafe { env.save_poseidon_state(x.clone(), i) };
+                    });
+                };
+
+                state
+            });
+        }
+        Instruction::PoseidonSpongeAbsorb => {
+            todo!()
+        }
         Instruction::NoOp => {}
     }
 

arrabbiata/src/lib.rs

@@ -64,7 +64,8 @@ pub const MAXIMUM_FIELD_SIZE_IN_BITS: u64 = 255;
 pub const NUMBER_OF_VALUES_TO_ABSORB_PUBLIC_IO: usize = NUMBER_OF_COLUMNS * 2;
 
 /// The number of selectors used in the circuit.
-pub const NUMBER_OF_SELECTORS: usize = column::Gadget::COUNT;
+pub const NUMBER_OF_SELECTORS: usize =
+    column::Gadget::COUNT + (PlonkSpongeConstants::PERM_ROUNDS_FULL / 5);
 
 /// The arity of the multivariate polynomials describing the constraints.
 /// We consider, erroneously, that a public input can be considered as a

arrabbiata/src/witness.rs

@@ -338,7 +338,8 @@ where
     /// Activate the gadget for the current row
     fn activate_gadget(&mut self, gadget: Gadget) {
         // IMPROVEME: it should be called only once per row
-        self.selectors[gadget as usize][self.current_row] = true;
+        let gadget = usize::from(gadget);
+        self.selectors[gadget][self.current_row] = true;
     }
 
     fn constrain_boolean(&mut self, x: Self::Variable) {
@@ -467,6 +468,18 @@ where
         self.write_public_input(pos, rc)
     }
 
+    fn get_poseidon_round_constant_as_constant(&self, round: usize, i: usize) -> Self::Variable {
+        if self.current_iteration % 2 == 0 {
+            E1::sponge_params().round_constants[round][i]
+                .to_biguint()
+                .into()
+        } else {
+            E2::sponge_params().round_constants[round][i]
+                .to_biguint()
+                .into()
+        }
+    }
+
     fn get_poseidon_mds_matrix(&mut self, i: usize, j: usize) -> Self::Variable {
         if self.current_iteration % 2 == 0 {
             E1::sponge_params().mds[i][j].to_biguint().into()
@@ -1182,6 +1195,12 @@ where
                 }
             }
             Instruction::NoOp => Instruction::NoOp,
+            Instruction::PoseidonSpongeAbsorb => {
+                todo!()
+            }
+            Instruction::PoseidonPermutation(_) => {
+                todo!()
+            }
         }
     }
 

arrabbiata/tests/constraint.rs

@@ -130,7 +130,6 @@ fn test_degree_of_constraints_ivc() {
 #[test]
 fn test_gadget_elliptic_curve_scaling() {
     let instr = Instruction::EllipticCurveScaling(0, 0);
-    // FIXME: update when the gadget is fnished
     helper_compute_constraints_gadget(instr, 10);
 
     let mut exp_degrees = HashMap::new();
@@ -143,6 +142,20 @@ fn test_gadget_elliptic_curve_scaling() {
     helper_check_gadget_activated(instr, Gadget::EllipticCurveScaling);
 }
 
+#[test]
+fn test_gadget_poseidon_permutation() {
+    let instr = Instruction::PoseidonPermutation(0);
+    helper_compute_constraints_gadget(instr, 15);
+
+    let mut exp_degrees = HashMap::new();
+    exp_degrees.insert(5, 15);
+    helper_check_expected_degree_constraints(instr, exp_degrees);
+
+    helper_gadget_number_of_columns_used(instr, 15, 0);
+
+    helper_check_gadget_activated(instr, Gadget::PoseidonPermutation(0));
+}
+
 #[test]
 fn test_get_mvpoly_equivalent() {
     // Check that each constraint can be converted to a MVPoly. The type of the

arrabbiata/tests/witness.rs

@@ -16,7 +16,7 @@ use poly_commitment::{commitment::CommitmentCurve, PolyComm};
 use rand::{CryptoRng, RngCore};
 
 #[test]
-fn test_unit_witness_poseidon_next_row_gadget_one_full_hash() {
+fn test_unit_witness_poseidon_gadget_one_full_hash() {
     let srs_log2_size = 6;
     let sponge: [BigInt; PlonkSpongeConstants::SPONGE_WIDTH] =
         std::array::from_fn(|_i| BigInt::from(42u64));
@@ -53,6 +53,61 @@ fn test_unit_witness_poseidon_next_row_gadget_one_full_hash() {
             .map(|x| x.to_biguint().into())
             .collect::<Vec<_>>()
     };
+
+    // Check correctness for current iteration
+    assert_eq!(env.sponge_e1.to_vec(), exp_output);
+    // Check the other sponge hasn't been modified
+    assert_eq!(env.sponge_e2, sponge.clone());
+
+    // Number of rows used by one full hash
+    assert_eq!(env.current_row, 12);
+}
+
+#[test]
+fn test_unit_witness_poseidon_permutation_gadget_one_full_hash() {
+    // Expected output:
+    // 13562506435502224548799089445428941958058503946524561166818119397766682137724
+    // 27423099486669760867028539664936216880884888701599404075691059826529320129892
+    // 736058628407775696076653472820678709906041621699240400715815852096937303940
+    let srs_log2_size = 6;
+    let indexed_relation: IndexedRelation<Fp, Fq, Vesta, Pallas> =
+        IndexedRelation::new(srs_log2_size);
+
+    let sponge: [BigInt; PlonkSpongeConstants::SPONGE_WIDTH] = [
+        BigInt::from(42u64),
+        BigInt::from(42u64),
+        BigInt::from(42u64),
+    ];
+
+    let mut env = Env::<Fp, Fq, Vesta, Pallas>::new(
+        BigInt::from(1u64),
+        sponge.clone(),
+        sponge.clone(),
+        indexed_relation,
+    );
+
+    env.current_instruction = Instruction::PoseidonPermutation(0);
+
+    (0..(PlonkSpongeConstants::PERM_ROUNDS_FULL / 5)).for_each(|i| {
+        interpreter::run_ivc(&mut env, Instruction::PoseidonPermutation(5 * i));
+        env.reset();
+    });
+    let exp_output = {
+        let mut state = sponge
+            .clone()
+            .to_vec()
+            .iter()
+            .map(|x| Fp::from_biguint(&x.to_biguint().unwrap()).unwrap())
+            .collect::<Vec<_>>();
+        poseidon_block_cipher::<Fp, PlonkSpongeConstants>(
+            poseidon_3_60_0_5_5_fp::static_params(),
+            &mut state,
+        );
+        state
+            .iter()
+            .map(|x| x.to_biguint().into())
+            .collect::<Vec<_>>()
+    };
     // Check correctness for current iteration
     assert_eq!(env.sponge_e1.to_vec(), exp_output);
     // Check the other sponge hasn't been modified