From a68609e3fbf525fec51f9257fd68c1976e85b731 Mon Sep 17 00:00:00 2001 From: Xinding Wei Date: Wed, 16 Aug 2023 15:26:53 -0700 Subject: [PATCH 1/4] Add Poseidon chip --- halo2-base/Cargo.toml | 3 + halo2-base/src/lib.rs | 2 + halo2-base/src/poseidon/mds.rs | 153 ++++++++++++++++++++++++++++++ halo2-base/src/poseidon/mod.rs | 110 ++++++++++++++++++++++ halo2-base/src/poseidon/spec.rs | 157 +++++++++++++++++++++++++++++++ halo2-base/src/poseidon/state.rs | 131 ++++++++++++++++++++++++++ halo2-base/src/poseidon/tests.rs | 99 +++++++++++++++++++ halo2-base/src/utils/mod.rs | 4 +- 8 files changed, 657 insertions(+), 2 deletions(-) create mode 100644 halo2-base/src/poseidon/mds.rs create mode 100644 halo2-base/src/poseidon/mod.rs create mode 100644 halo2-base/src/poseidon/spec.rs create mode 100644 halo2-base/src/poseidon/state.rs create mode 100644 halo2-base/src/poseidon/tests.rs diff --git a/halo2-base/Cargo.toml b/halo2-base/Cargo.toml index 93f0f21b..21a7add8 100644 --- a/halo2-base/Cargo.toml +++ b/halo2-base/Cargo.toml @@ -20,6 +20,9 @@ halo2_proofs_axiom = { git = "https://github.com/axiom-crypto/halo2.git", packag # Use PSE halo2 and halo2curves for compatibility when feature = "halo2-pse" is on halo2_proofs = { git = "https://github.com/privacy-scaling-explorations/halo2.git", rev = "f348757", optional = true } +# This is Scroll's audited poseidon circuit. We only use it for the Native Poseidon spec. We do not use the halo2 circuit at all (and it wouldn't even work because the halo2_proofs tag is not compatbile). +# We forked it to upgrade to ff v0.13 and removed the circuit module +poseidon-rs = { git = "https://github.com/axiom-crypto/poseidon-circuit.git", rev = "1aee4a1" } # plotting circuit layout plotters = { version = "0.3.0", optional = true } tabbycat = { version = "0.1", features = ["attributes"], optional = true } diff --git a/halo2-base/src/lib.rs b/halo2-base/src/lib.rs index 9f20386e..e5890fce 100644 --- a/halo2-base/src/lib.rs +++ b/halo2-base/src/lib.rs @@ -41,6 +41,8 @@ use utils::ScalarField; /// Module that contains the main API for creating and working with circuits. pub mod gates; +/// Module for the Poseidon hash function. +pub mod poseidon; /// Module for SafeType which enforce value range and realted functions. pub mod safe_types; /// Utility functions for converting between different types of field elements. diff --git a/halo2-base/src/poseidon/mds.rs b/halo2-base/src/poseidon/mds.rs new file mode 100644 index 00000000..d4bef74b --- /dev/null +++ b/halo2-base/src/poseidon/mds.rs @@ -0,0 +1,153 @@ +use crate::utils::ScalarField; + +/// The type used to hold the MDS matrix +pub(crate) type Mds = [[F; T]; T]; + +/// `MDSMatrices` holds the MDS matrix as well as transition matrix which is +/// also called `pre_sparse_mds` and sparse matrices that enables us to reduce +/// number of multiplications in apply MDS step +#[derive(Debug, Clone)] +pub struct MDSMatrices { + pub(crate) mds: MDSMatrix, + pub(crate) pre_sparse_mds: MDSMatrix, + pub(crate) sparse_matrices: Vec>, +} + +/// `SparseMDSMatrix` are in `[row], [hat | identity]` form and used in linear +/// layer of partial rounds instead of the original MDS +#[derive(Debug, Clone)] +pub struct SparseMDSMatrix { + pub(crate) row: [F; T], + pub(crate) col_hat: [F; RATE], +} + +/// `MDSMatrix` is applied to `State` to achive linear layer of Poseidon +#[derive(Clone, Debug)] +pub struct MDSMatrix(pub(crate) Mds); + +impl MDSMatrix { + pub(crate) fn mul_vector(&self, v: &[F; T]) -> [F; T] { + let mut res = [F::ZERO; T]; + for i in 0..T { + for j in 0..T { + res[i] += self.0[i][j] * v[j]; + } + } + res + } + + pub(crate) fn identity() -> Mds { + let mut mds = [[F::ZERO; T]; T]; + for i in 0..T { + mds[i][i] = F::ONE; + } + mds + } + + /// Multiplies two MDS matrices. Used in sparse matrix calculations + pub(crate) fn mul(&self, other: &Self) -> Self { + let mut res = [[F::ZERO; T]; T]; + for i in 0..T { + for j in 0..T { + for k in 0..T { + res[i][j] += self.0[i][k] * other.0[k][j]; + } + } + } + Self(res) + } + + pub(crate) fn transpose(&self) -> Self { + let mut res = [[F::ZERO; T]; T]; + for i in 0..T { + for j in 0..T { + res[i][j] = self.0[j][i]; + } + } + Self(res) + } + + pub(crate) fn determinant(m: [[F; N]; N]) -> F { + let mut res = F::ONE; + let mut m = m; + for i in 0..N { + let mut pivot = i; + while m[pivot][i] == F::ZERO { + pivot += 1; + assert!(pivot < N, "matrix is not invertible"); + } + if pivot != i { + res = -res; + m.swap(pivot, i); + } + res *= m[i][i]; + let inv = m[i][i].invert().unwrap(); + for j in i + 1..N { + let factor = m[j][i] * inv; + for k in i + 1..N { + m[j][k] -= m[i][k] * factor; + } + } + } + res + } + + /// See Section B in Supplementary Material https://eprint.iacr.org/2019/458.pdf + /// Factorises an MDS matrix `M` into `M'` and `M''` where `M = M' * M''`. + /// Resulted `M''` matrices are the sparse ones while `M'` will contribute + /// to the accumulator of the process + pub(crate) fn factorise(&self) -> (Self, SparseMDSMatrix) { + assert_eq!(RATE + 1, T); + // Given `(t-1 * t-1)` MDS matrix called `hat` constructs the `t * t` matrix in + // form `[[1 | 0], [0 | m]]`, ie `hat` is the right bottom sub-matrix + let prime = |hat: Mds| -> Self { + let mut prime = Self::identity(); + for (prime_row, hat_row) in prime.iter_mut().skip(1).zip(hat.iter()) { + for (el_prime, el_hat) in prime_row.iter_mut().skip(1).zip(hat_row.iter()) { + *el_prime = *el_hat; + } + } + Self(prime) + }; + + // Given `(t-1)` sized `w_hat` vector constructs the matrix in form + // `[[m_0_0 | m_0_i], [w_hat | identity]]` + let prime_prime = |w_hat: [F; RATE]| -> Mds { + let mut prime_prime = Self::identity(); + prime_prime[0] = self.0[0]; + for (row, w) in prime_prime.iter_mut().skip(1).zip(w_hat.iter()) { + row[0] = *w + } + prime_prime + }; + + let w = self.0.iter().skip(1).map(|row| row[0]).collect::>(); + // m_hat is the `(t-1 * t-1)` right bottom sub-matrix of m := self.0 + let mut m_hat = [[F::ZERO; RATE]; RATE]; + for i in 0..RATE { + for j in 0..RATE { + m_hat[i][j] = self.0[i + 1][j + 1]; + } + } + // w_hat = m_hat^{-1} * w, where m_hat^{-1} is matrix inverse and * is matrix mult + // we avoid computing m_hat^{-1} explicitly by using Cramer's rule: https://en.wikipedia.org/wiki/Cramer%27s_rule + let mut w_hat = [F::ZERO; RATE]; + let det = Self::determinant(m_hat); + let det_inv = Option::::from(det.invert()).expect("matrix is not invertible"); + for j in 0..RATE { + let mut m_hat_j = m_hat; + for i in 0..RATE { + m_hat_j[i][j] = w[i]; + } + w_hat[j] = Self::determinant(m_hat_j) * det_inv; + } + let m_prime = prime(m_hat); + let m_prime_prime = prime_prime(w_hat); + // row = first row of m_prime_prime.transpose() = first column of m_prime_prime + let row: [F; T] = + m_prime_prime.iter().map(|row| row[0]).collect::>().try_into().unwrap(); + // col_hat = first column of m_prime_prime.transpose() without first element = first row of m_prime_prime without first element + let col_hat: [F; RATE] = m_prime_prime[0][1..].try_into().unwrap(); + (m_prime, SparseMDSMatrix { row, col_hat }) + } +} diff --git a/halo2-base/src/poseidon/mod.rs b/halo2-base/src/poseidon/mod.rs new file mode 100644 index 00000000..a6584d0e --- /dev/null +++ b/halo2-base/src/poseidon/mod.rs @@ -0,0 +1,110 @@ +use crate::{ + gates::GateInstructions, + poseidon::{spec::OptimizedPoseidonSpec, state::PoseidonState}, + AssignedValue, Context, ScalarField, +}; + +#[cfg(test)] +mod tests; + +/// Module for maximum distance separable matrix operations. +pub mod mds; +/// Module for poseidon specification. +pub mod spec; +/// Module for poseidon states. +pub mod state; + +/// Chip for Poseidon hasher. The chip is stateful. +pub struct PoseidonHasherChip { + init_state: PoseidonState, + state: PoseidonState, + spec: OptimizedPoseidonSpec, + absorbing: Vec>, +} + +impl PoseidonHasherChip { + /// Create new Poseidon hasher chip. + pub fn new( + ctx: &mut Context, + ) -> Self { + let init_state = PoseidonState::::default(ctx); + let state = init_state.clone(); + Self { + init_state, + state, + spec: OptimizedPoseidonSpec::new::(), + absorbing: Vec::new(), + } + } + + /// Reset state to default and clear the buffer. + pub fn clear(&mut self) { + self.state = self.init_state.clone(); + self.absorbing.clear(); + } + + /// Store given `elements` into buffer. + pub fn update(&mut self, elements: &[AssignedValue]) { + self.absorbing.extend_from_slice(elements); + } + + /// Consume buffer and perform permutation, then output second element of + /// state. + pub fn squeeze( + &mut self, + ctx: &mut Context, + gate: &impl GateInstructions, + ) -> AssignedValue { + 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![]); + } + + 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.0; + + 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.0; + 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/halo2-base/src/poseidon/spec.rs b/halo2-base/src/poseidon/spec.rs new file mode 100644 index 00000000..24dcf7fc --- /dev/null +++ b/halo2-base/src/poseidon/spec.rs @@ -0,0 +1,157 @@ +use crate::{poseidon::mds::*, utils::ScalarField}; + +use poseidon_rs::poseidon::primitives::Spec as PoseidonSpec; // trait +use std::marker::PhantomData; + +// struct so we can use PoseidonSpec trait to generate round constants and MDS matrix +#[derive(Debug)] +pub(crate) struct Poseidon128Pow5Gen< + F: ScalarField, + const T: usize, + const RATE: usize, + const R_F: usize, + const R_P: usize, + const SECURE_MDS: usize, +> { + _marker: PhantomData, +} + +impl< + F: ScalarField, + const T: usize, + const RATE: usize, + const R_F: usize, + const R_P: usize, + const SECURE_MDS: usize, + > PoseidonSpec for Poseidon128Pow5Gen +{ + fn full_rounds() -> usize { + R_F + } + + fn partial_rounds() -> usize { + R_P + } + + fn sbox(val: F) -> F { + val.pow_vartime([5]) + } + + // see "Avoiding insecure matrices" in Section 2.3 of https://eprint.iacr.org/2019/458.pdf + // most Specs used in practice have SECURE_MDS = 0 + fn secure_mds() -> usize { + SECURE_MDS + } +} + +// We use the optimized Poseidon implementation described in Supplementary Material Section B of https://eprint.iacr.org/2019/458.pdf +// This involves some further computation of optimized constants and sparse MDS matrices beyond what the Scroll PoseidonSpec generates +// The implementation below is adapted from https://github.com/privacy-scaling-explorations/poseidon + +/// `OptimizedPoseidonSpec` holds construction parameters as well as constants that are used in +/// permutation step. +#[derive(Debug, Clone)] +pub struct OptimizedPoseidonSpec { + pub(crate) r_f: usize, + pub(crate) mds_matrices: MDSMatrices, + pub(crate) constants: OptimizedConstants, +} + +/// `OptimizedConstants` has round constants that are added each round. While +/// full rounds has T sized constants there is a single constant for each +/// partial round +#[derive(Debug, Clone)] +pub struct OptimizedConstants { + pub(crate) start: Vec<[F; T]>, + pub(crate) partial: Vec, + pub(crate) end: Vec<[F; T]>, +} + +impl OptimizedPoseidonSpec { + /// Generate new spec with specific number of full and partial rounds. `SECURE_MDS` is usually 0, but may need to be specified because insecure matrices may sometimes be generated + pub fn new() -> Self { + let (round_constants, mds, mds_inv) = + Poseidon128Pow5Gen::::constants(); + let mds = MDSMatrix(mds); + let inverse_mds = MDSMatrix(mds_inv); + + let constants = + Self::calculate_optimized_constants(R_F, R_P, round_constants, &inverse_mds); + let (sparse_matrices, pre_sparse_mds) = Self::calculate_sparse_matrices(R_P, &mds); + + Self { + r_f: R_F, + constants, + mds_matrices: MDSMatrices { mds, sparse_matrices, pre_sparse_mds }, + } + } + + fn calculate_optimized_constants( + r_f: usize, + r_p: usize, + constants: Vec<[F; T]>, + inverse_mds: &MDSMatrix, + ) -> OptimizedConstants { + let (number_of_rounds, r_f_half) = (r_f + r_p, r_f / 2); + assert_eq!(constants.len(), number_of_rounds); + + // Calculate optimized constants for first half of the full rounds + let mut constants_start: Vec<[F; T]> = vec![[F::ZERO; T]; r_f_half]; + constants_start[0] = constants[0]; + for (optimized, constants) in + constants_start.iter_mut().skip(1).zip(constants.iter().skip(1)) + { + *optimized = inverse_mds.mul_vector(constants); + } + + // Calculate constants for partial rounds + let mut acc = constants[r_f_half + r_p]; + let mut constants_partial = vec![F::ZERO; r_p]; + for (optimized, constants) in constants_partial + .iter_mut() + .rev() + .zip(constants.iter().skip(r_f_half).rev().skip(r_f_half)) + { + let mut tmp = inverse_mds.mul_vector(&acc); + *optimized = tmp[0]; + + tmp[0] = F::ZERO; + for ((acc, tmp), constant) in acc.iter_mut().zip(tmp).zip(constants.iter()) { + *acc = tmp + constant + } + } + constants_start.push(inverse_mds.mul_vector(&acc)); + + // Calculate optimized constants for ending half of the full rounds + let mut constants_end: Vec<[F; T]> = vec![[F::ZERO; T]; r_f_half - 1]; + for (optimized, constants) in + constants_end.iter_mut().zip(constants.iter().skip(r_f_half + r_p + 1)) + { + *optimized = inverse_mds.mul_vector(constants); + } + + OptimizedConstants { + start: constants_start, + partial: constants_partial, + end: constants_end, + } + } + + fn calculate_sparse_matrices( + r_p: usize, + mds: &MDSMatrix, + ) -> (Vec>, MDSMatrix) { + let mds = mds.transpose(); + let mut acc = mds.clone(); + let mut sparse_matrices = (0..r_p) + .map(|_| { + let (m_prime, m_prime_prime) = acc.factorise(); + acc = mds.mul(&m_prime); + m_prime_prime + }) + .collect::>>(); + + sparse_matrices.reverse(); + (sparse_matrices, acc.transpose()) + } +} diff --git a/halo2-base/src/poseidon/state.rs b/halo2-base/src/poseidon/state.rs new file mode 100644 index 00000000..81cdd3a7 --- /dev/null +++ b/halo2-base/src/poseidon/state.rs @@ -0,0 +1,131 @@ +use crate::{ + gates::GateInstructions, + poseidon::mds::SparseMDSMatrix, + utils::ScalarField, + AssignedValue, Context, + QuantumCell::{Constant, Existing}, +}; + +#[derive(Clone)] +pub(crate) struct PoseidonState { + pub(crate) s: [AssignedValue; T], +} + +impl PoseidonState { + pub fn default(ctx: &mut Context) -> Self { + let mut default_state = [F::ZERO; T]; + // from Section 4.2 of https://eprint.iacr.org/2019/458.pdf + // • Variable-Input-Length Hashing. The capacity value is 2^64 + (o−1) where o the output length. + // for our transcript use cases, o = 1 + default_state[0] = F::from_u128(1u128 << 64); + Self { s: default_state.map(|f| ctx.load_constant(f)) } + } + + pub 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)) + } + + pub 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); + } + } + + pub 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); + } + + pub 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 }), + ); + } + } + + pub 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(); + } + + pub 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 + } + } +} diff --git a/halo2-base/src/poseidon/tests.rs b/halo2-base/src/poseidon/tests.rs new file mode 100644 index 00000000..b15aba09 --- /dev/null +++ b/halo2-base/src/poseidon/tests.rs @@ -0,0 +1,99 @@ +use super::*; +use crate::{ + gates::{builder::GateThreadBuilder, GateChip}, + halo2_proofs::halo2curves::{bn256::Fr, ff::PrimeField}, +}; + +use itertools::Itertools; + +#[test] +fn test_mds() { + let spec = OptimizedPoseidonSpec::::new::<8, 57, 0>(); + + let mds = vec![ + vec![ + "7511745149465107256748700652201246547602992235352608707588321460060273774987", + "10370080108974718697676803824769673834027675643658433702224577712625900127200", + "19705173408229649878903981084052839426532978878058043055305024233888854471533", + ], + vec![ + "18732019378264290557468133440468564866454307626475683536618613112504878618481", + "20870176810702568768751421378473869562658540583882454726129544628203806653987", + "7266061498423634438633389053804536045105766754026813321943009179476902321146", + ], + vec![ + "9131299761947733513298312097611845208338517739621853568979632113419485819303", + "10595341252162738537912664445405114076324478519622938027420701542910180337937", + "11597556804922396090267472882856054602429588299176362916247939723151043581408", + ], + ]; + for (row1, row2) in mds.iter().zip_eq(spec.mds_matrices.mds.0.iter()) { + for (e1, e2) in row1.iter().zip_eq(row2.iter()) { + assert_eq!(Fr::from_str_vartime(e1).unwrap(), *e2); + } + } +} + +#[test] +fn test_poseidon_against_test_vectors() { + let mut builder = GateThreadBuilder::prover(); + let gate = GateChip::::default(); + let ctx = builder.main(0); + + // https://extgit.iaik.tugraz.at/krypto/hadeshash/-/blob/master/code/test_vectors.txt + // poseidonperm_x5_254_3 + { + const R_F: usize = 8; + const R_P: usize = 57; + const T: usize = 3; + const RATE: usize = 2; + + let mut hasher = PoseidonHasherChip::::new::(ctx); + + let state = [0u64, 1, 2]; + hasher.state = + PoseidonState:: { s: state.map(|v| ctx.load_constant(Fr::from(v))) }; + let inputs = [Fr::zero(); RATE].iter().map(|f| ctx.load_constant(*f)).collect_vec(); + hasher.permutation(ctx, &gate, inputs); // avoid padding + let state_0 = hasher.state.s; + let expected = [ + "7853200120776062878684798364095072458815029376092732009249414926327459813530", + "7142104613055408817911962100316808866448378443474503659992478482890339429929", + "6549537674122432311777789598043107870002137484850126429160507761192163713804", + ]; + for (word, expected) in state_0.into_iter().zip(expected.iter()) { + assert_eq!(word.value.evaluate(), Fr::from_str_vartime(expected).unwrap()); + } + } + + // https://extgit.iaik.tugraz.at/krypto/hadeshash/-/blob/master/code/test_vectors.txt + // poseidonperm_x5_254_5 + { + const R_F: usize = 8; + const R_P: usize = 60; + const T: usize = 5; + const RATE: usize = 4; + + let mut hasher = PoseidonHasherChip::::new::(ctx); + + let state = [0u64, 1, 2, 3, 4]; + hasher.state = + PoseidonState:: { s: state.map(|v| ctx.load_constant(Fr::from(v))) }; + let inputs = [Fr::zero(); RATE].iter().map(|f| ctx.load_constant(*f)).collect_vec(); + hasher.permutation(ctx, &gate, inputs); + let state_0 = hasher.state.s; + let expected = [ + "18821383157269793795438455681495246036402687001665670618754263018637548127333", + "7817711165059374331357136443537800893307845083525445872661165200086166013245", + "16733335996448830230979566039396561240864200624113062088822991822580465420551", + "6644334865470350789317807668685953492649391266180911382577082600917830417726", + "3372108894677221197912083238087960099443657816445944159266857514496320565191", + ]; + for (word, expected) in state_0.into_iter().zip(expected.iter()) { + assert_eq!(word.value.evaluate(), Fr::from_str_vartime(expected).unwrap()); + } + } +} + +// TODO: test clear()/squeeze(). +// TODO: test constraints actually work. diff --git a/halo2-base/src/utils/mod.rs b/halo2-base/src/utils/mod.rs index 7c91448f..29430345 100644 --- a/halo2-base/src/utils/mod.rs +++ b/halo2-base/src/utils/mod.rs @@ -1,6 +1,6 @@ use core::hash::Hash; -use crate::ff::PrimeField; +use crate::ff::{FromUniformBytes, PrimeField}; #[cfg(not(feature = "halo2-axiom"))] use crate::halo2_proofs::arithmetic::CurveAffine; use crate::halo2_proofs::circuit::Value; @@ -44,7 +44,7 @@ where /// Helper trait to represent a field element that can be converted into [u64] limbs. /// /// Note: Since the number of bits necessary to represent a field element is larger than the number of bits in a u64, we decompose the integer representation of the field element into multiple [u64] values e.g. `limbs`. -pub trait ScalarField: PrimeField + From + Hash + PartialEq + PartialOrd { +pub trait ScalarField: PrimeField + FromUniformBytes<64> + From + Hash + Ord { /// Returns the base `2bit_len` little endian representation of the [ScalarField] element up to `num_limbs` number of limbs (truncates any extra limbs). /// /// Assumes `bit_len < 64`. From 152e7cd6033cb92f164fe1b110342dc5845f3e4e Mon Sep 17 00:00:00 2001 From: Jonathan Wang <31040440+jonathanpwang@users.noreply.github.com> Date: Wed, 16 Aug 2023 18:25:01 -0700 Subject: [PATCH 2/4] chore: minor fixes --- halo2-base/src/poseidon/mds.rs | 1 + halo2-base/src/poseidon/state.rs | 2 +- halo2-base/src/poseidon/{tests.rs => tests/mod.rs} | 4 ++-- 3 files changed, 4 insertions(+), 3 deletions(-) rename halo2-base/src/poseidon/{tests.rs => tests/mod.rs} (95%) diff --git a/halo2-base/src/poseidon/mds.rs b/halo2-base/src/poseidon/mds.rs index d4bef74b..536fd7b3 100644 --- a/halo2-base/src/poseidon/mds.rs +++ b/halo2-base/src/poseidon/mds.rs @@ -1,3 +1,4 @@ +#![allow(clippy::needless_range_loop)] use crate::utils::ScalarField; /// The type used to hold the MDS matrix diff --git a/halo2-base/src/poseidon/state.rs b/halo2-base/src/poseidon/state.rs index 81cdd3a7..b533767d 100644 --- a/halo2-base/src/poseidon/state.rs +++ b/halo2-base/src/poseidon/state.rs @@ -124,7 +124,7 @@ impl PoseidonState Date: Wed, 16 Aug 2023 18:39:09 -0700 Subject: [PATCH 3/4] test(poseidon): add compatbility tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Cherry-picked from https://github.com/axiom-crypto/halo2-lib/pull/98 Co-authored-by: Antonio Mejías Gil --- halo2-base/Cargo.toml | 2 + .../src/poseidon/tests/compatibility.rs | 117 ++++++++++++++++++ halo2-base/src/poseidon/tests/mod.rs | 2 + 3 files changed, 121 insertions(+) create mode 100644 halo2-base/src/poseidon/tests/compatibility.rs diff --git a/halo2-base/Cargo.toml b/halo2-base/Cargo.toml index 21a7add8..183fb60f 100644 --- a/halo2-base/Cargo.toml +++ b/halo2-base/Cargo.toml @@ -38,6 +38,8 @@ criterion = "0.4" criterion-macro = "0.4" test-case = "3.1.0" proptest = "1.1.0" +# native poseidon for testing +pse-poseidon = { git = "https://github.com/axiom-crypto/pse-poseidon.git" } # memory allocation [target.'cfg(not(target_env = "msvc"))'.dependencies] diff --git a/halo2-base/src/poseidon/tests/compatibility.rs b/halo2-base/src/poseidon/tests/compatibility.rs new file mode 100644 index 00000000..383a83a0 --- /dev/null +++ b/halo2-base/src/poseidon/tests/compatibility.rs @@ -0,0 +1,117 @@ +use std::{cmp::max, iter::zip}; + +use crate::{ + gates::{builder::GateThreadBuilder, GateChip}, + halo2_proofs::halo2curves::bn256::Fr, + poseidon::PoseidonHasherChip, + utils::ScalarField, +}; +use pse_poseidon::Poseidon; +use rand::Rng; + +// make interleaved calls to absorb and squeeze elements and +// check that the result is the same in-circuit and natively +fn poseidon_compatiblity_verification< + F: ScalarField, + const T: usize, + const RATE: usize, + const R_F: usize, + const R_P: usize, +>( + // 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, +) { + let mut builder = GateThreadBuilder::prover(); + let gate = GateChip::default(); + + let ctx = builder.main(0); + + // constructing native and in-circuit Poseidon sponges + let mut native_sponge = Poseidon::::new(R_F, R_P); + // assuming SECURE_MDS = 0 + let mut circuit_sponge = PoseidonHasherChip::::new::(ctx); + + // 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(ctx, &gate); + + 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(ctx, &gate); + + 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); +} + +#[test] +fn test_poseidon_compatibility_absorbing_only() { + let absorptions = random_nested_list_f(8, 5); + let squeezings = Vec::new(); + + poseidon_compatiblity_verification::(absorptions, squeezings); +} + +#[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); +} + +#[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); +} diff --git a/halo2-base/src/poseidon/tests/mod.rs b/halo2-base/src/poseidon/tests/mod.rs index 5315f8a6..f4289ac0 100644 --- a/halo2-base/src/poseidon/tests/mod.rs +++ b/halo2-base/src/poseidon/tests/mod.rs @@ -6,6 +6,8 @@ use crate::{ use itertools::Itertools; +mod compatibility; + #[test] fn test_mds() { let spec = OptimizedPoseidonSpec::::new::<8, 57, 0>(); From 1fdcded5f5526889350ba03674ed73f49885b6dd Mon Sep 17 00:00:00 2001 From: Jonathan Wang <31040440+jonathanpwang@users.noreply.github.com> Date: Wed, 16 Aug 2023 19:08:22 -0700 Subject: [PATCH 4/4] chore: minor refactor to more closely match snark-verifier https://github.com/axiom-crypto/snark-verifier/blob/main/snark-verifier/src/util/hash/poseidon.rs --- halo2-base/src/poseidon/mod.rs | 28 +++++++++++++--------- halo2-base/src/poseidon/state.rs | 41 +++++++++++++++++--------------- 2 files changed, 39 insertions(+), 30 deletions(-) diff --git a/halo2-base/src/poseidon/mod.rs b/halo2-base/src/poseidon/mod.rs index a6584d0e..dcb1549a 100644 --- a/halo2-base/src/poseidon/mod.rs +++ b/halo2-base/src/poseidon/mod.rs @@ -1,3 +1,5 @@ +use std::mem; + use crate::{ gates::GateInstructions, poseidon::{spec::OptimizedPoseidonSpec, state::PoseidonState}, @@ -27,7 +29,7 @@ impl PoseidonHasherChip( ctx: &mut Context, ) -> Self { - let init_state = PoseidonState::::default(ctx); + let init_state = PoseidonState::default(ctx); let state = init_state.clone(); Self { init_state, @@ -37,6 +39,12 @@ impl PoseidonHasherChip, spec: OptimizedPoseidonSpec) -> Self { + let init_state = PoseidonState::default(ctx); + Self { spec, state: init_state.clone(), init_state, absorbing: Vec::new() } + } + /// Reset state to default and clear the buffer. pub fn clear(&mut self) { self.state = self.init_state.clone(); @@ -55,17 +63,13 @@ impl PoseidonHasherChip, gate: &impl GateInstructions, ) -> AssignedValue { - let mut input_elements = vec![]; - input_elements.append(&mut self.absorbing); - - let mut padding_offset = 0; + let input_elements = mem::take(&mut self.absorbing); + let exact = input_elements.len() % RATE == 0; for chunk in input_elements.chunks(RATE) { - padding_offset = RATE - chunk.len(); self.permutation(ctx, gate, chunk.to_vec()); } - - if padding_offset == 0 { + if exact { self.permutation(ctx, gate, vec![]); } @@ -80,25 +84,27 @@ impl PoseidonHasherChip PoseidonState PoseidonState PoseidonState, 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) { - *x = new_x - } + self.s = iter::once(gate.inner_product( + ctx, + self.s.iter().copied(), + mds.row.iter().map(|c| Constant(*c)), + )) + .chain( + mds.col_hat + .iter() + .zip(self.s.iter().skip(1)) + .map(|(coeff, state)| gate.mul_add(ctx, self.s[0], Constant(*coeff), *state)), + ) + .collect::>() + .try_into() + .unwrap(); } }