From b703ffd9518be168fb9d57b256cc73a0b4eb4d89 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Antonio=20Mej=C3=ADas=20Gil?= Date: Mon, 31 Jul 2023 09:03:51 +0200 Subject: [PATCH] Compatibility with native Poseidon sponge (#98) * fixed compatibility issue, added explanation * added comment, formatted * changed corrected computation into circuit constraint, removed old code * added first version of test * finalised tests * remove unnecessary builder construction from test * removed unnecessary variable k * added poseidon test to ci.yml --- .github/workflows/ci.yml | 4 + hashes/poseidon/src/lib.rs | 211 +++++++++++++++++++++++++++++++++++ hashes/poseidon/src/tests.rs | 119 ++++++++++++++++++++ 3 files changed, 334 insertions(+) create mode 100644 hashes/poseidon/src/lib.rs create mode 100644 hashes/poseidon/src/tests.rs diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 33cc92e1..eceb664a 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -21,6 +21,10 @@ jobs: working-directory: "halo2-base" run: | cargo test + - name: Run poseidon tests + working-directory: 'hashes/poseidon' + run: | + cargo test test_poseidon_compatibility - name: Run halo2-ecc tests (mock prover) working-directory: "halo2-ecc" run: | diff --git a/hashes/poseidon/src/lib.rs b/hashes/poseidon/src/lib.rs new file mode 100644 index 00000000..952b8288 --- /dev/null +++ b/hashes/poseidon/src/lib.rs @@ -0,0 +1,211 @@ +// impl taken from https://github.com/scroll-tech/halo2-snark-aggregator/tree/main/halo2-snark-aggregator-api/src/hash + +use ::poseidon::{SparseMDSMatrix, Spec, State}; +use halo2_base::halo2_proofs::plonk::Error; +use halo2_base::{ + gates::GateInstructions, + utils::ScalarField, + AssignedValue, Context, + QuantumCell::{Constant, Existing}, +}; + +pub mod tests; + +struct PoseidonState { + s: [AssignedValue; T], +} + +impl PoseidonState { + fn x_power5_with_constant( + ctx: &mut Context, + gate: &impl GateInstructions, + x: AssignedValue, + constant: &F, + ) -> AssignedValue { + let x2 = gate.mul(ctx, x, x); + let x4 = gate.mul(ctx, x2, x2); + gate.mul_add(ctx, x, x4, Constant(*constant)) + } + + fn sbox_full( + &mut self, + ctx: &mut Context, + gate: &impl GateInstructions, + constants: &[F; T], + ) { + for (x, constant) in self.s.iter_mut().zip(constants.iter()) { + *x = Self::x_power5_with_constant(ctx, gate, *x, constant); + } + } + + fn sbox_part(&mut self, ctx: &mut Context, gate: &impl GateInstructions, constant: &F) { + let x = &mut self.s[0]; + *x = Self::x_power5_with_constant(ctx, gate, *x, constant); + } + + fn absorb_with_pre_constants( + &mut self, + ctx: &mut Context, + gate: &impl GateInstructions, + inputs: Vec>, + pre_constants: &[F; T], + ) { + assert!(inputs.len() < T); + let offset = inputs.len() + 1; + + // Explanation of what's going on: before each round of the poseidon permutation, + // two things have to be added to the state: inputs (the absorbed elements) and + // preconstants. Imagine the state as a list of T elements, the first of which is + // the capacity: |--cap--|--el1--|--el2--|--elR--| + // - A preconstant is added to each of all T elements (which is different for each) + // - The inputs are added to all elements starting from el1 (so, not to the capacity), + // to as many elements as inputs are available. + // - To the first element for which no input is left (if any), an extra 1 is added. + + // adding preconstant to the distinguished capacity element (only one) + self.s[0] = gate.add(ctx, self.s[0], Constant(pre_constants[0])); + + // adding pre-constants and inputs to the elements for which both are available + for ((x, constant), input) in + self.s.iter_mut().skip(1).zip(pre_constants.iter().skip(1)).zip(inputs.iter()) + { + *x = gate.sum(ctx, [Existing(*x), Existing(*input), Constant(*constant)]); + } + + // adding only pre-constants when no input is left + for (i, (x, constant)) in + self.s.iter_mut().skip(offset).zip(pre_constants.iter().skip(offset)).enumerate() + { + *x = gate.add( + ctx, + Existing(*x), + Constant(if i == 0 { F::one() + constant } else { *constant }), + ); + } + } + + fn apply_mds( + &mut self, + ctx: &mut Context, + gate: &impl GateInstructions, + mds: &[[F; T]; T], + ) { + let res = mds + .iter() + .map(|row| { + gate.inner_product(ctx, self.s.iter().copied(), row.iter().map(|c| Constant(*c))) + }) + .collect::>(); + + self.s = res.try_into().unwrap(); + } + + fn apply_sparse_mds( + &mut self, + ctx: &mut Context, + gate: &impl GateInstructions, + mds: &SparseMDSMatrix, + ) { + let sum = + gate.inner_product(ctx, self.s.iter().copied(), mds.row().iter().map(|c| Constant(*c))); + let mut res = vec![sum]; + + for (e, x) in mds.col_hat().iter().zip(self.s.iter().skip(1)) { + res.push(gate.mul_add(ctx, self.s[0], Constant(*e), *x)); + } + + for (x, new_x) in self.s.iter_mut().zip(res.into_iter()) { + *x = new_x + } + } +} + +pub struct PoseidonChip { + init_state: [AssignedValue; T], + state: PoseidonState, + spec: Spec, + absorbing: Vec>, +} + +impl PoseidonChip { + pub fn new(ctx: &mut Context, r_f: usize, r_p: usize) -> Result { + let init_state = State::::default() + .words() + .into_iter() + .map(|x| ctx.load_constant(x)) + .collect::>>(); + Ok(Self { + spec: Spec::new(r_f, r_p), + init_state: init_state.clone().try_into().unwrap(), + state: PoseidonState { s: init_state.try_into().unwrap() }, + absorbing: Vec::new(), + }) + } + + pub fn clear(&mut self) { + self.state = PoseidonState { s: self.init_state }; + self.absorbing.clear(); + } + + pub fn update(&mut self, elements: &[AssignedValue]) { + self.absorbing.extend_from_slice(elements); + } + + pub fn squeeze( + &mut self, + ctx: &mut Context, + gate: &impl GateInstructions, + ) -> Result, Error> { + let mut input_elements = vec![]; + input_elements.append(&mut self.absorbing); + + let mut padding_offset = 0; + + for chunk in input_elements.chunks(RATE) { + padding_offset = RATE - chunk.len(); + self.permutation(ctx, gate, chunk.to_vec()); + } + + if padding_offset == 0 { + self.permutation(ctx, gate, vec![]); + } + + Ok(self.state.s[1]) + } + + fn permutation( + &mut self, + ctx: &mut Context, + gate: &impl GateInstructions, + inputs: Vec>, + ) { + let r_f = self.spec.r_f() / 2; + let mds = &self.spec.mds_matrices().mds().rows(); + + let constants = self.spec.constants().start(); + self.state.absorb_with_pre_constants(ctx, gate, inputs, &constants[0]); + for constants in constants.iter().skip(1).take(r_f - 1) { + self.state.sbox_full(ctx, gate, constants); + self.state.apply_mds(ctx, gate, mds); + } + + let pre_sparse_mds = &self.spec.mds_matrices().pre_sparse_mds().rows(); + self.state.sbox_full(ctx, gate, constants.last().unwrap()); + self.state.apply_mds(ctx, gate, pre_sparse_mds); + + let sparse_matrices = &self.spec.mds_matrices().sparse_matrices(); + let constants = &self.spec.constants().partial(); + for (constant, sparse_mds) in constants.iter().zip(sparse_matrices.iter()) { + self.state.sbox_part(ctx, gate, constant); + self.state.apply_sparse_mds(ctx, gate, sparse_mds); + } + + let constants = &self.spec.constants().end(); + for constants in constants.iter() { + self.state.sbox_full(ctx, gate, constants); + self.state.apply_mds(ctx, gate, mds); + } + self.state.sbox_full(ctx, gate, &[F::zero(); T]); + self.state.apply_mds(ctx, gate, mds); + } +} diff --git a/hashes/poseidon/src/tests.rs b/hashes/poseidon/src/tests.rs new file mode 100644 index 00000000..fd9625ad --- /dev/null +++ b/hashes/poseidon/src/tests.rs @@ -0,0 +1,119 @@ +#[cfg(test)] +mod tests { + use std::{cmp::max, iter::zip}; + + use halo2_base::{ + gates::{builder::GateThreadBuilder, GateChip}, + halo2_proofs::halo2curves::bn256::Fr, + utils::ScalarField, + }; + use poseidon::Poseidon; + use rand::Rng; + + use crate::PoseidonChip; + + // make interleaved calls to absorb and squeeze elements and + // check that the result is the same in-circuit and natively + fn poseidon_compatiblity_verification( + // elements of F to absorb; one sublist = one absorption + mut absorptions: Vec>, + // list of amounts of elements of F that should be squeezed every time + mut squeezings: Vec, + rounds_full: usize, + rounts_partial: usize, + ) { + let mut builder = GateThreadBuilder::prover(); + let gate = GateChip::default(); + + let mut ctx = builder.main(0); + + // constructing native and in-circuit Poseidon sponges + let mut native_sponge = Poseidon::::new(rounds_full, rounts_partial); + let mut circuit_sponge = + PoseidonChip::::new(&mut ctx, rounds_full, rounts_partial) + .expect("Failed to construct Poseidon circuit"); + + // preparing to interleave absorptions and squeezings + let n_iterations = max(absorptions.len(), squeezings.len()); + absorptions.resize(n_iterations, Vec::new()); + squeezings.resize(n_iterations, 0); + + for (absorption, squeezing) in zip(absorptions, squeezings) { + // absorb (if any elements were provided) + native_sponge.update(&absorption); + circuit_sponge.update(&ctx.assign_witnesses(absorption)); + + // squeeze (if any elements were requested) + for _ in 0..squeezing { + let native_squeezed = native_sponge.squeeze(); + let circuit_squeezed = + circuit_sponge.squeeze(&mut ctx, &gate).expect("Failed to squeeze"); + + assert_eq!(native_squeezed, *circuit_squeezed.value()); + } + } + + // even if no squeezings were requested, we squeeze to verify the + // states are the same after all absorptions + let native_squeezed = native_sponge.squeeze(); + let circuit_squeezed = circuit_sponge.squeeze(&mut ctx, &gate).expect("Failed to squeeze"); + + assert_eq!(native_squeezed, *circuit_squeezed.value()); + } + + fn random_nested_list_f(len: usize, max_sub_len: usize) -> Vec> { + let mut rng = rand::thread_rng(); + let mut list = Vec::new(); + for _ in 0..len { + let len = rng.gen_range(0..=max_sub_len); + let mut sublist = Vec::new(); + + for _ in 0..len { + sublist.push(F::random(&mut rng)); + } + list.push(sublist); + } + list + } + + fn random_list_usize(len: usize, max: usize) -> Vec { + let mut rng = rand::thread_rng(); + let mut list = Vec::new(); + for _ in 0..len { + list.push(rng.gen_range(0..=max)); + } + list + } + + #[test] + fn test_poseidon_compatibility_squeezing_only() { + let absorptions = Vec::new(); + let squeezings = random_list_usize(10, 7); + + poseidon_compatiblity_verification::(absorptions, squeezings, 8, 57); + } + + #[test] + fn test_poseidon_compatibility_absorbing_only() { + let absorptions = random_nested_list_f(8, 5); + let squeezings = Vec::new(); + + poseidon_compatiblity_verification::(absorptions, squeezings, 8, 57); + } + + #[test] + fn test_poseidon_compatibility_interleaved() { + let absorptions = random_nested_list_f(10, 5); + let squeezings = random_list_usize(7, 10); + + poseidon_compatiblity_verification::(absorptions, squeezings, 8, 57); + } + + #[test] + fn test_poseidon_compatibility_other_params() { + let absorptions = random_nested_list_f(10, 10); + let squeezings = random_list_usize(10, 10); + + poseidon_compatiblity_verification::(absorptions, squeezings, 8, 120); + } +}